#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Tue Apr 28 14:44:33 EDT 2020
if __name__ == '__main__':
    from misc.timer import Timer
    t = 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
from shared import regions
city = 'Wuhan'
latx, lonx = regions[city]
#
if __name__ == '__main__':
    t.check('end import')
    
#start from here
plt.figure()
# climq from Gabe
ifile = 'merra2q/merra_weekly_clim.nc'
da_gabe = xr.open_dataset(ifile).WEEK_QV2M.rename(TT='time', LAT='lat', LON='lon')
da_gabe = da_gabe.assign_coords(time=range(1,53)).rename(time='week')
da_gabe.sel(lon=lonx, lat=latx, method='nearest').plot(label='Gabe weekly clim')

# climq from WY's recalculation
ifile = 'merra2q/merra2.QV2M.wclim1981-2017.nc'
da_wy = xr.open_dataset(ifile).QV2M
da_wy.sel(lon=lonx, lat=latx, method='nearest').plot(label='Wenchang weekly clim (MERRA2)')

# climq of new version from MERRA2
ifile = 'merra2q/merra2.QV2M.w1clim1981-2017.nc'
da2 = xr.open_dataset(ifile).QV2M
da2.sel(lon=lonx, lat=latx, method='nearest').plot(label='MERRA2 new-version weekly clim')

# climq of new version from ERA5
ifile = 'q_data/clim/era5.q2m.w1clim1981-2017.nc'
da5 = xr.open_dataset(ifile).q2m.rename(longitude='lon', latitude='lat')
da5.sel(lon=lonx, lat=latx, method='nearest').plot(label='ERA5 new-version weekly clim')

# climq from Gabe fixed
ifile = 'merra_weekly_climo.nc'
da_gabe = xr.open_dataset(ifile).WEEK_QV2M.rename(TWEEK='week', LAT='lat', LON='lon')
da_gabe = da_gabe.assign_coords(week=range(1,53))
da_gabe.sel(lon=lonx, lat=latx, method='nearest').plot(label='Gabe weekly clim fixed', ls='--')

plt.legend()
plt.tight_layout()
plt.title(city, loc='left')
plt.savefig(f'fig.q2m.{city}.wclim.comparison.png')

if __name__ == '__main__':
    t.check(f'end {__file__}')