#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Wed Oct  5 11:09:59 EDT 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
from modelout import get_modelout_data, update_modelout_data
from modelout.getdata import funcs
import xfilter
nwindow, dimlp = 9, 'year'
#lowpass = lambda x: x.filter.lowpass(1/nwindow, dim=dimlp, padtype='odd')
lowpass = lambda x: x.rolling(year=nwindow, center=True, min_periods=1).mean() if x.year.size>9 else x
import geoxarray
from misc import get_kws_from_argv
#
if __name__ == '__main__':
    tt.check('end import')
#
#start from here
years_ctl =  range(2001,2100+1)
#years_pert = range(2101,2120+1)
years_pert = range(2101,2140+1)
daname = get_kws_from_argv('daname', 't_surf')
#func = lambda x: x.load().geo.fldmean()
#funcname = 'glbmean'
#func = funcs['nino34']
#funcname = 'nino34'
func = funcs['annualmean']
funcname = 'annualmean'


das = []
labels = []
dc = {}

model = 'FLOR'
#ctl1860_v202407
label = 'FLOR_ctl_1860_v202407'
expname = 'CTL1860_v202407_tigercpu_intelmpi_18_576PE'
da = get_modelout_data(daname=daname, model=model, expname=expname, func=func, funcname=funcname, years=years_ctl)
labels.append(label)
das.append(da)
dc[label] = da
#ctl1860_v202407_FA
label = 'FLOR_ctl_1860_v202407_FA'
expname = 'CTL1860_v202407_FA_tigercpu_intelmpi_18_576PE'
da = get_modelout_data(daname=daname, model=model, expname=expname, func=func, funcname=funcname, years=years_ctl)
labels.append(label)
das.append(da)
dc[label] = da

#wy2024-08-06, shenzhen
#ctl1860_v202407_2xCO2
label = 'FLOR_ctl_1860_v202407_2xCO2'
expname = 'CTL1860_v202407_2xCO2_tigercpu_intelmpi_18_576PE'
da = get_modelout_data(daname=daname, model=model, expname=expname, func=func, funcname=funcname, years=years_pert)
labels.append(label)
das.append(da)
dc[label] = da
#ctl1860_v202407_1pct2xCO2
label = 'FLOR_ctl_1860_v202407_1pct2xCO2'
expname = 'CTL1860_v202407_1pct2xCO2_tigercpu_intelmpi_18_576PE'
da = get_modelout_data(daname=daname, model=model, expname=expname, func=func, funcname=funcname, years=years_pert)
labels.append(label)
das.append(da)
dc[label] = da
#ctl1860_v202407_FA_2xCO2
label = 'FLOR_ctl_1860_v202407_FA_2xCO2'
expname = 'CTL1860_v202407_FA_2xCO2_tigercpu_intelmpi_18_576PE'
da = get_modelout_data(daname=daname, model=model, expname=expname, func=func, funcname=funcname, years=years_pert)
labels.append(label)
das.append(da)
dc[label] = da
#ctl1860_v202407_FA_1pct2xCO2
label = 'FLOR_ctl_1860_v202407_FA_1pct2xCO2'
expname = 'CTL1860_v202407_FA_1pct2xCO2_tigercpu_intelmpi_18_576PE'
da = get_modelout_data(daname=daname, model=model, expname=expname, func=func, funcname=funcname, years=years_pert)
labels.append(label)
das.append(da)
dc[label] = da

#units
units = da.attrs['units']
 
def wyplot(da, label, ax=None, **kws):
    if ax is None: ax = plt.subplots()[1]
    da_ctl = das[1] if '_FA_' in label else das[0]
    #da = da.groupby('time.year').mean('time')
    daa = da.mean('time') - da_ctl.mean('time') #anomal from the original control
    #da = da.pipe(lowpass)
    #da = da.sel(year=slice(1900,2000))
    daa.assign_attrs(units=units).plot.contourf(ax=ax, **kws)

    year_start = years_pert[0] - 2000 # e.g. year 2001 is the first year (year 1)
    year_stop = years_pert[-1] - 2000 # e.g. year 2001 is year 1
    ax.set_title(f"{label.split('_v202407_')[-1]} {daname} anomaly: {year_start}-{year_stop}")

