import xarray as xr
import numpy as np

ds_mass = xr.open_dataset("GEOSChem.AerosolMass.20190701_0000z.nc4")
ds_species = xr.open_dataset("GEOSChem.SpeciesConc.20190701_0000z.nc4")
ds_met = xr.open_dataset("GEOSChem.StateMet.20190701_0000z.nc4")
# 'ug m-3'
total_PM25 = ds_mass['PM25'][0,0,:,:].values 
BC = ds_mass['AerMassBC'][0,0,:,:].values  
NH4 = ds_mass['AerMassNH4'][0,0,:,:].values 
NIT = ds_mass['AerMassNIT'][0,0,:,:].values 
SO4 = ds_mass['AerMassSO4'][0,0,:,:].values 


# 'mol mol-1 dry' 
DST1 = ds_species['SpeciesConcVV_DST1'][0,0,:,:].values 
DST2 = ds_species['SpeciesConcVV_DST2'][0,0,:,:].values 
DST3 = ds_species['SpeciesConcVV_DST3'][0,0,:,:].values
DST4 = ds_species['SpeciesConcVV_DST4'][0,0,:,:].values 
SOAS = ds_species['SpeciesConcVV_SOAS'][0,0,:,:].values 
OCPO = ds_species['SpeciesConcVV_OCPO'][0,0,:,:].values
OCPI = ds_species['SpeciesConcVV_OCPI'][0,0,:,:].values 
SALA = ds_species['SpeciesConcVV_SALA'][0,0,:,:].values 
BCPI = ds_species['SpeciesConcVV_BCPI'][0,0,:,:].values
BCPO = ds_species['SpeciesConcVV_BCPO'][0,0,:,:].values 
SO4_SP = ds_species['SpeciesConcVV_SO4'][0,0,:,:].values 
NH4_SP = ds_species['SpeciesConcVV_NH4'][0,0,:,:].values 
NIT_SP = ds_species['SpeciesConcVV_NIT'][0,0,:,:].values 
HMS_SP = ds_species['SpeciesConcVV_HMS'][0,0,:,:].values 


air_density = ds_met['Met_AIRDEN'][0,0,:,:].values # 'kg m-3'
AD = ds_met['Met_AD'][0,0,:,:].values  # 'Dry air mass'，'kg'
AIRVOL = ds_met['Met_AIRVOL'][0,0,:,:].values  # 'Volume of dry air in grid box'，'m3'
AIR_DEN = AD / AIRVOL
PMID = ds_met['Met_PMID'][0,0,:,:].values  # 'Pressure (w/r/t moist air) at level centers' 'hPa'
T = ds_met['Met_T'][0,0,:,:].values  # 'Temperature' 'K'

# 'mol mol-1 dry' to 'ug m-3'
def dst_convert_unit(dust1_molar,dust2_molar,soas_molar,ocpo_molar,ocpi_molar,sala_molar,bcpi_molar,bcpo_molar,so4_molar,nh4_molar,nit_molar,hms_molar,air_density):
    MW_dust = 29  # g/mol
    MW_soas = 150
    MW_ocpo = 12.01
    MW_ocpi = 12.01
    MW_sala = 31.4
    MW_bc = 12.01
    MW_so4 = 96.6
    MW_nh4 = 18.05
    MW_nit = 62.01
    MW_hms = 111.10
    MW_air = 28.9644 # was 28.97 g/mol
    # to µg/m³
    dust1_mass = dust1_molar * (MW_dust / MW_air) * air_density * 1e9 * (1013.25 / PMID) * (T / 298)
    dust2_mass = dust2_molar * (MW_dust / MW_air) * air_density * 1e9 * (1013.25 / PMID) * (T / 298)
    soas_mass = soas_molar * (MW_soas / MW_air) * air_density * 1e9 * (1013.25 / PMID) * (T / 298)
    ocpo_mass = ocpo_molar * (MW_ocpo / MW_air) * air_density * 1e9 * (1013.25 / PMID) * (T / 298)
    ocpi_mass = ocpi_molar * (MW_ocpi / MW_air) * air_density * 1e9 * (1013.25 / PMID) * (T / 298)
    sala_mass = sala_molar * (MW_sala / MW_air) * air_density * 1e9 * (1013.25 / PMID) * (T / 298)
    bcpi_mass = bcpi_molar * (MW_bc / MW_air) * air_density * 1e9 * (1013.25 / PMID) * (T / 298)
    bcpo_mass = bcpo_molar * (MW_bc / MW_air) * air_density * 1e9 * (1013.25 / PMID) * (T / 298)
    so4_mass = so4_molar  * (MW_so4 / MW_air) * air_density * 1e9 * (1013.25 / PMID) * (T / 298)
    nh4_mass= nh4_molar  * (MW_nh4 / MW_air) * air_density * 1e9 * (1013.25 / PMID) * (T / 298)
    nit_mass= nit_molar  * (MW_nit / MW_air) * air_density * 1e9 * (1013.25 / PMID) * (T / 298)
    hms_mass= hms_molar  * (MW_hms / MW_air) * air_density * 1e9 * (1013.25 / PMID) * (T / 298)

    return dust1_mass, dust2_mass, soas_mass, ocpo_mass, ocpi_mass, sala_mass, bcpi_mass,bcpo_mass,so4_mass,nh4_mass,nit_mass,hms_mass

