#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Mon Aug  3 19:27:25 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
from misc.landmask import flagland
from xtc import tc_basins
import geoxarray
#
if __name__ == '__main__':
    tt.check('end import')
#
#start from here
ofile = __file__.replace('.py', '.nc')
if os.path.exists(ofile):
    da = xr.open_dataarray(ofile)
    print('[loaded]:', ofile)
else:
    lat = xr.DataArray(np.arange(-89.5,90), dims='lat')
    lon = xr.DataArray(np.arange(0.5,360), dims='lon')
    da = (lat + lon).assign_coords(lat=lat, lon=lon)
    land = flagland(da)
    bs = tc_basins()
    ocean = (1 - land).where(land<0.5)
    
    basins = ['NA', 'EP', 'WP', 'NI', 'SI', 'AU', 'SP']
    basin_areas = [ocean.where(bs.mask(ocean)==bs.map_keys(basin)).sel(lat=slice(10,30)).geo.fldint()*1e-12 if basin not in ('SI', 'AU', 'SP')
        else ocean.where(bs.mask(ocean)==bs.map_keys(basin)).sel(lat=slice(-30,-10)).geo.fldint()*1e-12
        for basin in basins]
    da = xr.DataArray(basin_areas, dims='basin', coords=[basins])
    da = da.assign_attrs(units='million km$^2$', long_name='TC basin area')
    ds = da.to_dataset(name='area')
    ds.to_netcdf(ofile)
    print('[newly saved]:', ofile)

if __name__ == '__main__':
    from wyconfig import * #my plot settings
    plt.figure()
    da.to_pandas().plot.bar()
    plt.ylabel(da.attrs['units'])
    
    tt.check(f'**Done**')
    plt.show()