#!/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 = [
'precip',
't_surf',
][-1]
model = 'FLOR'
tspans = [slice('0221', '0240'), slice('0221', '0240'), slice('2121', '2140')]
#tspan = slice('0101', '0150')
ifiles = """
t_surf_FLOR_co2x2_CTL1860_tigercpu_intelmpi_18_576PE_0101-0600_zonalmean.nc
t_surf_FLOR_co2x4_CTL1860_tigercpu_intelmpi_18_576PE_0101-0250_zonalmean.nc
t_surf_FLOR_CTL1860_newdiag_1pct4xCO2_tigercpu_intelmpi_18_576PE_2001-2140_zonalmean.nc
""".split()

ifiles_ctl = """
t_surf_FLOR_CTL1860_newdiag_tigercpu_intelmpi_18_576PE_0101-0200_zonalmean.nc
t_surf_FLOR_CTL1860_newdiag_tigercpu_intelmpi_18_576PE_0101-0200_zonalmean.nc
t_surf_FLOR_CTL1860_newdiag_tigercpu_intelmpi_18_576PE_2001-2100_zonalmean.nc
""".split()

ifiles = [ifile.replace('t_surf', daname) for ifile in ifiles]
ifiles_ctl = [ifile.replace('t_surf', daname) for ifile in ifiles_ctl]
#labels = [ifile.split('FLOR_')[-1].split('_')[0] for ifile in ifiles]
labels = ['co2x2', 'co2x4', '1pct4xCO2']

das = dict()
#da_ctl = xr.open_dataarray(ifile_ctl).load().mean('time')
for label,ifile,ifile_ctl,tspan in zip(labels, ifiles, ifiles_ctl, tspans):
    da_ctl = xr.open_dataarray(ifile_ctl).load().mean('time')
    da = xr.open_dataarray(ifile).load().sel(time=tspan).mean('time')
    daa = da - da_ctl
    wgt = np.cos(daa.lat/180*np.pi)
    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('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_glbmean
    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 = '--' if 'm' in label or 'xp' in label else '-'
        color = None#'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.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 norm by trop mean Ts [mm/day/K]')
    else:
        ax.set_ylabel(f'Ts anom norm by 30S-30N mean')
    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, f'years 120-140 mean', 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 '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()