#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Thu Apr 20 11:45:38 EDT 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
wython = '/tigress/wenchang/wython'
if wython not in sys.path: sys.path.append(wython); print('added to python path:', wython)
#from misc import get_kws_from_argv
import xlinregress
#
if __name__ == '__main__':
    tt.check('end import')
#
#start from here
def cal_ratio(ifile):
    """calculate the ratio of Gregory plot slope over years 21-150 to years 1-21""" 
    #dT
    #ifile = 't_surf_FLOR_co2x2_CTL1860_tigercpu_intelmpi_18_576PE_0101-0600_glbmean.nc'
    print(f'{ifile =}')
    ifile_ctl = glob.glob(ifile.replace('_co2x2_CTL1860', '_CTL1860_newdiag').replace('0101-0600', '????-????'))[0]
    print(f'{ifile_ctl =}')
    dT = xr.open_dataarray(ifile).groupby('time.year').mean('time') \
        - xr.open_dataarray(ifile_ctl).groupby('time.year').mean('time').sel(year=slice(101,200)).mean('year')
    
    #dN
    ifile_N = 'netrad_toa_FLOR_co2x2_CTL1860_tigercpu_intelmpi_18_576PE_0101-0600_glbmean.nc'
    print(f'{ifile_N =}')
    ifile_N_ctl = glob.glob(ifile_N.replace('_co2x2_CTL1860', '_CTL1860_newdiag').replace('0101-0600', '????-????'))[0]
    print(f'{ifile_N_ctl =}')
    dN = xr.open_dataarray(ifile_N).groupby('time.year').mean('time') \
        - xr.open_dataarray(ifile_N_ctl).groupby('time.year').mean('time').sel(year=slice(101,200)).mean('year')
    
    #years 1-20
    yearspan = slice(0,20)
    reg = dN.isel(year=yearspan).linregress.on(dT.isel(year=yearspan), dim='year')
    b = reg.slope.item()
    b1 = b
    print(f'{b1 =}')
     
    #years 21-150
    yearspan = slice(20,150)
    reg = dN.isel(year=yearspan).linregress.on(dT.isel(year=yearspan), dim='year')
    b = reg.slope.item()
    b2 = b
    print(f'{b2 =}')
    
    print(f'{b2/b1 =}')
    print()
    return b1,b2,b2/b1

ofile = __file__.replace('.py', '.nc')
if os.path.exists(ofile):
    ds = xr.open_dataset(ofile)
    print('[loaded]:', ofile)
else:
    lats = range(20,90+1,10)
    b1s = []
    b2s = []
    ratios = []
    
    for lat in lats:
        if lat == 30:
            ifile = f'tropmean/t_surf_FLOR_co2x2_CTL1860_tigercpu_intelmpi_18_576PE_0101-0600_tropmean.nc'
        elif lat == 90:
            ifile = f't_surf_FLOR_co2x2_CTL1860_tigercpu_intelmpi_18_576PE_0101-0600_glbmean.nc'
        else:
            ifile = f'trop{lat}mean/t_surf_FLOR_co2x2_CTL1860_tigercpu_intelmpi_18_576PE_0101-0600_trop{lat}mean.nc'
        b1,b2,r = cal_ratio(ifile)
        #print(f'{r =}')
        b1s.append(b1)
        b2s.append(b2)
        ratios.append(r)
    
    b1 = xr.DataArray(b1s, dims='lat_bound', coords=[lats,]) 
    b2 = xr.DataArray(b2s, dims='lat_bound', coords=[lats,]) 
    ratio = xr.DataArray(ratios, dims='lat_bound', coords=[lats,]) 
    ds = xr.Dataset(dict(b1=b1, b2=b2, ratio=ratio))
    ds.to_netcdf(ofile)
    #da.to_dataset(name='ratio').to_netcdf(ofile)
    print('[saved]:', ofile)

if __name__ == '__main__':
    from wyconfig import * #my plot settings
    fig,ax = plt.subplots()
    func = lambda x: -x
    ds['b1'].pipe(func).plot(marker='o', fillstyle='none', label='slope over years 1-20 (b1, W/m^2/K)')
    ds['b2'].pipe(func).plot(marker='o', fillstyle='none', label='slope over years 21-150 (b2, W/m^2/K)')
    ds['ratio'].plot(marker='o', fillstyle='none', label='b2/b1')

    ax.legend()
    ax.axhline(1, color='gray', ls='--')
    ax.set_title('ratio of slope over years 21-150 to 1-20 in Gregory plot')
    ax.set_xlabel('upper latitude used in tropical mean surface temperature')
    ax.set_ylabel('')
    
    #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)

    #reciprocal of slope
    fig,ax = plt.subplots()
    func = lambda x: -1/x
    ds['b1'].pipe(func).plot(marker='o', fillstyle='none', label='slope over years 1-20 (1/b1, K/W/m^2)')
    ds['b2'].pipe(func).plot(marker='o', fillstyle='none', label='slope over years 21-150 (1/b2, K/W/m^2)')
    ds['ratio'].pipe(lambda x: 1/x).plot(marker='o', fillstyle='none', label='(1/b2)/(1/b1)')

    ax.legend()
    ax.axhline(1, color='gray', ls='--')
    ax.set_title('ratio of 1/slope over years 21-150 to 1-20 in Gregory plot')
    ax.set_xlabel('upper latitude used in tropical mean surface temperature')
    ax.set_ylabel('')
    
    #savefig
    if 'savefig' in sys.argv or 's' in sys.argv:
        figname = __file__.replace('.py', f'_reciprocal.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()