#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Wed Dec 18 23:32:43 EST 2019
import sys, os.path, os, datetime
import xarray as xr, numpy as np, pandas as pd
#import matplotlib.pyplot as plt
#print()

#ifiles = 'LME_SST/b.e11.B1850C5CN.f19_g16.0850cntl.001.*.nc'
#ofile = 'LME_TC.0850cntl.nc'
#sstName = 'SST'
#xName = 'nlon'
#yName = 'nlat'
#latName = 'TLAT'
#lonName = 'TLONG'
#areaName = 'TAREA'
#shiftTimeAxis = True #CESM data time axis is one month late and needs a backward shift
#yearsRef = slice(1901, 2000) # reference years in calculation of sst anomaly for TC indices
def sst_to_vecchi_indices(ifiles, ofile=None, sstName='SST', xName='nlon', yName='nlat', lonName='TLONG', latName='TLAT', areaName='TAREA', shiftTimeAxis=True, yearsRef=None):
    print('[ifiles]:', ifiles)
    ds = xr.open_mfdataset(ifiles)
    if shiftTimeAxis:
        ds['time'] = ds.indexes['time'].shift(-1, 'MS')
        print('[time shifted]:', 'one month back')
    sst = ds[sstName].squeeze()
    if 'z_t' in list(sst.coords):
        sst = sst.drop('z_t')
        print('[coords dropped]:', 'z_t')
    area = ds[areaName] #+ sst*0 # apply the sst mask to area
    lon = ds[lonName]
    lat = ds[latName]

    #MDR(main development region)
    inMDR = (lat>=10)&(lat<=25)&(lon>=280)&(lon<=340)
    sst_ = sst.where(inMDR)
    area_ = area.where(inMDR)
    MDR = (sst_*area_).sum([xName, yName])/area_.sum([xName, yName]) # area mean
    MDR = MDR.groupby('time.year').mean('time') # yearly mean
    MDR.attrs = sst.attrs
    MDR.attrs['long_name'] = 'MDR ' + MDR.attrs['long_name']
    if yearsRef is not None:
        MDRa = MDR - MDR.sel(year=yearsRef).mean('year')
    else:
        MDRa = MDR - MDR.mean('year')

    #TROP (tropical reagion)
    inTROP = (lat>=-30)&(lat<=30)
    sst_ = sst.where(inTROP)
    area_ = area.where(inTROP)
    TROP = (sst_*area_).sum([xName, yName])/area_.sum([xName, yName]) # area mean
    TROP = TROP.groupby('time.year').mean('time') # yearly mean
    TROP.attrs = sst.attrs
    TROP.attrs['long_name'] = 'TROP ' + TROP.attrs['long_name']
    if yearsRef is not None:
        TROPa = TROP - TROP.sel(year=yearsRef).mean('year')
    else:
        TROPa = TROP - TROP.mean('year')

    #HU (hurricane #)
    HU = np.exp(1.707 + 1.388*MDRa - 1.521*TROPa)
    HU.attrs['long_name'] = 'HU#: EXP(1.707 + 1.388*MDRa - 1.521*TROPa)'
    if yearsRef:
        HU.attrs['yearsRef'] = f'{yearsRef.start}-{yearsRef.stop}'

    #TS (tropical storm #)
    TS = np.exp(2.10356 + 0.97612*MDRa - 0.97102*TROPa)
    TS.attrs['long_name'] = 'TS#: EXP(2.10356 + 0.97612*MDRa - 0.97102*TROPa)'
    if yearsRef:
        TS.attrs['yearsRef'] = f'{yearsRef.start}-{yearsRef.stop}'

    #PDI
    PDI = np.exp(0.76 + 1.87*MDRa - 1.58*TROPa)
    PDI.attrs['long_name'] = 'PDI: EXP(0.76 + 1.87*MDRa - 1.58*TROPa)'
    if yearsRef:
        PDI.attrs['yearsRef'] = f'{yearsRef.start}-{yearsRef.stop}'

    # output dataset
    ds_o = xr.Dataset(dict(
        HU=HU, TS=TS, PDI=PDI, MDR=MDR, TROP=TROP))

    # save to nc file
    if ofile and not os.path.exists(ofile):
        encoding = {v:{'_FillValue': None, 'dtype': 'float32'} 
            for v in ['HU', 'TS', 'PDI', 'MDR', 'TROP']}
        encoding['year'] = {'_FillValue': None, 'dtype': 'int32'}
        ds_o.to_netcdf(ofile, encoding=encoding)
        print('[saved]:', ofile)
    if ofile and os.path.exists(ofile):
        print('[exists]:', ofile)

    return ds_o
     
if __name__ == '__main__':
    tformat = '%Y-%m-%d %H:%M:%S'
    t0 = datetime.datetime.now()
    print('[start]:', t0.strftime(tformat))


    ifiles = 'LME_SST/b.e11.B1850C5CN.f19_g16.0850cntl.001.*.nc'
    ofile = 'LME_TC.0850cntl.nc'
    yearsRef = slice(1901, 2000) # reference years in calculation of sst anomaly for TC indices

    sst_to_vecchi_indices(ifiles, ofile, yearsRef=yearsRef)
    
    t1 = datetime.datetime.now()
    print('[total time used]:', f'{(t1-t0).seconds:,} seconds')
    print('[end]:', t1.strftime(tformat))
    print()