#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Thu Jun  4 09:44:59 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
source = 'HadISST'
sst_name = 'sst'
label = 'sstindex'
def get_sstindex(ifile=None):
    if ifile is None:
        ifile = '/tigress/wenchang/HIRAM/input_wy/HadISST/HadISST_sst.187001-201912.nc'
    with xr.open_dataset(ifile) as ds_in:
        sst = ds_in[sst_name]
    #sst = sst.groupby('time.year').mean('time')

    lat_trop = slice(30,-30)
    sst_trop = sst.sel(latitude=lat_trop).geo.fldmean()
    lat_mdr = slice(25,10)
    lon_mdr = slice(-80, -20)
    sst_mdr = sst.sel(latitude=lat_mdr, longitude=lon_mdr).geo.fldmean()

    ds = xr.Dataset(dict(
        sst_trop=sst_trop,
        sst_mdr=sst_mdr
        ))
    ds.sst_trop.attrs['long_name'] = f'tropical mean sst'
    ds.sst_trop.attrs['units'] = 'deg C'
    ds.sst_trop.attrs['note'] = f'lat: {lat_trop.start} to {lat_trop.stop}'
    ds.sst_mdr.attrs['long_name'] = 'main development region mean sst'
    ds.sst_mdr.attrs['units'] = 'deg C'
    ds.sst_mdr.attrs['note'] = f'lat: {lat_mdr.start} to {lat_mdr.stop}; lon: {lon_mdr.start} to {lon_mdr.stop}'

    return ds


 
if __name__ == '__main__':
    from wyconfig import * #my plot settings
    source = 'HadISST'
    sst_name = 'sst'
    label = 'sstindex'
    ifile = '/tigress/wenchang/HIRAM/input_wy/HadISST/HadISST_sst.187001-201912.nc'
    ofile = f'{source}.{label}.{ifile.split(".")[-2]}.nc'
    if os.path.exists(ofile):
        print('[exists]:', ofile)
        ds = xr.open_dataset(ofile) 
    else:
        ds = get_sstindex(ifile)
        ds.to_netcdf(ofile)
        print('[saved]:', ofile)

    figname = ofile.replace('.nc', f'.{tt.today()}.png')
    ds.pipe(lambda x: x.groupby('time.month') - x.groupby('time.month').mean('time') ).drop('month').to_dataframe().plot(alpha=0.5)
    plt.ylabel(f'{source}[deg C]')
    plt.savefig(figname)
    print('[saved]:', figname)

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