#!/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)
years1 = range(1920, 1930)
ens = range(1, 6)
figname = f'{data_name}.dust_bowl.{years[0]}:{years[-1]}-{years1[0]}:{years1[-1]}.{t0.strftime(fmt_short)}.png'
label = f'{data_name} diff: {years[0]}:{years[-1]} - {years1[0]}:{years1[-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}/amipHadISSTlong_tigercpu_intelmpi_18_540PE/en{en:02d}/POSTP/{year:04d}0101.atmos_month.nc' for year in years1] for en in ens]
da1 = xr.open_mfdataset(ifiles1, concat_dim=['en', 'time'], combine='nested')[data_name]

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

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

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

# US
plt.xlim(360-150, 360-40)
plt.ylim(10, 70)
plt.savefig(figname.replace('.png', '.US.png'), dpi=150)

ofile = 'precip_anom.dust_bowl.nc'
daa.to_dataset(name='precip_anom').to_netcdf(ofile)
print('[saved]:', ofile)

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