#!/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
co2_only = True
daname = [
'precip',
't_surf',
][-1]
model = 'FLOR'
tspan = slice('0201', '0250')
#tspan = slice('0101', '0150')
ifiles = """
t_surf_FLOR_co2x4_CTL1860_tigercpu_intelmpi_18_576PE_0101-0250_zonaloceanmean.nc
t_surf_FLOR_co2xp5_CTL1860_tigercpu_intelmpi_18_576PE_0101-0600_zonaloceanmean.nc
t_surf_FLOR_co2xp25_CTL1860_tigercpu_intelmpi_18_576PE_0101-0250_zonaloceanmean.nc
t_surf_FLOR_p2p0sol_CTL1860_tigercpu_intelmpi_18_576PE_0101-0400_zonaloceanmean.nc
t_surf_FLOR_p6p0sol_CTL1860_tigercpu_intelmpi_18_576PE_0101-0600_zonaloceanmean.nc
t_surf_FLOR_m2p0sol_CTL1860_tigercpu_intelmpi_18_576PE_0101-0400_zonaloceanmean.nc
t_surf_FLOR_m6p0sol_CTL1860_tigercpu_intelmpi_18_576PE_0101-0597_zonaloceanmean.nc
t_surf_FLOR_co2x2_CTL1860_tigercpu_intelmpi_18_576PE_0101-0600_zonaloceanmean.nc
""".split()[-1:]
ifile_ctl = 't_surf_FLOR_CTL1860_newdiag_tigercpu_intelmpi_18_576PE_0101-0200_zonaloceanmean.nc'

ifiles = [ifile.replace('t_surf', daname) for ifile in ifiles]
ifile_ctl = ifile_ctl.replace('t_surf', daname)
labels = [ifile.split('FLOR_')[-1].split('_')[0] for ifile in ifiles]

das = dict()
da_ctl = xr.open_dataarray(ifile_ctl).load().mean('time')
das_eGMST = dict()
das_glbmean = dict()
das_tropmean = dict()
for label,ifile in zip(labels, ifiles):
    da = xr.open_dataarray(ifile).load().sel(time=tspan).mean('time')
    daa = da - da_ctl
    wgt = np.cos(daa.lat/180*np.pi)
    #glbmean
    daa_glbmean = daa.weighted(wgt).mean('lat')
    lattrop = slice(-30,30)
    #weighted lat mean
    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('zonaloceanmean', 'tropmean')
    ifile_ctl_tropmean = ifile_ctl.replace(daname, f'../tropmean/t_surf').replace('zonaloceanmean', 'tropmean')
    daa_tropmean = xr.open_dataarray(ifile_tropmean).sel(time=tspan).mean('time') \
        - xr.open_dataarray(ifile_ctl_tropmean).mean('time')
    #daa = daa/daa_glbmean
    daa = daa/daa_tropmean
    #glbmean precip
    ifile_pr = ifile.replace(daname, f'precip')
    da_pr = xr.open_dataarray(ifile_pr).sel(time=tspan).mean('time').weighted(wgt).mean('lat')
    ifile_pr_ctl = ifile_ctl.replace(daname, f'precip')
    da_pr_ctl = xr.open_dataarray(ifile_pr_ctl).mean('time').weighted(wgt).mean('lat')
    daa_pr = da_pr/da_pr_ctl*100 - 100 #units %
    eGMST = (daa_pr + 2.13)/3.28 # equivalent uniform GMST change
    """

    das[label] = daa
    das_glbmean[label] = daa_glbmean
    das_tropmean[label] = daa_tropmean
    #das_eGMST[label] = eGMST
    
if __name__ == '__main__':
    from wyconfig import * #my plot settings
    from geoplots import xticks2lat
    fig,ax = plt.subplots()
    for label in labels:
        #ls = '--' if 'm' in label or 'xp' in label else '-'
        #color = 'C0' if 'co2' in label else 'C1'
        #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)
        #ax.axhline(das_glbmean[label], label='GMST anom', color='C1')
        #ax.axhline(das_tropmean[label], label='TROPMST anom', color='C2')
        #ax.axhline(das_eGMST[label], label='equivalent GMST anom', color='C3', ls='--')

    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'precip anom [mm/day]')
    else:
        ax.set_ylabel(f'SST  anom [K]')
    xticks2lat(ax=ax)
    ax.set_xlabel('')
    ax.text(1, 1, f'{tspan.start}-{tspan.stop}', ha='right', va='top', transform=ax.transAxes)
    #if daname in ('t_surf',): ax.axhline(1, color='gray', ls='--')
    #if daname in ('precip',): ax.axhline(0, color='gray', ls='--')
    
    #savefig
    if 'savefig' in sys.argv or 's' in sys.argv:
        figname = __file__.replace('.py', f'__{daname}.png')
        if co2_only: figname = figname.replace('.png', '_co2only.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()