#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Fri Sep 27 12:01:27 EDT 2019
from datetime import datetime
import os.path, sys, os, glob
import xarray as xr, numpy as np, pandas as pd
#import matplotlib.pyplot as plt
import salem

cntry_name = 'Mozambique'
model = 'FLOR'
expname = 'nudgeLongAll_newdiag_tigercpu_intelmpi_18_576PE'
ens = range(1, 7)
data_name = 'precip'
 
def get_region(cntry_name, expname, data_name, model='FLOR', years=None, en=None):
    '''get model output over a specified country from FLOR-like model experiments'''
    #cntry_name = 'Mozambique'
    #expname = 'CTL1860_newdiag_tigercpu_intelmpi_18_576PE'
    #data_name = "precip"
    cntry_sname = cntry_name.replace(' ', '')
    ofile = f'{model}.{expname}.{data_name}.{cntry_sname}.nc'
    if en is not None:
        ofile = ofile.replace(expname, f'{expname}.en{en:02d}')
    if os.path.exists(ofile):
        print('[exists]:', ofile)
        return
    # country name and lon/lat ranges
    #cntry_name = 'Mozambique'
    #cntry_sname = cntry_name.replace(' ', '')
    shdf = salem.read_shapefile(salem.get_demo_file('world_borders.shp'))
    shdf = shdf[ shdf['CNTRY_NAME']==cntry_name ]
    xmin = shdf.min_x.min() - 1
    xmax = shdf.max_x.max() + 1
    ymin = shdf.min_y.min() - 1
    ymax = shdf.max_y.max() + 1
    
    # IO files
    #expname = 'CTL1860_newdiag_tigercpu_intelmpi_18_576PE'
    ifiles = f'/tigress/wenchang/MODEL_OUT/{expname}/POSTP/????0101.atmos_daily.nc'
    if en is not None:
        ifiles = ifiles.replace(expname, f'{expname}/en{en:02d}')
    if model != 'FLOR':
        ifiles = ifiles.replace('MODEL_OUT', 'MODEL_OUT/{model}')
    if years is None: # get all available years
        ifiles_list = sorted(glob.glob(ifiles))
    else: # only get the specified years
        ifiles_list = [ifiles.replace('????', f'{year:04d}') for year in years]
        
    xname = 'grid_xt'
    xlim = (xmin, xmax)
    yname = 'grid_yt'
    ylim = (ymin, ymax)
    yincrease = True
    
    # lon range
    x = xr.open_dataset(ifiles_list[0])[xname].values
    ix = x.argsort()
    L = (x>=xlim[0]) & (x<=xlim[1])
    ixlim = ix[L][0], ix[L][-1]
    print(f'xlim: {xlim}; ixlim: {ixlim}')
    
    # lat range
    y = xr.open_dataset(ifiles_list[0])[yname].values
    iy = y.argsort()
    if yincrease is False:
        iy = iy[-1::-1]# reverse the indices when y is decreasing along axis
    L = (y>=ylim[0]) & (y<=ylim[1])
    iylim = iy[L][0], iy[L][-1]
    print(f'ylim: {ylim}; iylim: {iylim}')
    
    # extract the data
    ifiles = ' '.join(ifiles_list) # list -> string
    cmd = f'ncrcat -h -v {data_name} -d {xname},{ixlim[0]},{ixlim[1]} -d {yname},{iylim[0]},{iylim[1]} {ifiles} {ofile}'
    print(cmd, '...')
    s = os.system(cmd)
    if s==0:
        print('[OK]:', cmd)
    
    # mask the grids outside of the country
    da = xr.open_dataset(ofile)[data_name].rename({xname:'lon', yname:'lat'}).salem.roi(shape=shdf)
    
    # save the masked data to the ofile
    ofile_mask = ofile.replace(cntry_sname, cntry_sname+'.masked')
    da.to_dataset().to_netcdf(ofile_mask, unlimited_dims='time')
    print('[saved]:', ofile_mask)
    print('*done*')
    print()

if __name__ == '__main__':
    tformat = '%Y-%m-%d %H:%M:%S'
    t0 = datetime.now()
    print('[start]:', t0.strftime(tformat))

    print(model, expname, data_name, cntry_name)
    for en in ens:
        print('en: {en:02d}/{len(ens):02d}')
        get_region(cntry_name=cntry_name, expname=expname, model=model, data_name=data_name, en=en)
    
    t1 = datetime.now()
    print('[end]:', t1.strftime(tformat))
    print('[total time]:', f'{(t1-t0).seconds:,} seconds')