#!/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 misc.modelout_old import get_modelout_data, get_modelout_files
import xfilter
nwindow, dimlp = 9, 'year'
lowpass = lambda x: x.filter.lowpass(1/nwindow, dim=dimlp, padtype='odd')
#
if __name__ == '__main__':
    tt.check('end import')
#
#start from here
model = 'CM2.1p1'
daname = 't_surf'
func = lambda x: x.load().geo.fldmean()
funcname = 'gmst'

def get_run_files(model, expname, dsname='atmos_month', ens=None, loc='modelout'):
    """get list of model output files given daname, model, expname, and optionally, dsname, years, ens. """
    #daname = 'netrad_toa'
    #dsname = 'atmos_month'
    #model = 'AM2.5'
    #expname = 'CTL1990s_tigercpu_intelmpi_18_540PE'
    if ens is None:
        idir = os.path.join('/tigress/wenchang/MODEL_OUT', model, expname, 'POSTP')
    else:
        idir = os.path.join('/tigress/wenchang/MODEL_OUT', model, expname, f'en{ens:02d}', 'POSTP')
    if model in ('FLOR',):
        idir = idir.replace(f'{model}/', '')
    if loc == 'modelout':
        if not os.path.exists(idir):
            return
    elif loc == 'scratch': #data not in modelout dir; maybe in work dir
        idir = idir.replace('/tigress/wenchang/MODEL_OUT', '/home/wenchang/scratch') \
            .replace(f'{expname}', f'work/{expname}')
        if model in ('FLOR',):
            idir = idir.replace('home/wenchang/scratch', 'home/wenchang/scratch/FLOR')
        if ens is not None:
            idir = idir.replace(f'en{ens:02d}/', '') \
                .replace('_tigercpu_', f'_e{ens}_tigercpu_')
        if not os.path.exists(idir): #data not in work dir either
            idir = idir.replace('home/wenchang/scratch','home/wenchang/sGEOCLIM/wenchang')
            if not os.path.exists(idir): #data not in GEOCLIM scratch either
                return
    #get the year range based on all files under the POSTP dir
    ncfiles = [ncfile for ncfile in os.listdir(idir) if not ncfile.startswith('.') and ncfile.endswith('.nc')]
    ncfiles.sort()
    year0, year1 = int(ncfiles[0][:4]), int(ncfiles[-1][:4])
    years = range(year0, year1+1)
    ifiles = [os.path.join(idir, f'{year:04d}0101.{dsname}.nc') for year in years]

    return ifiles, years


def get_run_data(daname, model, expname, func, funcname, loop_files=True):
    """get data from the still running experiment"""
    print(daname, model, expname, funcname)
    s = f'{daname}_{model}_{expname}_????-????_{funcname}.nc'
    ifiles = glob.glob(s)
    #print(ifiles); sys.exit()
    if ifiles:
        ifiles.sort()
        da = xr.open_mfdataset(ifiles)[daname].load()
        year_end_cached = da.time.dt.year[-1].item()
        da_ = da
    else:
        da = get_modelout_data(daname=daname, model=model, expname=expname, func=func, funcname=funcname, loop_files=loop_files)
        return da
    _, years = get_modelout_files(model=model, expname=expname)
    if year_end_cached >= years[-1]:
        da = da_
        print('[loaded]:', s)
    else:
        da = get_modelout_data(daname=daname, model=model, expname=expname, func=func, funcname=funcname, loop_files=True, 
            years=range(year_end_cached+1, years[-1]+1))
        da = xr.concat([da_, da], dim='time')
    print()
    return da

expname = 'CTL1860_tigercpu_intelmpi_18_80PE'
da = get_modelout_data(daname=daname, model=model, expname=expname, func=func, funcname=funcname, loop_files=True)
da_ctl = da.sel(time=slice('0100', '0500'))