def wyplot_FA_minus_nonFA(case, ax=None, **kws):
    """response difference between nonFA and FA"""
    if ax is None: ax = plt.subplots()[1]
    if case == 'abrupt2xCO2':
        daa = ( dc['FLOR_ctl_1860_v202407_FA_2xCO2'].mean('time') - dc['FLOR_ctl_1860_v202407_FA'].mean('time') ) \
            - ( dc['FLOR_ctl_1860_v202407_2xCO2'].mean('time') - dc['FLOR_ctl_1860_v202407'].mean('time') )
    elif case == '1pct2xCO2':
        daa = ( dc['FLOR_ctl_1860_v202407_FA_1pct2xCO2'].mean('time') - dc['FLOR_ctl_1860_v202407_FA'].mean('time') ) \
            - ( dc['FLOR_ctl_1860_v202407_1pct2xCO2'].mean('time') - dc['FLOR_ctl_1860_v202407'].mean('time') )
    
    daa.assign_attrs(units=units).plot.contourf(ax=ax, **kws)
 
    year_start = years_pert[0] - 2000 # e.g. year 2001 is the first year (year 1)
    year_stop = years_pert[-1] - 2000 # e.g. year 2001 is year 1
    ax.set_title(f"{case} {daname} FA $-$ nonNA: {year_start}-{year_stop}")

if __name__ == '__main__':
    from wyconfig import * #my plot settings
    from geoplots import mapplot

    #wyplot
    if daname == 't_surf':
        vmax = 4
    elif daname == 'slp':
        vmax = 3
    elif daname == 'precip':
        vmax = 2 
    else:
        vmax = None
    if daname == 'precip':
        cmap = 'BrBG'
    else:
        cmap = None
    for da,label in zip(das[2:], labels[2:]):
        fig,ax = plt.subplots()
        wyplot(da, label, ax=ax, robust=True, levels=21, vmax=vmax, cmap=cmap, extend='both')
        plt.sca(ax)
        mapplot()


        #savefig
        if 'savefig' in sys.argv or 's' in sys.argv:
            figname = __file__.replace('.py', f'__{daname}_{label.split("_v202407_")[-1]}.png')
            if 'overwritefig' in sys.argv or 'o' in sys.argv:
                wysavefig(figname, overwritefig=True)
            else:
                wysavefig(figname)

    #wyplot_FA_minus_nonFA
    if daname == 't_surf':
        vmax = 2 
    elif daname == 'slp':
        vmax = 3
    elif daname == 'precip':
        vmax = 2 
    else:
        vmax = None
    for case in ('abrupt2xCO2', '1pct2xCO2'):
        fig,ax = plt.subplots()
        wyplot_FA_minus_nonFA(case, ax, robust=True, levels=21, vmax=vmax, cmap=cmap,extend='both')
        plt.sca(ax)
        mapplot()
        #savefig
        if 'savefig' in sys.argv or 's' in sys.argv:
            figname = __file__.replace('.py', f'__{daname}_FA_minus_nonFA_{case}.png')
            if 'overwritefig' in sys.argv or 'o' in sys.argv:
                wysavefig(figname, overwritefig=True)
            else:
                wysavefig(figname)

    #panels
    fig,axes = plt.subplots(1, 3, figsize=(12,3))
    vmax = 4 
    ax = axes[0]
    label = 'FLOR_ctl_1860_v202407_2xCO2'
    da = dc[label]
    wyplot(da, label, ax=ax, robust=True, levels=21, vmax=vmax, cmap=cmap, extend='both')

    ax = axes[1]
    label = 'FLOR_ctl_1860_v202407_FA_2xCO2'
    da = dc[label]
    wyplot(da, label, ax=ax, robust=True, levels=21, vmax=vmax, cmap=cmap, extend='both')

    ax = axes[2]
    vmax = 2
    case = 'abrupt2xCO2'
    wyplot_FA_minus_nonFA(case, ax, robust=True, levels=21, vmax=vmax, cmap=cmap,extend='both')
    ax.set_title('FA $-$ nonFA')

    for ax in axes:
        plt.sca(ax)
        mapplot()

    #savefig
    if 'savefig' in sys.argv or 's' in sys.argv:
        figname = __file__.replace('.py', f'__{daname}_FA_minus_nonFA_{case}_full.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:
        pass
    else:
        plt.show()