#!/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 = [
    'co2x2_CTL1860_tigercpu_intelmpi_18_576PE',
    'co2xp5_CTL1860_tigercpu_intelmpi_18_576PE',
    'm6p0sol_CTL1860_tigercpu_intelmpi_18_576PE',
    'p6p0sol_CTL1860_tigercpu_intelmpi_18_576PE',
    'm6p0sol_CTL1860_recover0591_tigercpu_intelmpi_18_576PE',
    'm6p0sol_CTL1860_recover0501_tigercpu_intelmpi_18_576PE',
    'm6p0sol_CTL1860_recover0251_tigercpu_intelmpi_18_576PE',
    'm1p0sol_CTL1860_tigercpu_intelmpi_18_576PE',
    'p1p0sol_CTL1860_tigercpu_intelmpi_18_576PE',
    'CTL1860_newdiag_tigercpu_intelmpi_18_576PE',
    ]
"""
#deprecated
year_ranges = [range(101, 597+1),
    range(101, 400+1),
    range(501, 900+1),
    range(251, 400+1),
    range(101, 600+1),
    range(101, 600+1),
    range(1, 1000+1),
    ]
"""
data_names = ['netrad_toa', 't_surf'][1:2]
tag = 'glbMean'

def get_glbMean(expname, data_name, years=None):
    '''get global mean of a variable from model output'''
    # input files
    if years is None:
        ifiles = glob.glob(f'/tigress/wenchang/MODEL_OUT/{expname}/POSTP/????0101.atmos_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.atmos_month.nc'
            for year in years]
    print('[years]:', f'{years[0]:04d} - {years[-1]:04d}')
    
    # output files    
    ofile = f'{expname}.{data_name}.{tag}.{years[0]:04d}-{years[-1]:04d}.nc'
    if os.path.exists(ofile):
        print('[exists]:', ofile)
        return
    
    # open dataset
    with  xr.open_mfdataset(ifiles) as ds:
        da = ds[data_name]
    attrs = da.attrs # save the raw attrs
    
    # computation
    da = da.rename({'grid_xt': 'lon', 'grid_yt': 'lat'}) \
        .geo.fldmean().load()
    
    # save to netcdf
    da.attrs = attrs
    encoding = {data_name: {'zlib': True, 'complevel': 1}}
    da.to_dataset(name=data_name) \
        .to_netcdf(ofile, encoding=encoding)
    print('[saved]:', ofile)

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

    #for expname,years in zip(expnames, year_ranges):
    for expname in expnames:
        for data_name in data_names:
            print()
            print(expname, data_name, tag)
            get_glbMean(expname, data_name, years=None)

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