#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Tue Mar 28 19:47:00 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',
'precip',
][-1]
model = 'AM2.5'
#tspan = slice('0201', '0250')
tspan = slice('0101', '0150')
ifiles = """
t_surf_AM2.5_CTL1990s_plus2K_tigercpu_intelmpi_18_540PE_0101-0150_zonalmean.nc
t_surf_AM2.5_CTL1990s_2xCO2_tigercpu_intelmpi_18_540PE_0101-0150_zonalmean.nc
t_surf_AM2.5_CTL1990s_plus2K_2xCO2_tigercpu_intelmpi_18_540PE_0101-0150_zonalmean.nc
""".split()
ifile_ctl = 't_surf_AM2.5_CTL1990s_tigercpu_intelmpi_18_540PE_0101-0200_zonalmean.nc'

ifiles = [ifile.replace('t_surf', daname) for ifile in ifiles]
ifile_ctl = ifile_ctl.replace('t_surf', daname)
labels = [ifile.split('AM2.5_CTL1990s_')[-1].split('_tigercpu')[0] for ifile in ifiles]

das = dict()
da_ctl = xr.open_dataarray(ifile_ctl).load().mean('time')
for label,ifile in zip(labels, ifiles):
    da = xr.open_dataarray(ifile).load().sel(time=tspan).mean('time')
    daa = da - da_ctl
    #weighted lat mean
    #wgt = np.cos(daa.lat/180*np.pi)
    #daa_glbmean = daa.weighted(wgt).mean('lat')
    #daa = daa/daa_glbmean
    #lattrop = slice(-30,30)
    #daa_tropmean = daa.sel(lat=lattrop).weighted(wgt.sel(lat=lattrop)).mean('lat')
    #tropmean t_surf
    #ifile_tropmean = ifile.replace(daname, f'../tropmean/t_surf').replace('zonalmean', 'tropmean')
    #ifile_ctl_tropmean = ifile_ctl.replace(daname, f'../tropmean/t_surf').replace('zonalmean', 'tropmean')
    #daa_tropmean = xr.open_dataarray(ifile_tropmean).sel(time=tspan).mean('time') \
    #    - xr.open_dataarray(ifile_ctl_tropmean).mean('time')
    #daa = daa/daa_tropmean
    das[label] = daa
    
if __name__ == '__main__':
    from wyconfig import * #my plot settings
    from geoplots import xticks2lat
    fig,ax = plt.subplots()
    for label in labels:
        ls = None#'--' if 'm' in label or 'xp' in label else '-'
        color = None #'C0'
        lw = 1.5 #1 if 'x2' in label or 'xp5' in label or '2p0' in label else 1.5
        das[label].plot(label=label, ls=ls, color=color, lw=lw)
    (das['2xCO2'] + das['plus2K']).plot(label='plus2K + 2xCO2', ls='--', color=color, lw=lw)

    ax.axvspan(-30, 30, color='gray', alpha=0.1)
    ax.legend()
    ax.set_xticks(range(-90,91,15))
    if daname in ('precip',):
        ax.set_ylabel(f'{model} precip anom[mm/day]')
    else:
        ax.set_ylabel(f'{model} Ts anom [K]')
    xticks2lat(ax=ax)
    ax.set_xlabel('')
    #ax.text(1, 1, f'{tspan.start}-{tspan.stop}', ha='right', va='top', transform=ax.transAxes)
    #ax.text(1, 1, model, ha='right', va='top', transform=ax.transAxes)
    if daname in ('t_surf',): 
        ax.axhline(0, color='gray', ls='--')
    elif daname in ('precip',): 
        ax.axhline(0, color='gray', ls='--')
        #ax.set_ylim(-0.1, 0.4) #mm/day/K
    
    #savefig
    if 'savefig' in sys.argv or 's' in sys.argv:
        figname = __file__.replace('.py', f'_{daname}.png')
        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()