#!/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
import xfilter
#nwindow, dimlp = 1, 'year'
nwindow, dimlp = 5, 'month'
#lowpass = lambda x: x.filter.lowpass(1/nwindow, dim=dimlp, padtype='odd')
lowpass = lambda x: x.rolling(time=nwindow, center=True, min_periods=1).mean() if x.time.size>nwindow else x
import geoxarray
from misc.comp1samp import comp1samp
from modelout.getdata import funcs
#
if __name__ == '__main__':
    tt.check('end import')
#
#start from here
daname = 't_surf'
#func = lambda x: x.load().geo.fldmean()
funcname = 'indexEW' #'glbmean'
func = funcs[funcname]

volc = 'Novarupta' # 'StMaria'
year0 = 1912 #1902
imon_erupt = 6 -1 #10 - 1 #oct

das = []
labels = []

model = 'FLOR'
#ctl1860_tiger3
label = 'FLOR_ctl_1860_tg3'
expname = 'CTL1860_tiger3_intelmpi_24_1116PE'
da = update_modelout_data(daname=daname, model=model, expname=expname, func=func, funcname=funcname)#, years=range(100,201))
da_ctl = da

"""
#ctl1860_tiger3_FAtrop
label = 'FLOR_ctl_1860_tg3_FAtrop'
expname = 'CTL1860_FAtrop_tiger3_intelmpi_24_1116PE'
da = update_modelout_data(daname=daname, model=model, expname=expname, func=func, funcname=funcname)#, years=range(100,201))
labels.append(label)
das.append(da)
"""
exps = glob.glob(f'/scratch/gpfs/GEOCLIM/wenchang/tiger3/FLOR/work/CTL1860_volc{volc}_tiger3_intel24ifort_openmpi_1116PE_e*')
members = [int(e.split('_e')[-1]) for e in exps if os.listdir(f'{e}/POSTP')]
members.sort()
nens = len(members)
#for m in members: print(m)
#sys.exit()
das_ref = []
for m in members:
    expname = f'CTL1860_volc{volc}_tiger3_intel24ifort_openmpi_1116PE_e{m}'
    label = ''.join( expname.split('_tiger3_intel24ifort_openmpi_1116PE') )
    da = update_modelout_data(daname=daname, model=model, expname=expname, func=func, funcname=funcname)#, years=range(100,201))
    labels.append(label)
    das.append(da)

    nyears = da.time.size//12
    year_start_ctl = 101 + (m-1)*10
    year_end_ctl = year_start_ctl + nyears - 1
    da_ref = da_ctl.sel(time=slice(f'{year_start_ctl:04d}', f'{year_end_ctl:04d}'))
    #print(m, da_ref.time.values)
    das_ref.append(da_ref.assign_coords(time=da.time))
das = xr.concat(das, dim=pd.Index(members, name='ens'))
damean = das.mean('ens')
das_ref = xr.concat(das_ref, dim=pd.Index(members, name='ens'))
damean_ref = das_ref.mean('ens')
ncount = das.count(dim='ens') #number of ensemble members that have finished simulation for each month
#print(ncount); sys.exit()
ds = comp1samp(das-damean_ref, dim='ens')
err = ds.err

date_erupt = damean.time.values[imon_erupt]

def wyplot(da, m, label,**kws):
    nyears = da.time.size//12
    year_start_ctl = 101 + (m-1)*10
    year_end_ctl = year_start_ctl + nyears - 1
    da_ref = da_ctl.sel(time=slice(f'{year_start_ctl:04d}', f'{year_end_ctl:04d}'))
    da = da - damean_ref
    lw = 1; alpha = 0.3
    da.plot(lw=lw, alpha=alpha, **kws)
    
 
if __name__ == '__main__':
    from wyconfig import * #my plot settings
    #ctl
    fig,ax = plt.subplots(figsize=(8,4.5))
    for da,m,label in zip(das, members, labels):
        wyplot(da, m, label, ax=ax)
    #(damean - damean_ref + err).pipe(lowpass).plot(ax=ax, color='gray', ls='--')
    #(damean - damean_ref - err).pipe(lowpass).plot(ax=ax, color='gray', ls='--')
    ax.fill_between(damean.time.values, (damean - damean_ref - err).pipe(lowpass), (damean - damean_ref + err).pipe(lowpass), color='gray', alpha=0.3)
    (damean - damean_ref).pipe(lowpass).plot(ax=ax, color='gray', ls='--')
    (damean - damean_ref).pipe(lowpass).where(ncount==nens).plot(ax=ax, color='k', label=f'{volc}({date_erupt.year}-{date_erupt.month:02d}) {nens}ens mean', ls='-')
    ax.legend(ncol=1)
    ax.set_ylabel(f'{funcname} SSTA [degC]')
    #ax.set_xlim(1900, 3300)
    ax.axhline(0, color='gray', ls='--')
    ax.axvline(damean.time.values[imon_erupt], color='gray', ls='--')

    #savefig
    if 'savefig' in sys.argv or 's' in sys.argv:
        figname = __file__.replace('.py', f'.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()