#!/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.groupby('time.year').mean('time')
daname = 'netrad_toa'
ifile = f'{daname}_AM2.5_CTL1860_florCTL1860noleapSST_tigercpu_intelmpi_18_540PE_0101-0150_glbmean.nc'
da0 = xr.open_dataarray(ifile).pipe(func)
label0 = 'CTL1860'
ifile = f'{daname}_AM2.5_CTL1860_florCTL1860noleapSST_2xCO2_tigercpu_intelmpi_18_540PE_0101-0150_glbmean.nc'
da1 = xr.open_dataarray(ifile).pipe(func)
label1 = '2xCO2'
#ifile = f'{daname}_AM2.5_CTL1990s_solar1pctplus_tigercpu_intelmpi_18_540PE_0101-0150_glbmean.nc'
#da2 = xr.open_dataarray(ifile).pipe(func)
#label2 = '+1pct solar'
ifile = f'{daname}_AM2.5_CTL1860_florCTL1860noleapSST_minus2pctSolar_tigercpu_intelmpi_18_540PE_0101-0130_glbmean.nc'
da3 = xr.open_dataarray(ifile).pipe(func)
label3 = '-2pct solar'
 
 
if __name__ == '__main__':
    from wyconfig import * #my plot settings
    #ratio of 2xCO2 to 1pctSolar
    #years = slice(101,150)
    def print_ratio(years):
        print(years)
        tmean = lambda x: x.sel(year=years).mean()
        da0mean = da0.pipe(tmean)
        da1_ = da1.pipe(tmean) - da0mean
        print(f'2xCO2 forcing: {da1_.item():.2f}W/m^2')
        #da2_ = da2.pipe(tmean) - da0mean
        #print(f'1pctSolar forcing: {da2_.item():.2f}W/m^2')
        da3_ = da3.pipe(tmean) - da0mean
        print(f'-2pctSolar forcing: {da3_.item():.2f}W/m^2')
        #r = da2_/da1_
        #print(f'ratio of 1pctSolar/2xCO2 r = {r.item():.2f}')
        #print(f'ratio of 2xCO2/1pctSolar r = {1/r.item():.2f}')
        r2 = da3_/da1_
        print(f'ratio of m2pctSolar/2xCO2 r = {r2.item():.2f}')
        print(f'ratio of 2xCO2/m2pctSolar r = {1/r2.item():.2f}')
        print()
        return da1_.item(), da3_.item(), r2.item() #2xCO2, 1pctSolar, and ratio
    print(daname)
    years = slice(101,150)
    co2,solar2,r2 = print_ratio(years)
    years = slice(101,130)
    print_ratio(years)
    
    da0.plot(label=label0)
    da1.plot(label=label1 + f' {co2:.2f}W/m$^2$')
    #da2.assign_attrs(units='W/m$^2$').plot(label=label2 + f' {solar:.2f}W/m$^2$')
    da3.assign_attrs(units='W/m$^2$').plot(label=label3 + f' {solar2:.2f}W/m$^2$')

    plt.legend(ncol=2, loc='upper left')
    plt.title(f'AM2.5 ratio of m2pctSolar:2xCO2 = {r2:.2f}')
    plt.ylim(None, 6)
    

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