#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Tue Jun 30 22:10:12 EDT 2020
if __name__ == '__main__':
    from misc.timer import Timer
    tt = Timer(f'start {__file__}')
import sys, os.path, os, glob
import xarray as xr, numpy as np, pandas as pd
#import matplotlib.pyplot as plt
#more imports
import geoxarray
#
if __name__ == '__main__':
    tt.check('end import')
#
#start from here
idata_name = 'ts'
#ifile = 'ts.historical_ssp585.1850-2100.ens26.ocean.nc'
ifile = 'ts.past1000.0850-1849.ens02.ocean.nc'
ofile = ifile.replace('.nc', '.tc.nc')
if os.path.exists(ofile):
    print('**exists**:', ofile)
    sys.exit()
lons = dict(
    mdr=slice(280, 340)
    )
lats = dict(
    mdr=slice(10, 25),
    trop=slice(-30,30)
    )
#years_base = range(1981, 2011)
years_base = range(1800, 1850)
 
if __name__ == '__main__':
    #from wyconfig import * #my plot settings
    with xr.open_dataset(ifile) as ds:
        sst = ds[idata_name]

    #MDR
    mdr = sst.sel(lon=lons['mdr'], lat=lats['mdr']).geo.fldmean()
    mdr = mdr.groupby('time.year').mean('time')
    mdr = mdr.assign_attrs(long_name='MDR SST', units='K')
    mdr_a = mdr - mdr.sel(year=years_base).mean('year')

    #TROP
    trop = sst.sel(lat=lats['trop']).geo.fldmean()
    trop = trop.groupby('time.year').mean('time')
    trop = trop.assign_attrs(long_name='Tropical mean SST', units='K')
    trop_a = trop - trop.sel(year=years_base).mean('year')

    #rename
    MDRa = mdr_a
    TROPa = trop_a

    #HU
    HU = np.exp(1.707 + 1.388*MDRa - 1.521*TROPa)
    HU = HU.assign_attrs(long_name='HU#: EXP(1.707 + 1.388*MDRa - 1.521*TROPa)')

    #TS
    TS = np.exp(2.10356 + 0.97612*MDRa - 0.97102*TROPa)
    TS = TS.assign_attrs(long_name='TS#: EXP(2.10356 + 0.97612*MDRa - 0.97102*TROPa)')

    #PDI
    PDI = np.exp(0.76 + 1.87*MDRa - 1.58*TROPa)
    PDI = PDI.assign_attrs(long_name='PDI: EXP(0.76 + 1.87*MDRa - 1.58*TROPa)')

    ds = xr.Dataset(dict(
        HU=HU, TS=TS, PDI=PDI, MDR=mdr, TROP=trop))
    ds = ds.assign_attrs(years_base=f'{years_base[0]}-{years_base[-1]}')

    ds.to_netcdf(ofile)
    print('**saved**:', ofile)

    tt.check(f'**Done**')