#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Mon Aug 19 10:32:34 EDT 2019
import os, os.path, sys
import xarray as xr, numpy as np

from xtc.basins import tc_basins
from xtc.counts import tc_count
#from density_new import tc_density
from xtc.density import tc_density

years = range(1980, 2021+1)

# TC tracks
#ifile = f'ibtracs.tracks.{years[0]:04d}-{years[-1]:04d}.nc'
ifile = f'IBTrACS.ALL.v04r00.2022-01-25.tracksByYear.{years[0]}-{years[-1]}.nc'
ds = xr.open_dataset(ifile)
#only TS17: max of windmax at least 17m/s
minWindmax = 17
ds = ds.where(ds['windmax'].max('stage') > minWindmax)
#minLife = 2 days
minLife = 2 #days
#ds = ds.where((ds['windmax']>17).count('stage') > minLife*8) #3 hourly data, each day has 8 track points
ds = ds.where( (ds['windmax']>17).sum('stage') > minLife*8 )

# TC # counts in global domain and each basin
#ofile = f'ibtracs.counts.{years[0]:04d}-{years[-1]:04d}.nc'
ofile = ifile.replace('tracksByYear', f'counts.TS17.{minLife}days17')
if not os.path.exists(ofile):
    basins = ('global', 'NA', 'EP', 'WP', 'NI', 'SI', 'AU', 'SP', 'SA')
    n_tc = { basin: tc_count(ds, basin=basin)
        for basin in basins} # dict comprehension
    ds_counts = xr.Dataset(n_tc)

    # save to netcdf
    dnames = basins + ds_counts['global'].dims 
    encoding = {dname: {'dtype': 'int32'}
        for dname in dnames}
    ds_counts.to_netcdf(ofile, encoding=encoding)
    print('[Saved]:', ofile)
else:
    print('[Exists]:', ofile)
# monthly to yearly
ofile_yearly = ofile.replace('.nc', '.yearly.nc')
if not os.path.exists(ofile_yearly):
    if 'ds_counts' not in globals():
        ds_counts = xr.load_dataset(ofile)
    ds_counts_yearly = ds_counts.groupby('time.year').sum('time')
    dnames = ('global', 'NA', 'EP', 'WP', 'NI', 'SI', 'AU', 'SP', 'SA')
    encoding = {dname: {'dtype': 'int32'}
        for dname in dnames + ds_counts_yearly['global'].dims}
    for dname in dnames:
        ds_counts_yearly[dname].attrs['long_name'] = f'{dname} yearly TC #'
    ds_counts_yearly.to_netcdf(ofile_yearly, encoding=encoding)
    print('[Saved]:', ofile_yearly)
else:
    print('[Exists]:', ofile_yearly)
# monthly to monthly climatology
ofile_clim = ofile.replace('.nc', '.climatology.nc')
if not os.path.exists(ofile_clim):
    if 'ds_counts' not in globals():
        ds_counts = xr.load_dataset(ofile)
    ds_counts_clim = ds_counts.groupby('time.month').mean('time')
    dnames = ('global', 'NA', 'EP', 'WP', 'NI', 'SI', 'AU', 'SP', 'SA')
    encoding = {dname: {'dtype': 'float32'}
        for dname in dnames + ds_counts_clim['global'].dims}
    for dname in dnames:
        ds_counts_clim[dname].attrs['long_name'] = f'{dname} TC # monthly climatology'
    ds_counts_clim.to_netcdf(ofile_clim, encoding=encoding)
    print('[Saved]:', ofile_clim)
else:
    print('[Exists]:', ofile_clim)

# TC density
#ofile = f'ibtracs.density.{years[0]:04d}-{years[-1]:04d}.nc'
ofile = ifile.replace('tracksByYear', f'density.TS17.{minLife}days17')
if not os.path.exists(ofile):
    density = tc_density(ds, lowpass_on=True, genesis_on=False)
    genesis_density = tc_density(ds, lowpass_on=True, genesis_on=True)
    ds_density = xr.Dataset(dict(density=density,
                            density_g=genesis_density))

    # save to netcdf
    encoding = {dname: {'dtype': 'int32', 'zlib': True, 'complevel': 1}
        if dname in ['time',]
        else {'dtype': 'float32', 'zlib': True, 'complevel': 1}
        for dname in ['time', 'lat', 'lon', 'density', 'density_g']}
    ds_density.to_netcdf(ofile, encoding=encoding)
    print('[Saved]:', ofile)
else:
    print('[Exists]:', ofile)
# monthly to yearly
ofile_yearly = ofile.replace('.nc', '.yearly.nc')
if not os.path.exists(ofile_yearly):
    if 'ds_density' not in globals():
        ds_density = xr.load_dataset(ofile)
    ds_density_yearly = ds_density.groupby('time.year').sum('time')
    # save to netcdf
    ds_density_yearly['density'].attrs['long_name'] = 'TC density'
    ds_density_yearly['density'].attrs['units'] = 'TC days per year per 10x10deg box'
    ds_density_yearly['density_g'].attrs['long_name'] = 'TC genesis density'
    ds_density_yearly['density_g'].attrs['units'] = 'TC genesis # per year per 10x10deg box'
    dnames_float = ('density', 'density_g', 'lat', 'lon')
    dnames_int = ('year', 'en') if 'en' in ds_density_yearly.dims else ('year',)
    encoding = {dname: {'dtype': 'int32', 'zlib': True, 'complevel': 1}
        if dname in dnames_int
        else  {'dtype': 'float32', 'zlib': True, 'complevel': 1}
        for dname in dnames_float + dnames_int}
    ds_density_yearly.to_netcdf(ofile_yearly, encoding=encoding)
    print('[Saved]:', ofile_yearly)
else:
    print('[Exists]:', ofile_yearly)
# monthly to monthly climatology
ofile_clim = ofile.replace('.nc', '.climatology.nc')
if not os.path.exists(ofile_clim):
    if 'ds_density' not in globals():
        ds_density = xr.load_dataset(ofile)
    ds_density_clim = ds_density.groupby('time.month').mean('time')
    # save to netcdf
    ds_density_clim['density'].attrs['long_name'] = 'TC density'
    ds_density_clim['density'].attrs['units'] = 'TC days per month per 10x10deg box'
    ds_density_clim['density_g'].attrs['long_name'] = 'TC genesis density'
    ds_density_clim['density_g'].attrs['units'] = 'TC genesis # per month per 10x10deg box'
    dnames_float = ('density', 'density_g', 'lat', 'lon')
    dnames_int = ('month', 'en') if 'en' in ds_density_clim.dims else ('month',)
    encoding = {dname: {'dtype': 'int32', 'zlib': True, 'complevel': 1}
        if dname in dnames_int
        else  {'dtype': 'float32', 'zlib': True, 'complevel': 1}
        for dname in dnames_float + dnames_int}
    ds_density_clim.to_netcdf(ofile_clim, encoding=encoding)
    print('[Saved]:', ofile_clim)
else:
    print('[Exists]:', ofile_clim)