"""
#co2 out of model range
expname = 'CTL1860_8xCO2_tigercpu_intelmpi_18_80PE'
da = get_run_data(daname=daname, model=model, expname=expname, func=func, funcname=funcname, loop_files=True)
da_8xco2 = da
"""
 
expname = 'CTL1860_4xCO2_tigercpu_intelmpi_18_80PE'
da = get_run_data(daname=daname, model=model, expname=expname, func=func, funcname=funcname, loop_files=True)
da_4xco2 = da
 
expname = 'CTL1860_2xCO2_tigercpu_intelmpi_18_80PE'
da = get_run_data(daname=daname, model=model, expname=expname, func=func, funcname=funcname, loop_files=True)
da_2xco2 = da
 
expname = 'CTL1860_p5xCO2_tigercpu_intelmpi_18_80PE'
da = get_run_data(daname=daname, model=model, expname=expname, func=func, funcname=funcname, loop_files=True)
da_p5xco2 = da
 
expname = 'CTL1860_p25xCO2_tigercpu_intelmpi_18_80PE'
da = get_run_data(daname=daname, model=model, expname=expname, func=func, funcname=funcname, loop_files=True)
da_p25xco2 = da
 
expname = 'CTL1860_p125xCO2_tigercpu_intelmpi_18_80PE'
da = get_run_data(daname=daname, model=model, expname=expname, func=func, funcname=funcname, loop_files=True)
da_p125xco2 = da
 
 
expname = 'CTL1860_p6pctSolar_tigercpu_intelmpi_18_80PE'
da = get_run_data(daname=daname, model=model, expname=expname, func=func, funcname=funcname)
da_p6pctSolar = da
 
expname = 'CTL1860_p4pctSolar_tigercpu_intelmpi_18_80PE'
da = get_run_data(daname=daname, model=model, expname=expname, func=func, funcname=funcname)
da_p4pctSolar = da
 
expname = 'CTL1860_p2pctSolar_tigercpu_intelmpi_18_80PE'
da = get_run_data(daname=daname, model=model, expname=expname, func=func, funcname=funcname)
da_p2pctSolar = da
 
expname = 'CTL1860_p1pctSolar_tigercpu_intelmpi_18_80PE'
da = get_run_data(daname=daname, model=model, expname=expname, func=func, funcname=funcname)
da_p1pctSolar = da
 
expname = 'CTL1860_m1pctSolar_tigercpu_intelmpi_18_80PE'
da = get_run_data(daname=daname, model=model, expname=expname, func=func, funcname=funcname)
da_m1pctSolar = da
 
expname = 'CTL1860_m2pctSolar_tigercpu_intelmpi_18_80PE'
da = get_run_data(daname=daname, model=model, expname=expname, func=func, funcname=funcname)
da_m2pctSolar = da
 
expname = 'CTL1860_m4pctSolar_tigercpu_intelmpi_18_80PE'
da = get_run_data(daname=daname, model=model, expname=expname, func=func, funcname=funcname)
da_m4pctSolar = da
 
expname = 'CTL1860_m6pctSolar_tigercpu_intelmpi_18_80PE'
da = get_run_data(daname=daname, model=model, expname=expname, func=func, funcname=funcname)
da_m6pctSolar = da
 
"""
expname = 'CTL1860_p1pctSolar_tigercpu_intelmpi_18_80PE'
da = get_modelout_data_chunked(daname=daname, model=model, expname=expname, func=func, funcname=funcname)
da_p1pctSolar = da
"""
 
