#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Mon Aug 12 11:28:48 EDT 2019

import xarray as xr, numpy as np, pandas as pd

# tc lme reconstruction
ens = range(1, 13+1)
HU = []
TS = []
PDI = []
MDR = []
TROP = []
for en in ens:
    print(en, end='; ')
    ifile = f'/tigress/gvecchi/DATA/CESM/LME/SST/indices.b.e11.BLMTRC5CN.f19_g16.{en:03d}.pop.h.SST.nc'
    ds = xr.open_dataset(ifile).rename({'TIME': 'time'})
    ds = ds.mean('Z_T').assign_coords(time=ds.indexes['time'].shift(-1, 'MS')) \
        .groupby('time.year').mean('time')
    ssta_mdr = ds['MDRSST'].pipe(lambda x: x-x.sel(year=slice(1901, 2001)).mean())
    ssta_trop = ds['TROPSST'].pipe(lambda x: x-x.sel(year=slice(1901, 2001)).mean())
    HU.append( np.exp(1.707+1.388*ssta_mdr-1.521*ssta_trop) )
    TS.append( np.exp(2.10356+0.97612*ssta_mdr-0.97102*ssta_trop) )
    PDI.append( np.exp(0.76+1.87*ssta_mdr-1.58*ssta_trop) )
    MDR.append( ssta_mdr )
    TROP.append( ssta_trop )

HU = xr.concat(HU, dim=pd.Index(ens, name='en')) \
    .transpose('year', 'en')
HU.attrs['long_name'] = 'EXP(1.707+1.388*MDR-1.521*TROP)'

TS = xr.concat(TS, dim=pd.Index(ens, name='en')) \
    .transpose('year', 'en')
TS.attrs['long_name'] = 'EXP(2.10356+0.97612*MDR-0.97102*TROP)'

PDI = xr.concat(PDI, dim=pd.Index(ens, name='en')) \
    .transpose('year', 'en')
PDI.attrs['long_name'] = 'EXP(0.76+1.87*MDR-1.58*TROP)'

MDR = xr.concat(MDR, dim=pd.Index(ens, name='en')) \
    .transpose('year', 'en')
MDR.attrs['long_name'] = 'x-x.sel(year=slice(1901, 2001)).mean()'
MDR.attrs['history'] = '/tigress/gvecchi/DATA/CESM/LME/SST/indices.b.e11.BLMTRC5CN.f19_g16.[1-13].pop.h.SST.nc'

TROP = xr.concat(TROP, dim=pd.Index(ens, name='en')) \
    .transpose('year', 'en')
TROP.attrs['long_name'] = 'x-x.sel(year=slice(1901, 2001)).mean()'
TROP.attrs['history'] = '/tigress/gvecchi/DATA/CESM/LME/SST/indices.b.e11.BLMTRC5CN.f19_g16.[1-13].pop.h.SST.nc'

ds = xr.Dataset(dict(HU=HU, TS=TS, PDI=PDI, MDR=MDR, TROP=TROP))
ds.to_netcdf('LME_TC.nc')