#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Mon Apr 25 17:02:51 EDT 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
ds = xr.open_dataset('../RH.global.tcseason3.hist-nat.97ens.nc')
da = ds.timeseries
models = [mm.split('_')[0] for mm in ds.model_member.values]
ds['model'] = xr.DataArray(models, dims='model_member')
n_models = ds.model.groupby(ds.model).count()
model_count = [f'{m}({n})' for m,n in zip(n_models.model.values, n_models.values)]

if __name__ == '__main__':
    from wyconfig import * #my plot settings
    fig, ax = plt.subplots(figsize=(8,6))
    da.plot(hue='model_member', color='gray', alpha=0.5, lw=1, ls='-')
    da_ = da.groupby(ds.model).mean('model_member').assign_coords(model=model_count)
    for model in da_.model.values:
        da_.sel(model=model).plot(label=model)
    da.groupby(ds.model).mean('model_member').mean('model').plot(color='k', ls='-', label='ens mean', lw=2)

    plt.legend(bbox_to_anchor=(1,1), loc='upper left')
    plt.ylabel('RH [%]')
    plt.title('hist-nat relative humidity', loc='left')
    
    #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**')
    print()
    plt.show()