#!/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_gpi(expname, dataname, year_start, n_ens, n_years=5):
    '''extract a varialbe from ctl experiment and tranform into ensembles'''
    ofile = f'data/{expname}.GPI.{dataname}.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}/analysis_wy/GPI/{year:04d}0101.GPI.nc'
            for year in range(year_start + en -1, year_start + en -1 + n_years)]
        #print(ifiles)
        ds = xr.open_mfdataset(ifiles)
        da = ds[dataname]
        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_gpi(expname, dataname, ens):
    '''extract a variable from volcanic dataset and combine all the ensemble members'''
    ofile = f'data/{expname}.GPI.{dataname}.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}/analysis_wy/GPI/*.GPI.nc'
        #print(ifiles)
        ds = xr.open_mfdataset(ifiles)
        da = ds[dataname]
        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)