def wyplot(da, flip=False, **kws):
    ax = kws.pop('ax', plt.gca())
    da = da.groupby('time.year').mean('time') - da_ctl.groupby('time.year').mean('time').sel(year=slice(101,200)).mean('year') #anom
    if da.year[0] == 101:
        da = xr.concat([da_ctl.sel(time='0100').groupby('time.year').mean('time')*0, da], dim='year')
    #da = da.groupby('time.month') - da_ctl.sel(time=slice('0101','0200')).groupby('time.month').mean('time') #anom
    if flip: da = -da
    da = da.pipe(lowpass)
    da = da.assign_coords(year=da.year-100) #shift the year axis to start with 0 (instead of 100)
    da.plot(**kws)
    #ax.plot(da.isel(year=-1).year, da.isel(year=-1), marker='o', fillstyle='none', color='gray')
    #ax.text(da.year.values[-1], da.values[-1], f'{da.year.values[-1]}', color='gray')
    
 
if __name__ == '__main__':
    from wyconfig import * #my plot settings
    fig,ax = plt.subplots(figsize=(8,4))

    for da,label in [(da_ctl, 'ctl'), (da_4xco2, '4xCO2'), (da_2xco2, '2xCO2'),
        (da_p5xco2, '0.5xCO2'), (da_p25xco2, '0.25xCO2'), (da_p125xco2, '0.125xCO2'),
        (da_p6pctSolar, '+6% solar'), (da_p4pctSolar, '+4% solar'), (da_p2pctSolar, '+2% solar'), (da_p1pctSolar, '+1% solar'),
        (da_m1pctSolar, '-1% solar'), (da_m2pctSolar, '-2% solar'), (da_m4pctSolar, '-4% solar'), (da_m6pctSolar, '-6% solar')]:
        if label=='ctl':
            da.pipe(wyplot, label=label, color='k')
            ax.set_prop_cycle(None)
        else:
            da.pipe(wyplot, label=label)

    ax.legend(ncol=1)
    #ax.set_xlim(100, 250)
    ax.set_ylabel('CM2.1 GMST anom [K]')
    
    #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)

    #mirror the negative forcing
    fig,ax = plt.subplots(figsize=(8,4))

    da, label = da_ctl, 'ctl'
    da.pipe(wyplot, label=label, color='k')

    da, label = da_2xco2, '2xCO2'
    da.pipe(wyplot, label=label, color='C0', ls='-')
    da, label = da_p5xco2, '0.5xCO2'
    da.pipe(wyplot, label=label, color='C0', ls='--', flip=True)

    da, label = da_4xco2, '4xCO2'
    da.pipe(wyplot, label=label, color='C1', ls='-')
    da, label = da_p25xco2, '0.25xCO2'
    da.pipe(wyplot, label=label, color='C1', ls='--', flip=True)

    da, label = da_p125xco2, '0.125xCO2'
    da.pipe(wyplot, label=label, color='C2', ls='--', flip=True)

    da, label = da_p1pctSolar, '+1% solar'
    da.pipe(wyplot, label=label, color='C3', ls='-')
    da, label = da_m1pctSolar, '-1% solar'
    da.pipe(wyplot, label=label, color='C3', ls='--', flip=True)

    da, label = da_p2pctSolar, '+2% solar'
    da.pipe(wyplot, label=label, color='C4', ls='-')
    da, label = da_m2pctSolar, '-2% solar'
    da.pipe(wyplot, label=label, color='C4', ls='--', flip=True)

    da, label = da_p4pctSolar, '+4% solar'
    da.pipe(wyplot, label=label, color='C5', ls='-')
    da, label = da_m4pctSolar, '-4% solar'
    da.pipe(wyplot, label=label, color='C5', ls='--', flip=True)

    da, label = da_p6pctSolar, '+6% solar'
    da.pipe(wyplot, label=label, color='C6', ls='-')
    da, label = da_m6pctSolar, '-6% solar'
    da.pipe(wyplot, label=label, color='C6', ls='--', flip=True)

    ax.legend(ncol=1)
    #ax.set_xlim('0100', '0250')
    ax.set_ylabel(f'CM2.1 GMST anom [K], {nwindow}-{dimlp}-lp')
    
    #savefig
    if 'savefig' in sys.argv or 's' in sys.argv:
        figname = __file__.replace('.py', f'_mirror.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()