#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Fri Mar  3 11:52:49 EST 2023
if __name__ == '__main__':
    import sys,os
    from misc.timer import Timer
    tt = Timer(f'[{os.getcwd()}] 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 import get_kws_from_argv
#
if __name__ == '__main__':
    tt.check('end import')
#
#start from here
#precip
ifile = 'precip_AM2.5_CTL1990s_tigercpu_intelmpi_18_540PE_0101-0200_glbmean.nc'
da = xr.open_dataarray(ifile)
da_ctl = da
 
ifile = 'precip_AM2.5_CTL1990s_plus2K_tigercpu_intelmpi_18_540PE_0101-0150_glbmean.nc'
da = xr.open_dataarray(ifile)
da_2k = da
 
ifile = 'precip_AM2.5_CTL1990s_2xCO2_tigercpu_intelmpi_18_540PE_0101-0150_glbmean.nc'
da = xr.open_dataarray(ifile)
da_2xco2 = da

ifile = 'precip_AM2.5_CTL1990s_plus4K_tigercpu_intelmpi_18_540PE_0101-0150_glbmean.nc'
da = xr.open_dataarray(ifile)
da_4k = da
 
#ts
ifile = 't_surf_AM2.5_CTL1990s_tigercpu_intelmpi_18_540PE_0101-0200_glbmean.nc'
da = xr.open_dataarray(ifile)
da_ctl_ts = da
#plus2K
ifile = 't_surf_AM2.5_CTL1990s_plus2K_tigercpu_intelmpi_18_540PE_0101-0150_glbmean.nc'
da = xr.open_dataarray(ifile)
da_ts = da
dgmst2k = ( da_ts.mean('time') - da_ctl_ts.mean('time') ).item()
print('plus2K dgmst:', dgmst2k, 'K')
#plus4K
ifile = 't_surf_AM2.5_CTL1990s_plus4K_tigercpu_intelmpi_18_540PE_0101-0150_glbmean.nc'
da = xr.open_dataarray(ifile)
da_ts = da
dgmst4k = ( da_ts.mean('time') - da_ctl_ts.mean('time') ).item()
print('plus4K dgmst:', dgmst4k, 'K')


func = lambda x: (x.mean('time')/da_ctl.mean('time')*100-100).assign_attrs(units='%').item()
print('plus2K precip change:', func(da_2k), '%')
print('plus2K gmst change:', dgmst2k, 'K')
print('plus2K precip sensitivity:', func(da_2k)/dgmst2k, '%/K gmst')
print()

print('2xCO2 precip sensitivity:', func(da_2xco2), '%/2xCO2')

print()
print('plus4K precip change:', func(da_4k), '%')
print('plus4K gmst change:', dgmst4k, 'K')
print('plus4K precip sensitivity:', func(da_4k)/dgmst4k, '%/K gmst')
 
if __name__ == '__main__':
    #from wyconfig import * #my plot settings
    
    #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 or 'n' in sys.argv:
        pass
    else:
        if 'plt' in globals(): plt.show()