#!/usr/bin/env python
import os, os.path, sys, datetime
import numpy as np
import xarray as xr
import pandas as pd
from .geopotential import geopotential_height_make

#expname = sys.argv[1] #'Agung_PI_ens_noleap'
#compname = sys.argv[2] #'atmos_month'
#dataname = sys.argv[3] #'t_surf'
#en_start, en_end = sys.argv[4], sys.argv[5]

def year_shift(da, n, to_datetime=False):
    '''shift the time coords of da by n years'''
    da = da.copy()

    new_time = [t.replace(year=t.year+n) for t in da.time.values]
    da['time'] = new_time

    if to_datetime:
        new_time = [datetime.datetime(*t.timetuple()[0:6]) for t in da.time.values]
        da['time'] = new_time

    return da

def get_ctl_data(expname, compname, dataname, year_start, n_ens, n_years=5, plev=None):
    '''extract a varialbe from ctl experiment and tranform into ensembles'''
    ofile = f'data/{expname}.{compname}.{dataname}.nc'
    if plev:
        ofile = ofile.replace('.nc', f'.{plev}hPa.nc')
    if os.path.exists(ofile):
        print('\tExists:', ofile)
        return

    dirin = f'/tigress/wenchang/MODEL_OUT/{expname}'
    
    ens = range(1, n_ens+1)
    # combine ensemble members
    das = []
    for en in ens:
        ifiles = [f'{dirin}/POSTP/{year:04d}0101.{compname}.nc'
            for year in range(year_start + en -1, year_start + en -1 + n_years)]
        #print(ifiles)
        with xr.set_options(enable_cftimeindex=True):
            ds = xr.open_mfdataset(ifiles)
        if dataname == 'z_full':
            da = geopotential_height_make(ds)
        else:
            da = ds[dataname]
        if plev:
            #da = da.sel(pfull=plev, method='nearest')
            da = da.interp(pfull=plev) # WY: use interp to get the exact level
        da = year_shift(da, n=1-en+1-year_start, to_datetime=False)
        das.append(da)
    dae = xr.concat(das, pd.Index(ens, name='en'))
    
    new_dims = dict(grid_xt='lon', grid_yt='lat')
    dae.rename(new_dims).to_dataset().to_netcdf(ofile)
    print('Created:', ofile)

def get_volc_data(expname, compname, dataname, ens, plev=None):
    '''extract a variable from volcanic dataset and combine all the ensemble members'''
    ofile = f'data/{expname}.{compname}.{dataname}.nc'
    if plev:
        ofile = ofile.replace('.nc', f'.{plev}hPa.nc')
    if os.path.exists(ofile):
        print('\tExists:', ofile)
        return

    dirin = f'/tigress/wenchang/MODEL_OUT/{expname}'

    # combine ensemble members
    das = []
    for en in ens:
        ifiles = f'{dirin}/en{en:02d}/POSTP/*.{compname}.nc'
        #print(ifiles)
        with xr.set_options(enable_cftimeindex=True):
            ds = xr.open_mfdataset(ifiles)
        if dataname == 'z_full':
            da = geopotential_height_make(ds)
        else:
            da = ds[dataname]
        if plev:
            #da = da.sel(pfull=plev, method='nearest')
            da = da.interp(pfull=plev) # WY: use interp to get the exact level
        das.append(da)
    dae = xr.concat(das, pd.Index(ens, name='en'))

    # convert to Python datetime
    new_time = [datetime.datetime(*t.replace(year=t.year).timetuple()[0:6]) 
                            for t in da.time.values]
    dae['time'] = new_time
    
    new_dims = dict(grid_xt='lon', grid_yt='lat')
    dae.rename(new_dims).to_dataset().to_netcdf(ofile)
    print('Created:', ofile)

if __name__ == '__main__':
    print('Functions included in wylib:')
    funcs = ('year_shift', 'get_ctl_data', 'get_volc_data',)
    for f in funcs:
        print(f)