#!/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_HIRAM_CTL1990s_v201910_tigercpu_intelmpi_18_540PE_0101-0300_glbmean.nc'
da = xr.open_dataarray(ifile)
da_ctl = da
 
ifile = 'precip_HIRAM_CTL1990s_v201910_plus2K_tigercpu_intelmpi_18_540PE_0101-0150_glbmean.nc'
da = xr.open_dataarray(ifile)
da_2k = da
 
ifile = 'precip_HIRAM_CTL1990s_v201910_2xCO2_tigercpu_intelmpi_18_540PE_0101-0150_glbmean.nc'
da = xr.open_dataarray(ifile)
da_2xco2 = da

#ts
ifile = 't_surf_HIRAM_CTL1990s_v201910_tigercpu_intelmpi_18_540PE_0101-0300_glbmean.nc'
da = xr.open_dataarray(ifile)
da_ctl_ts = da
ifile = 't_surf_HIRAM_CTL1990s_v201910_plus2K_tigercpu_intelmpi_18_540PE_0101-0150_glbmean.nc'
da = xr.open_dataarray(ifile)
da_ts = da
dgmst = ( da_ts.mean('time') - da_ctl_ts.mean('time') ).item()
print('plus2K dgmst:', dgmst, 'K')


func = lambda x: (x.mean('time')/da_ctl.mean('time')*100-100).assign_attrs(units='%').item()
print('plus2K precip sensitivity:', func(da_2k)/dgmst, '%/K gmst')
print('2xCO2 precip sensitivity:', func(da_2xco2), '%/2xCO2')
 
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()