#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Mon Jul  3 10:58:28 EDT 2023
if __name__ == '__main__':
    import sys,os
    from misc.timer import Timer
    tt = Timer(f'[{os.getcwd()}] 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
wython = '/tigress/wenchang/wython'
if wython not in sys.path: sys.path.append(wython); print('added to python path:', wython)
#from misc import get_kws_from_argv
#
if __name__ == '__main__':
    tt.check('end import')
#
#start from here
daname = 't_surf'
idir = '/tigress/wenchang/analysis/FLOR'
ofile = __file__.replace('.py', '.nc')
if not os.path.exists(ofile) or 'od' in sys.argv:
    #zonalmean
    ifile = os.path.join(idir, 'zonalmean', 't_surf_FLOR_co2x2_CTL1860_tigercpu_intelmpi_18_576PE_0101-0600_zonalmean.nc')
    print(ifile)
    da = xr.open_dataarray(ifile)
    ifile_ctl = os.path.join(idir, 'zonalmean', 't_surf_FLOR_CTL1860_newdiag_tigercpu_intelmpi_18_576PE_0101-0200_zonalmean.nc')
    print(ifile_ctl)
    da_ctl = xr.open_dataarray(ifile_ctl)
    #anomaly
    daa = da.groupby('time.month') - da_ctl.groupby('time.month').mean('time')
    daa_zonalmean = daa
     
    #zonaoceanlmean
    ifile = os.path.join(idir, 'zonaloceanmean', 't_surf_FLOR_co2x2_CTL1860_tigercpu_intelmpi_18_576PE_0101-0600_zonaloceanmean.nc')
    print(ifile)
    da = xr.open_dataarray(ifile)
    ifile_ctl = os.path.join(idir, 'zonaloceanmean', 't_surf_FLOR_CTL1860_newdiag_tigercpu_intelmpi_18_576PE_0101-0200_zonaloceanmean.nc')
    print(ifile_ctl)
    da_ctl = xr.open_dataarray(ifile_ctl)
    #anomaly
    daa = da.groupby('time.month') - da_ctl.groupby('time.month').mean('time')
    daa_zonaloceanmean = daa
    
    ds = xr.Dataset({f'{daname}_zonalmean': daa_zonalmean, f'{daname}_zonaloceanmean': daa_zonaloceanmean})
    print(ds)
    ds.to_netcdf(ofile)
    print('[saved]:', ofile)
else:
    ds = xr.open_dataset(ofile)
    print(ds)
    print('[loaded]:', ofile)
 
#convert from lat to sin(lat)
sinlat = np.sin(ds.lat.values/180*np.pi)
ds  = ds.rename(lat='sinlat').assign_coords(sinlat=sinlat)
print(ds)

def wyplot(ax=None):
    if ax is None: fig,ax = plt.subplots()
    color = 'C0'
    tspan = slice('0101', '0120')
    daa = ds[f'{daname}_zonalmean']
    ls = '-'
    daa.sel(time=tspan).mean('time').plot(ax=ax, color=color, ls=ls,  label='zonal mean, years 1-20')
    daa = ds[f'{daname}_zonaloceanmean']
    ls  = '--'
    daa.sel(time=tspan).mean('time').plot(ax=ax, color=color, ls=ls, label='ocean zonal mean, years 1-20')
    
    color = 'C1'
    tspan = slice('0231', '0250')
    daa = ds[f'{daname}_zonalmean']
    ls = '-'
    daa.sel(time=tspan).mean('time').plot(ax=ax, color=color, ls=ls,  label='zonal mean, years 131-150')

    daa = ds[f'{daname}_zonaloceanmean']
    ls  = '--'
    daa.sel(time=tspan).mean('time').plot(ax=ax, color=color, ls=ls,  label='ocean zonal mean, years 131-150')

    color = 'C3'
    tspan = slice('0581', '0600')
    daa = ds[f'{daname}_zonalmean']
    ls = '-'
    daa.sel(time=tspan).mean('time').plot(ax=ax, color=color, ls=ls,  label=f'zonal mean, years 481-500')

    daa = ds[f'{daname}_zonaloceanmean']
    ls  = '--'
    daa.sel(time=tspan).mean('time').plot(ax=ax, color=color, ls=ls,  label='ocean zonal mean, years 481-500')

    ax.axvspan(-0.5, 0.5, color='gray', alpha=0.1)
    ax.legend()
    #lats = range(-90, 91, 30)
    lats = [-90, -60, -45, -30, -15, 0, 15, 30, 45, 60, 90]
    sinlats = [np.sin(lat/180*np.pi) for lat in lats]
    def latlabel(lat):
        if lat > 0:
            return f'{lat}N'
        elif lat <0: 
            return f'{-lat}S'
        else:
            return f'{lat}'
    labels = [latlabel(lat) for lat in lats]
    ax.set_xticks(sinlats)
    ax.set_xticklabels(labels)
    ax.set_xlim(-1,1)
    ax.set_xlabel('')
    ax.set_ylabel('Ts anomaly [K]')
if __name__ == '__main__':
    from wyconfig import * #my plot settings
    wyplot()

    
    #savefig
    if 'savefig' in sys.argv or 's' in sys.argv:
        figname = __file__.replace('.py', f'.png')
        if 'pdf' in sys.argv: figname = figname.replace('.png', '.pdf')
        if 'overwritefig' in sys.argv or 'o' in sys.argv:
            wysavefig(figname, overwritefig=True)
        else:
            wysavefig(figname)
    tt.check(f'**Done**')
    print()
    if 'notshowfig' in sys.argv or 'n' in sys.argv:
        pass
    else:
        if 'plt' in globals(): plt.show()