#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Mon Jul  3 13:22:33 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
from fig_lines_precip import wyplot as plot_lines_precip
#
if __name__ == '__main__':
    tt.check('end import')
#
#start from here
 
 
if __name__ == '__main__':
    from wyconfig import * #my plot settings
    exclude_2xCO2 = True
    figsize = (10,8) if exclude_2xCO2 else (10,11)
    dpi = 80 
    nrows = 3 if exclude_2xCO2 else 4
    fig,axes = plt.subplots(nrows, 2, figsize=figsize, dpi=dpi, sharey='row')
    cases = ['co2x2', 'co2x4', 'co2xp5', 'co2xp25']
    labels = ['2xCO2', '4xCO2', '0.5xCO2', '0.25xCO2']
    if exclude_2xCO2:
        cases = cases[1:]
        labels = labels[1:]
    TsOptions = ['glbmean', 'tropoceanmean']

    for ii,case in enumerate(cases):
        for jj,TsOption in enumerate(TsOptions):
            ax = axes[ii, jj]
            plot_lines_precip(case=case, TsOption=TsOption, ax=ax, lowpass_on=True, residual_on=False)
            
            #ylabel
            if jj==0:
                ax.set_ylabel(f'{labels[ii]} global precip change [%]')
            else:
                ax.set_ylabel('')
            #xlabel
            if ii == 3:
                ax.set_xlabel('year')
            else:
                ax.set_xlabel('')
    #col title
    ax = axes[0,0]
    ax.set_title('$T_s$: global mean surface temperature')
    ax = axes[0,1]
    ax.set_title('$T_s$: tropical mean sea surface temperature')
            
    #savefig
    if 'savefig' in sys.argv or 's' in sys.argv:
        figname = __file__.replace('.py', f'.png')
        if exclude_2xCO2: figname = figname.replace('.png', '__no2xCO2.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()