#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Wed Sep 18 21:10:30 EDT 2019
from datetime import datetime
import os.path, sys, os, glob
#import matplotlib.pyplot as plt
import xarray as xr, numpy as np, pandas as pd
#import geoxarray

expname = 'CTL1860_noleap_tigercpu_intelmpi_18_576PE'
data_name = 'SIC'
year_start_ctl = 11
n_years = 5
ens = range(1, 31)
n_ens = len(ens) 
time_new = xr.cftime_range(start='0001-01', periods=n_years*12, freq='MS', calendar='noleap')
ofile = f'{expname}.{data_name}.ens.nc'

def get_sic(expname, years=None):
    '''get sea ice extent from model output'''
    # ifiles
    if years is None:
        ifiles = glob.glob(f'/tigress/wenchang/MODEL_OUT/{expname}/POSTP/????0101.ice_month.nc')
        ifiles.sort()
        years = [int( s.split('/')[-1][0:4] ) for s in ifiles]
    else:
        ifiles = [f'/tigress/wenchang/MODEL_OUT/{expname}/POSTP/{year:04d}0101.ice_month.nc'
            for year in years]
    
    # open dataset
    with  xr.open_mfdataset(ifiles) as ds:
        da = ds['CN']
        attrs = da.attrs
    with xr.open_dataset(ifiles[0]) as ds:
        CELL_AREA, GEOLAT, GEOLON = ds['CELL_AREA'], ds['GEOLAT'], ds['GEOLON']
    
    # CN -> SIC 
    sic = da.sum('ct')
    sic.attrs = attrs

    # new ds
    ds = sic.to_dataset(name=data_name)

    return ds, CELL_AREA, GEOLAT, GEOLON 

if __name__ == '__main__':
    if os.path.exists(ofile):
        print('[exists]:', ofile)
        sys.exit()

    # record the start time
    tformat = '%Y-%m-%d %H:%M:%S'
    t0 = datetime.now()
    print('[time start]:', t0.strftime(tformat))


    dss = []
    for en in ens:
        print(f'en = {en:02d}')
        years = range(year_start_ctl+en-1, year_start_ctl+en-1+5)
        ds, CELL_AREA, GEOLAT, GEOLON = get_sic(expname=expname, years=years)
        ds = ds.assign_coords(time=time_new)
        dss.append(ds)

    # concatenated dataset
    ds = xr.concat(dss, dim=pd.Index(ens, name='en'))
    ds['en'].attrs['long_name'] = 'ensemble member'

    # 
    ds['CELL_AREA'] = CELL_AREA
    ds['GEOLAT'] = GEOLAT
    ds['GEOLON'] = GEOLON
    
    # save the new dataset
    ds.to_netcdf(ofile)
    print('[saved]:', ofile)
    
    # the end
    t1 = datetime.now()
    dt = t1 - t0
    print('[time used]:', f'{dt.seconds:,} seconds')
    print('[time end]:', t1.strftime(tformat))
    print()