#!/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

expnames = [
    'Pinatubo_PI_ens_noleap',
    'Agung_PI_ens_noleap',
    'StMaria_PI_ens_noleap',
    'Chichon_PI_ens_noleap',
    ]
data_name_o = 'SIE' # data name of output

def get_sie(expname, en, years=None):
    '''get sea ice extent from model output'''
    # ifiles
    if years is None:
        ifiles = glob.glob(f'/tigress/wenchang/MODEL_OUT/{expname}/en{en:02d}/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}/en{en:02d}/POSTP/{year:04d}0101.ice_month.nc'
            for year in years]
    
    # open dataset
    data_name = 'CN'
    with  xr.open_mfdataset(ifiles) as ds:
        da = ds[data_name]
        cell_area = ds['CELL_AREA'] # units are the Earth surface area
    attrs = da.attrs # save the raw attrs
    
    # CN -> SIC -> SIE 
    # global
    sie_glb = (da.sum('ct')*cell_area).sum(['xt', 'yt'])*100 # units are the percent of the Earth surface area
    sie_glb.attrs['long_name'] = f'global sea ice extent'
    sie_glb.attrs['units'] = '% of the Earth surface area'
    # NH
    sie_nh = (da.sum('ct')*cell_area).sel(yt=slice(0, None)).sum(['xt', 'yt'])*100 # units are the percent of the Earth surface area
    sie_nh.attrs['long_name'] = f'NH sea ice extent'
    sie_nh.attrs['units'] = '% of the Earth surface area'
    # SH
    sie_sh = (da.sum('ct')*cell_area).sel(yt=slice(None, 0)).sum(['xt', 'yt'])*100 # units are the percent of the Earth surface area
    sie_sh.attrs['long_name'] = f'SH sea ice extent'
    sie_sh.attrs['units'] = '% of the Earth surface area'

    # new dataset
    ds = xr.Dataset(dict(
        sie_glb=sie_glb,
        sie_nh=sie_nh,
        sie_sh=sie_sh
        ))
    
    return ds 

def get_sie_ens(expname):
    '''get sea ice extent ensembles from a given experiment name'''
    ens = range(1, 31)
    # ofile
    ofile = f'{expname}.{data_name_o}.nc'
    if os.path.exists(ofile):
        print('[exists]:', ofile)
        return
    dss = []
    for en in ens:
        print(f'en = {en:02d}')
        ds = get_sie(expname=expname, en=en, years=None)
        dss.append(ds)

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

    # save to netcdf
    ds.to_netcdf(ofile)
    print('[saved]:', ofile)

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

    for expname in expnames: 
        print(expname, data_name_o, tag)
        get_sie_ens(expname)

    # the end
    t1 = datetime.now()
    dt = t1 - t0
    print('[time used]:', f'{dt.seconds:,} seconds')
    print('[time end]:', t1.strftime(tformat))
    print()