#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Thu Apr 28 14:58:14 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
#
if __name__ == '__main__':
    tt.check('end import')
#
#start from here
func = lambda x: x.load().groupby('time.year').mean('time')
daname = 'netrad_toa'
#daname = 't_surf'
if len(sys.argv)>1:
    daname = sys.argv[1]
if daname in ('t_surf',):
    units = 'K'
elif daname in ('netrad_toa',):
    units = 'W/m$^2$'
ifile = f'{daname}_FLOR_CTL1990_v201905_tigercpu_intelmpi_18_576PE_0101-0300_glbmean.nc'
da0 = xr.open_dataarray(ifile).pipe(func)
label0 = 'CTL1990s'
#ifiles = sorted( glob.glob(f'{daname}_FLOR_CTL1990_v201905_2xCO2_tigercpu_intelmpi_18_576PE_????-????_glbmean.nc') )
#da1 = xr.open_mfdataset(ifiles)[daname].pipe(func)
ifile = f'{daname}_FLOR_CTL1990_v201905_2xCO2_tigercpu_intelmpi_18_576PE_0101-0300_glbmean.nc'
da1 = xr.open_dataarray(ifile).pipe(func)
label1 = f'2xCO2, {da1.year.size} years now'
#ifiles = sorted( glob.glob(f'{daname}_FLOR_CTL1990_v201905_2xCO2_m1p7pctSolar_tigercpu_intelmpi_18_576PE_????-????_glbmean.nc') )
#da2 = xr.open_mfdataset(ifiles)[daname].pipe(func)
ifile = f'{daname}_FLOR_CTL1990_v201905_2xCO2_m1p7pctSolar_tigercpu_intelmpi_18_576PE_0101-0300_glbmean.nc'
da2 = xr.open_dataarray(ifile).pipe(func)
label2 = f'2xCO2 and 1.7% solar reduction, {da2.year.size} years now'

ifiles = sorted( glob.glob(f'{daname}_FLOR_CTL1990_v201905_2xCO2_m2pctSolar_tigercpu_intelmpi_18_576PE_????-????_glbmean.nc') )
da3 = xr.open_mfdataset(ifiles)[daname].pipe(func)
label3 = f'2xCO2 and 2% solar reduction, {da3.year.size} years now'

tag = f'_{daname}'
 
 
if __name__ == '__main__':
    from wyconfig import * #my plot settings
    #ratio of 2xCO2 to 1pctSolar
    #years = slice(101,150)
    
    da0.plot(label=label0)
    da1.plot(label=label1)
    da2.assign_attrs(units=units).plot(label=label2)
    da3.assign_attrs(units=units).plot(label=label3)

    plt.legend()
    plt.title(f'FLOR 1990 {daname}')
    

    #savefig
    if len(sys.argv)>1 and 'savefig' in sys.argv[1:]:
        figname = __file__.replace('.py', f'{tag}.png')
        if 'overwritefig' in sys.argv[1:]:
            wysavefig(figname, overwritefig=True)
        else:
            wysavefig(figname)
    tt.check(f'**Done**')
    print()
    if 'notshowfig' in sys.argv:
        pass
    else:
        plt.show()