#!/usr/bin/env python
import xarray as xr
from geoplots import mapplot
import datetime

# time info
t0 = datetime.datetime.now()
fmt = '%Y-%m-%dT%H:%M:%S'
fmt_short = '%Y-%m-%d'
print('[start]:', t0.strftime(fmt))

# params
data_name = 'precip'
units = 'mm/day'
long_name = data_name
scale = 24*3600 
years = range(1932, 1940)
ens = range(1, 6)
figname = f'{data_name}.diff.{years[0]}-{years[-1]}.{t0.strftime(fmt_short)}.png'
label = f'{data_name} diff: ChanCorr - CTL, {years[0]}-{years[-1]}'
cmap = 'BrBG'

# data read
idir = '/tigress/wenchang/MODEL_OUT/HIRAM'
ifiles = [[f'{idir}/amipHadISSTlong_tigercpu_intelmpi_18_540PE/en{en:02d}/POSTP/{year:04d}0101.atmos_month.nc' for year in years] for en in ens]
da = xr.open_mfdataset(ifiles, concat_dim=['en', 'time'], combine='nested')[data_name]

ifiles1 = [[f'{idir}/gv_amipHadISSTlongChancorr_tigercpu_intelmpi_18_540PE/en{en:02d}/POSTP/{year:04d}0101.atmos_month.nc' for year in years] for en in ens]
da1 = xr.open_mfdataset(ifiles1, concat_dim=['en', 'time'], combine='nested')[data_name]

# compute
daa = da1.mean(['en', 'time']) - da.mean(['en', 'time'])
daa.compute()

# fig
plt.close('all')
daa.pipe(lambda x: x*scale).assign_attrs(dict(long_name=long_name, units=units)).plot(robust=True, levels=20, cmap=cmap)
mapplot()
plt.title(label)

plt.tight_layout()
plt.savefig(figname, dpi=150)

t1 = datetime.datetime.now()
print('[end]:', t1.strftime(fmt))
print('[time used]:', f'{(t1-t0):,} seconds')