#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Thu Jan 13 18:14:26 EST 2022
if __name__ == '__main__':
    import sys
    from misc.timer import Timer
    tt = Timer('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
#
if __name__ == '__main__':
    tt.check('end import')
#
#start from here
ifile_e5 = '/tigress/wenchang/data/era5/analysis_wy/TCI_v202201/wyanalysis/era5.yearly.PI.vmax.global.FMA-ASO.1979-2020.nc'
ifile_m2 = '/tigress/wenchang/data/merra2/analysis/TCI/wyanalysis/merra2.yearly.PI.vmax.global.FMA-ASO.1980-2020.nc'
season = ifile_e5.split('.')[-3]
ds_e5 = xr.open_dataset(ifile_e5)
ds_m2 = xr.open_dataset(ifile_m2)

if __name__ == '__main__':
    from wyconfig import * #my plot settings
    fig, ax = plt.subplots()
    da = ds_e5['vmax']
    ds_e5.vmax.plot(color='C0', marker='o', fillstyle='none', label='ERA5 global', ax=ax)
    ds_e5.vmaxNH.plot(color='C0', ls='--', label='ERA5 NH', ax=ax)
    ds_e5.vmaxSH.plot(color='C0', ls=':', label='ERA5 SH', ax=ax)

    ds_m2.vmax.plot(color='C1', marker='o', fillstyle='none', label='MERRA2 global', ax=ax)
    ds_m2.vmaxNH.plot(color='C1', ls='--', label='MERRA2 NH', ax=ax)
    ds_m2.vmaxSH.plot(color='C1', ls=':', label='MERRA2 SH', ax=ax)

    ax.legend()
    ax.set_title(f'{season} potential intensity from ERA5 and MERRA2')
    ax.set_ylabel('PI vmax [m/s]')
    
    #savefig
    if len(sys.argv)>1 and 'savefig' in sys.argv[1:]:
        figname = __file__.replace('.py', f'.png')
        if 'overwritefig' in sys.argv[1:]:
            wysavefig(figname, overwritefig=True)
        else:
            wysavefig(figname)
    tt.check(f'**Done**')
    plt.show()