DST1_converted,DST2_converted,SOAS_converted,OCPO_converted,OCPI_converted,SALA_converted,BCPI_converted,BCPO_converted,SO4_converted,NH4_converted,NIT_converted,HMS_converted = \
    dst_convert_unit(DST1,DST2,SOAS,OCPO,OCPI,SALA,BCPI,BCPO,SO4_SP,NH4_SP,NIT_SP,HMS_SP,AIR_DEN)

# ======== Calculate PM25 concentration based on the formula for PM2.5 in GC ============#
total_pm25_cal = 1.1* (SO4_converted+NH4_converted+NIT_converted+HMS_converted) +BCPI_converted+BCPO_converted + \
                 2.1* (OCPO_converted + 1.05 * OCPI_converted) + 1.05 * SOAS_converted + \
                 DST1_converted + 0.30 *  DST2_converted + 1.86 * SALA_converted


diff = total_pm25_cal - total_PM25
diff_perc = (total_pm25_cal - total_PM25)/total_PM25
print('diff_pm25_perc_max = %.2f' % (np.nanmax(diff_perc)))
print('diff_pm25_perc_min = %.2f' % (np.nanmin(diff_perc)))
print('diff_pm25_perc_mean = %.2f' % (np.nanmean(diff_perc)))
print('----------------')
BC_cal=BCPI_converted+BCPO_converted
diff_bc = BC_cal-BC
diff_bc_perc = diff_bc/BC_cal
print('diff_BC_perc_max = %.2f' % (np.nanmax(diff_bc_perc)))
print('diff_BC_perc_min = %.2f' % (np.nanmin(diff_bc_perc)))
print('diff_BC_perc_mean = %.2f' % (np.nanmean(diff_bc_perc)))
print('----------------')

diff_SO4 = SO4_converted-SO4
diff_SO4_perc = diff_SO4/SO4_converted
print('diff_SO4_perc_max = %.2f' % (np.nanmax(diff_SO4_perc)))
print('diff_SO4_perc_min = %.2f' % (np.nanmin(diff_SO4_perc)))
print('diff_SO4_perc_mean = %.2f' % (np.nanmean(diff_SO4_perc)))
print('----------------')

diff_NH4 = NH4_converted-NH4
diff_NH4_perc = diff_NH4/NH4_converted
print('diff_NH4_perc_max = %.2f' % (np.nanmax(diff_NH4_perc)))
print('diff_NH4_perc_min = %.2f' % (np.nanmin(diff_NH4_perc)))
print('diff_NH4_perc_mean = %.2f' % (np.nanmean(diff_NH4_perc)))
print('----------------')

diff_NIT = NIT_converted-NIT
diff_NIT_perc = diff_NIT/NIT_converted
print('diff_NIT_perc_max = %.2f' % (np.nanmax(diff_NIT_perc)))
print('diff_NIT_perc_min = %.2f' % (np.nanmin(diff_NIT_perc)))
print('diff_NIT_perc_mean = %.2f' % (np.nanmean(diff_NIT_perc)))