#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Tue Jun 29 16:14:32 EDT 2021
#SBATCH -N 1  # 1 node
#SBATCH -n 40 # 40 cores
#SBATCH -t 5:00:00
#SBATCH --mail-type=begin
#SBATCH --mail-type=end
#SBATCH --mail-user=wenchang@princeton.edu
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
#import salem
from misc.modelout import get_modelout_data_ens, get_modelout_data
#
if __name__ == '__main__':
    tt.check('end import')
#
#start from here
#shapefile for the region
#cntry_name = 'Madagascar'
#shdf = salem.read_shapefile(salem.get_demo_file('world_borders.shp'))
#shdf = shdf[ shdf['CNTRY_NAME']==cntry_name ]
#shapefile = '../shapefile/eastcoast.shp'
#shdf = salem.read_shapefile(shapefile)

daname = 'precip'
model = 'AM2.5C360'
expname = 'amipHadISSTlong_chancorr_tigercpu_intelmpi_18_1080PE'
dsname = 'atmos_month'
years = range(1981,2020)
expname_extend2020 = 'amipHadISSTlong_chancorr_tigercpu_intelmpi_18_1080PE_extend2020'
years_extend2020 = range(2020,2021)
years_ext = years_extend2020
func = lambda x: x.load().rename(grid_xt='lon', grid_yt='lat') \
    .sel(lon=slice(15, 50), lat=slice(-35, -10))
funcname = 'SouthAfricaLarge'
ens = range(1,11)
ofile = f'data_{daname}_{model}_{expname}_{len(ens)}ens_{years[0]:04d}-{years_ext[-1]:04d}_{funcname}.nc'
if os.path.exists(ofile):
    da = xr.open_dataarray(ofile)
    print('[loaded]:', ofile)
else:
    ofile_ = ofile.replace('data_', 'tmp_').replace(f'{years[0]:04d}-{years_ext[-1]:04d}', f'{years[0]:04d}-{years[-1]:04d}')
    da = get_modelout_data_ens(daname, model, expname, func=func, funcname=funcname, ens=ens, years=years, ofile=ofile_)
    ofile_ = ofile.replace('data_', 'tmp_').replace(f'{years[0]:04d}-{years_ext[-1]:04d}', f'{years_ext[0]:04d}-{years_ext[-1]:04d}')
    da_ext = get_modelout_data_ens(daname, model, expname_extend2020, func=func, funcname=funcname, ens=ens, years=years_extend2020, ofile=ofile_)
    da = xr.concat([da, da_ext], dim='time')
    da.to_dataset(name=daname).to_netcdf(ofile)
    print('[saved]:', ofile)
    """
    das = []
    for en in ens:
        ofile_en = ofile.replace('data_', 'tmp_').replace(f'{len(ens)}ens', f'ens{en:02d}') \
            .replace(f'{years[0]:04d}-{years_ext[-1]:04d}', f'{years[0]:04d}-{years[-1]:04d}')
        da = get_modelout_data(daname, model, expname, ens=en, years=years, func=func, funcname=funcname, ofile=ofile_en, dsname=dsname)
        ofile_en_ext = ofile.replace('data_', 'tmp_').replace(f'{len(ens)}ens', f'ens{en:02d}') \
            .replace(f'{years[0]:04d}-{years_ext[-1]:04d}', f'{years_ext[0]:04d}-{years_ext[-1]:04d}')
        da_ext = get_modelout_data(daname, model, expname_extend2020, ens=en, years=years_extend2020, func=func, funcname=funcname, ofile=ofile_en_ext, dsname=dsname)
        da = xr.concat([da, da_ext], dim='time')
        das.append(da)
    da = xr.concat(das, dim=pd.Index(ens, name='ens'))
    ds = da.to_dataset(name=daname)
    ds.to_netcdf(ofile)
    print('[saved]:', ofile)
    """
 
if __name__ == '__main__':
    """
    from wyconfig import * #my plot settings
    from geoplots import mapplot
    fig, ax = plt.subplots(figsize=(6,6))
    da.mean('year').pipe(lambda x: x-273.15).assign_attrs(units='$^\circ$C').plot.contourf(ax=ax, levels=np.arange(20,41), extend='both')
    mapplot(ax=ax, xticks=np.arange(360-125, 360-117+1), yticks=np.arange(43, 54+1))
    """
    
    #savefig
    if 'savefig' in sys.argv:
        figname = __file__.replace('.py', f'.png')
        wysavefig(figname)
    tt.check(f'**Done**')
    #plt.show()