#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Wed Aug 19 21:58:56 EDT 2020
"""
README
forecast q2m anomaly given t2m anomaly and model hindcast climatology if RH uses ERA5 climatology.
ERA5 climatology surface pressure ps is also used to convert vapor pressure to specific humidity.
ERA5 climatology specific humidity q2m is used to get full specific humidity forecast (ERA5 q2m_clim + forecast q2m anomaly)
"""
if __name__ == '__main__':
    from misc.timer import Timer
    tt = Timer(f'start {__file__}')
import sys, os.path, os, glob
import xarray as xr, numpy as np, pandas as pd
#import matplotlib.pyplot as plt
#more imports
#
if __name__ == '__main__':
    tt.check('end import')
#
#start from here
nmme_dir = '/tigress/wenchang/data/NMME'
#model = 'GFDL_FLOR'
years_validate = [2019,]

def ta2q(T, RH, p):
    '''calculate specific humidity q (kg/kg) given air temperature T (K), 
    relative humidity RH (0-1) and air pressure p (Pa).

    steps:
    1) T -> es: es = es0*exp(a*(T-T0)/(T-T0+Tb) 
        where es0=610.94pa, a=17.625, T0=273.15, Tb=243.04
        see Magnus equation in https://en.wikipedia.org/wiki/Vapour_pressure_of_water
    2) es -> e: e = es*RH
    3) e -> q: q =  (Rd/Rv)*e/( p - (1-Rd/Rv)*e )'''
    # physical parameters
    es0 = 610.94 # Pa
    a = 17.625
    T0 = 273.15 # K
    Tb = 243.04 # K
    Rd, Rv = 287, 461 # J/kg/K

    # 3 steps to calculate specific humidity
    es = es0*np.exp(a*(T-T0)/(T-T0+Tb))
    #print(es)
    try:
        e = es.groupby('time.month')*RH
        e = e.drop('month')
    except:
        e = es*RH
    try:
        q = (Rd/Rv)*e/(p - ((1-Rd/Rv)*e).groupby('time.month'))
        q = q.drop('month')
    except:
        q = (Rd/Rv)*e/(p - (1-Rd/Rv)*e)

    return q

# obs q2m climatology
ifile_q2mclim = '/tigress/wenchang/data/era5/analysis_wy/2m_specific_humidity/monthly/clim/era5.2m_specific_humidity.mclim.1982-2010.nmme.nc'
q2mclim = xr.open_dataarray(ifile_q2mclim) # units is kg/kg

# obs RH2m
ifile_RH = '/tigress/wenchang/data/era5/analysis_wy/2m_relative_humidity/monthly/clim/era5.2m_relative_humidity.mclim.1982-2010.nmme.nc'
RH = xr.open_dataarray(ifile_RH)/100 # units is % -> 0-1

# obs surface pressure
ifile_ps = '/tigress/wenchang/data/era5/analysis_wy/surface_pressure/monthly/clim/era5.surface_pressure.mclim.1982-2010.nmme.nc'
ps = xr.open_dataarray(ifile_ps) # units is Pa

# climatology t2m
ifile_clim = '/tigress/wenchang/data/era5/analysis_wy/2m_temperature/monthly/clim/era5.2m_temperature.mclim.1982-2010.nmme.nc'
clim = xr.open_dataset(ifile_clim)['t2m']

for year in years_validate:
    print(year)
    ofile = f'q2m.validate.{year:04d}.nc'
    if os.path.exists(ofile):
        print('[exists]:', ofile)
        continue

    # true q2m in validation year
    ifile_q2m = f'/tigress/wenchang/data/era5/analysis_wy/2m_specific_humidity/monthly/era5.2m_specific_humidity.monthly.{year:04d}.nc'
    da = xr.open_dataarray(ifile_q2m)
    q2m_true = da.rename(longitude='lon', latitude='lat').sel(lon=range(0,360), lat=range(-90,91))

    # t2m in validation year
    ifile_fcst = f'/tigress/wenchang/data/era5/analysis_wy/2m_temperature/monthly/era5.2m_temperature.monthly.{year:04d}.nc'
    da = xr.open_dataarray(ifile_fcst)
    fcst = da.rename(longitude='lon', latitude='lat').sel(lon=range(0,360), lat=range(-90,91))

    
    # forecast q2m anomaly if RH keeps the same as the obs monthly climatology
    q2m_a = ta2q(fcst, RH, ps).groupby('time.month') - ta2q(clim, RH, ps)
    q2m_a = q2m_a.assign_attrs(units='kg/kg', long_name='2m specific humidity anomaly')
     
    # obs q2m monthly climatology plus forecast q2m anomaly under obs constant RH
    q2m = q2mclim + q2m_a.groupby('time.month')
    q2m = q2m.assign_attrs(units='kg/kg', long_name='2m specific humidity')
    
    # save data
    ds = xr.Dataset(dict(q2m=q2m, anom=q2m_a, q2m_true=q2m_true, anom_true=q2m_true.groupby('time.month')-q2mclim))
    ds.to_netcdf(ofile)
    print('[saved]:', ofile)
    print()
    

if __name__ == '__main__':
    #from wyconfig import * #my plot settings
    
    tt.check(f'**Done**')
    #plt.show()