#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Fri Jan 14 15:48:37 EST 2022
if __name__ == '__main__':
    import sys
    from misc.timer import Timer
    tt = Timer('start ' + ' '.join(sys.argv))
import sys, os.path, os, glob, datetime
import xarray as xr, numpy as np, pandas as pd, matplotlib.pyplot as plt
#more imports
import xlinregress
#
if __name__ == '__main__':
    tt.check('end import')
#
#start from here
ifile = 'PI_historical_FMA-ASO_ens417_1982-2014.nc'
da_hist = xr.open_dataarray(ifile)
ifile = 'PI_hist-nat_FMA-ASO_ens167_1982-2014.nc'
da_nat = xr.open_dataarray(ifile)
ifile = 'PI_hist-GHG_FMA-ASO_ens109_1982-2014.nc'
da_ghg = xr.open_dataarray(ifile)
model_members = [m for m in da_ghg.model_member.values 
    if m in da_nat.model_member.values
    and m in da_hist.model_member.values]
season = ifile.split('_')[2]

da = da_hist.sel(model_member=model_members)
models = [model_member.split('_')[0] for model_member in da.model_member.values]
models = xr.DataArray(models, dims='model_member', coords=[da.model_member.values])
da = da.assign_coords(model=models)
#rg = da.mean('model_member').linregress.on(da.year)
rg = da.groupby(da.model).mean('model_member').mean('model').linregress.on(da.year)
lons = list(rg.lon.values)
lats = list(rg.lat.values)

#rg1 = da.linregress.on(da.year)
#rg1 = rg1.assign_coords(model=models)
#rg1 = rg1.groupby(rg.model).mean('model_member').mean('model')

#era5
ifile = '/tigress/wenchang/data/era5/analysis_wy/TCI_v202201/wyanalysis/era5.yearly.PI_2d.vmax.FMA-ASO.1979-2020.nc'
da = xr.open_dataarray(ifile).sel(year=slice(1982,2014))
rg_era5 = da.linregress.on(da.year)

 
if __name__ == '__main__':
    from wyconfig import * #my plot settings
    fig, ax = plt.subplots()
    da = (rg_era5.slope.sel(lat=lats, lon=lons) - rg.slope).assign_attrs(units='m/s/year', long_name=f'ERA5 $-$ CMIP6 mmm {season} PI vmax linear trend over 1982-2014')
    da.plot.contourf(levels=np.arange(-0.4, 0.41, 0.04), extend='both', ax=ax)

    latmin, latmax = 10, 30
    lonmin, lonmax = 40, 360-20
    ax.plot([lonmin, lonmax, lonmax, lonmin, lonmin], [latmin, latmin, latmax, latmax, latmin], color='gray')
    latmin, latmax = -30, -10
    lonmin, lonmax = 30, 360-150
    ax.plot([lonmin, lonmax, lonmax, lonmin, lonmin], [latmin, latmin, latmax, latmax, latmin], color='gray')

    
    #savefig
    if len(sys.argv)>1 and 'savefig' in sys.argv[1:]:
        figname = __file__.replace('.py', f'.png')
        if 'overwritefig' in sys.argv[1:]:
            wysavefig(figname, overwritefig=True)
        else:
            wysavefig(figname)
    tt.check(f'**Done**')
    plt.show()