#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Tue Jun 29 16:14:32 EDT 2021
if __name__ == '__main__':
    from misc.timer import Timer
    tt = Timer(f'start {__file__}')
import sys, os.path, os, glob, datetime
import xarray as xr, numpy as np, pandas as pd, matplotlib.pyplot as plt
#more imports
import geoxarray
from misc.landmask import flagland
import salem
#
if __name__ == '__main__':
    tt.check('end import')
#
#start from here
shapefile = '../shapefiles/sf_hawkes-gisborne/sf_hawkes-gisborne.shp'
shdf = salem.read_shapefile(shapefile)
ifile = 'precip_AM2.5C360_amipHadISSTrcp45_tigercpu_intelmpi_18_1080PE_3ens_1871-2050_NewZealand.nc'
dsname =  'precip'
daname = dsname
units = 'mm' # 'mm/day'
years = slice('1991', '2020')
title = f'AM2.5C360 AMIP 3ens New Zealand {years.start}-{years.stop} Rx2day {daname}'
tag = '_' + title.replace(' ', '_')
ofile_nc = __file__.replace('.py', f'{tag}.nc')

if os.path.exists(ofile_nc):
    da = xr.open_dataarray(ofile_nc)
    print('[loaded]:', ofile_nc)
else:
    print('calculating...')
    da = xr.open_dataarray(ifile).sel(time=years).load() \
        .rolling(time=2).sum().groupby('time.year').max('time') \
        .mean(['year', 'ens']) \
        .assign_attrs(units=units, long_name=title)
    print('saving...')
    ds = da.to_dataset(name=daname)
    ds.to_netcdf(ofile_nc)
    print('[saved]:', ofile_nc)
 
if __name__ == '__main__':
    from wyconfig import * #my plot settings
    from geoplots import mapplot
    figsize = None#(6,4)# None
    fig, ax = plt.subplots(figsize=figsize)
    vmax = 200
    vmin = 50
    levels = 31
    xticks = None# range(15,50+1,5)
    shdf.boundary.plot(ax=ax, color='C3', lw=1)
    da.assign_attrs(units=units).plot(cmap='YlGnBu', vmin=vmin, vmax=vmax, levels=levels)
    mapplot(xticks=xticks)
    #mapplot()
    #plt.title(title, loc='left')
    plt.xlabel('')
    plt.ylabel('')

    
    #savefig
    if 'savefig' in sys.argv or 's' in sys.argv:
        figname = __file__.replace('.py', f'{tag}.png')
        wysavefig(figname)
    tt.check(f'**Done**')
    plt.show()