#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Wed Jun 21 15:36:03 EDT 2023
#WY: from /tigress/wenchang/analysis/AM2.5/fig_bars_dprecip_dF_florSST.py
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
from misc.comp2samp import comp2samp
#
if __name__ == '__main__':
    tt.check('end import')
#
model = 'AM2.5'
case = None# '1034'
#start from here
def func_anom(ifile, ifile_ctl): 
    #tspan = slice('0111', None) # omit the first 10 years 0101-0110
    tspan = slice(None, None)
    #da = xr.open_dataarray(ifile).sel(time=tspan).mean('time') 
    #da0 = xr.open_dataarray(ifile_ctl).sel(time=tspan).mean('time')
    #daa = (da/da0-1)*100
    #diff and err
    da = xr.open_dataarray(ifile).sel(time=tspan).groupby('time.year').mean('time') 
    da0 = xr.open_dataarray(ifile_ctl).sel(time=tspan).groupby('time.year').mean('time')
    ds = comp2samp(da, da0, dim='year')
    daa = ds.x1m/ds.x2m*100 - 100 #% diff
    daa.attrs['units'] = '%'
    err = ds.err/ds.x2m * 100
    err.attrs['unit'] = '%'
    #print(da, da0, ds, daa, err); sys.exit()
    """
    #dgmst
    da = xr.open_dataarray(ifile.replace('precip_', 't_surf_')).mean('time') 
    da0 = xr.open_dataarray(ifile_ctl.replace('precip_', 't_surf_')).mean('time')
    dgmst = da - da0
    #normalized by dgmst
    daa = daa/dgmst
    daa.attrs['units'] = '%/K'
    """
    #return daa
    return xr.Dataset(dict(dprecip=daa, err=err))

ofile = __file__.replace('.py', '.nc')
if not os.path.exists(ofile) or 'od' in sys.argv: #ofile does not exist or overwrite data (od) is specified
    dss = []
    labels  = []
    
    #plus2K
    label  = '+2K'
    ifile_ctl = '/tigress/wenchang/analysis/AM2.5/precip_AM2.5_CTL1860_florSST_tigercpu_intelmpi_18_540PE_0101-0200_glbmean.nc'
    ifile = '/tigress/wenchang/analysis/AM2.5/precip_AM2.5_CTL1860_florSST_plus2K_tigercpu_intelmpi_18_540PE_0101-0150_glbmean.nc'
    ds = func_anom(ifile, ifile_ctl)
    dss.append(ds)
    labels.append(label)
     
    #plus2Ktrop
    label  = '+2Ktrop'
    ifile_ctl = '/tigress/wenchang/analysis/AM2.5/precip_AM2.5_CTL1860_florSST_tigercpu_intelmpi_18_540PE_0101-0200_glbmean.nc'
    ifile = '/tigress/wenchang/analysis/AM2.5/precip_AM2.5_CTL1860_florSST_plus2Ktrop_tigercpu_intelmpi_18_540PE_0101-0150_glbmean.nc'
    ds = func_anom(ifile, ifile_ctl)
    dss.append(ds)
    labels.append(label)
     
    #plus2Kextr
    label = '+2Kextr'
    ifile_ctl = '/tigress/wenchang/analysis/AM2.5/precip_AM2.5_CTL1860_florSST_tigercpu_intelmpi_18_540PE_0101-0200_glbmean.nc'
    ifile = '/tigress/wenchang/analysis/AM2.5/precip_AM2.5_CTL1860_florSST_plus2Kextr_tigercpu_intelmpi_18_540PE_0101-0150_glbmean.nc'
    ds = func_anom(ifile, ifile_ctl)
    dss.append(ds)
    labels.append(label)

    """
    #2xCO2
    label  = '2xCO2'
    ifile_ctl = '/tigress/wenchang/analysis/AM2.5/precip_AM2.5_CTL1860_florSST_tigercpu_intelmpi_18_540PE_0101-0200_glbmean.nc'
    ifile = '/tigress/wenchang/analysis/AM2.5/precip_AM2.5_CTL1860_florSST_2xCO2_tigercpu_intelmpi_18_540PE_0101-0150_glbmean.nc'
    ds = func_anom(ifile, ifile_ctl)
    dss.append(ds)
    labels.append(label)
     
    #+2%solar
    label  = '+2%solar'
    ifile_ctl = '/tigress/wenchang/analysis/AM2.5/precip_AM2.5_CTL1860_florSST_tigercpu_intelmpi_18_540PE_0101-0200_glbmean.nc'
    ifile = '/tigress/wenchang/analysis/AM2.5/precip_AM2.5_CTL1860_florSST_solar2pctplus_tigercpu_intelmpi_18_540PE_0101-0150_glbmean.nc'
    ds = func_anom(ifile, ifile_ctl)
    dss.append(ds)
    labels.append(label)
     
    """
     
    ds = xr.concat(dss, dim=pd.Index(labels, name='label'))
    #save
    ds.to_netcdf(ofile)
    print('[saved]:', ofile)
else:
    ds = xr.open_dataset(ofile)
    print('[loaded]:', ofile)

da = ds['dprecip']
#cal the residual
residual = da.sel(label='+2K') - da.sel(label='+2Ktrop') - da.sel(label='+2Kextr')
residual = xr.DataArray(residual.item(), dims='label', coords=[['residual']]) 
da = xr.concat([da, residual], dim='label')
df = da.to_pandas()

def wyplot(ax=None, **kws):
    if ax is None: fig,ax = plt.subplots()
    n_labels = da.label.size
    #df.plot.bar(rot=0, color=[f'C{ii}' for ii in range(n_labels)], yerr=ds['err'], **kws)
    df.plot.bar(rot=0, color=[f'C{ii}' for ii in range(n_labels)], ax=ax,  **kws) #yerr too small

    for ii in range(n_labels):
        if ii == 0:
            y = da.isel(label=ii).item()
            ax.text(ii, y, f'{y:.2f}', va='bottom', ha='center')
        else:
            y = da.isel(label=ii).item()
            y_total = da.isel(label=0).item()
            frac = y/y_total
            ax.text(ii, max(y,0), f'{y:.2f}({frac:.2f})', va='bottom', ha='center')

    ax.set_ylabel(f'global mean precip change [%]')
    ax.set_xlabel(f'{model} experiment')
    ax.axhline(0, color='gray', ls='--')
    
    
if __name__ == '__main__':
    from wyconfig import * #my plot settings
    wyplot()

    #savefig
    if 'savefig' in sys.argv or 's' in sys.argv:
        figname = __file__.replace('.py', f'.png')
        if case is not None: figname = figname.replace('.png', f'__{case}.png')
        if 'pdf' in sys.argv:
            figname = figname.replace('.png', '.pdf')
        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()