#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Thu Apr 28 12:26:29 EDT 2022
if __name__ == '__main__':
    import sys
    from misc.timer import Timer
    tt = Timer('start ' + ' '.join(sys.argv))
import sys, os.path, os, glob, datetime
import xarray as xr, numpy as np, pandas as pd, matplotlib.pyplot as plt
#more imports
import geoxarray
#
if __name__ == '__main__':
    tt.check('end import')
#
#start from here
daname = 'netrad_toa'
dsname = 'atmos_month'
model = 'AM2.5'
#model = 'FLOR'
expname = 'CTL1990s_tigercpu_intelmpi_18_540PE'
#years = range(101,151)
#years = range(101,119)
years = None
years = range(101,301)
#years = slice('101', '150')
scratch_on = False

if len(sys.argv)>1: daname = sys.argv[1]
if len(sys.argv)>2: model = sys.argv[2]
if len(sys.argv)>3: expname = sys.argv[3]
if 'scratch' in sys.argv: scratch_on = True
if expname in ('CTL1990_v201905_tigercpu_intelmpi_18_576PE',
    'CTL1990_v201905_2xCO2_m1p7pctSolar_tigercpu_intelmpi_18_576PE',
    'CTL1990_v201905_2xCO2_tigercpu_intelmpi_18_576PE'):
    scratch_on = False
    #years = range(101,201)
elif expname in ('CTL1990_v201905_2xCO2_m2pctSolar_tigercpu_intelmpi_18_576PE', ):
    #scratch_on = True
    #years = None
    pass
    

ens = None
if ens is None:
    idir = os.path.join('/tigress/wenchang/MODEL_OUT', model, expname, 'POSTP')
else:
    idir = os.path.join('/tigress/wenchang/MODEL_OUT', model, expname, f'en{ens:02d}', 'POSTP')
if scratch_on:
    idir = os.path.join('/home/wenchang/scratch', model, 'work', expname, 'POSTP')
if model in ('FLOR',) and not scratch_on:
    idir = idir.replace(f'{model}/', '')
#get the year range based on all files under the POSTP dir
if years is None:
    ncfiles = os.listdir(idir)
    ncfiles.sort()
    year0, year1 = int(ncfiles[0][:4]), int(ncfiles[-1][:4])
    if expname in ('CTL1990_v201905_2xCO2_m2pctSolar_tigercpu_intelmpi_18_576PE', ):
        year0 = int( sorted( glob.glob(f'{daname}_{model}_{expname}_*_glbmean.nc') )[-1].split('_')[-2].split('-')[-1] ) +1 #last year that has been downloaded
        if year0 > year1:
            print('data is updated')
            sys.exit()
    years = range(year0, year1+1)
ifiles = [os.path.join(idir, f'{year:04d}0101.{dsname}.nc') for year in years]

ofile = f'{daname}_{model}_{expname}_{years[0]:04d}-{years[-1]:04d}_glbmean.nc'
info = ' '.join( ofile.split('.')[0].split('_') )
if os.path.exists(ofile):
    da = xr.open_dataarray(ofile)
    print('[loaded]:', ofile)
else:
    print('loading...')
    da = xr.open_mfdataset(ifiles)[daname].load()
    attrs = da.attrs
    print('calculating...')
    da = da.geo.fldmean()
    da.attrs = attrs
    ds = da.to_dataset(name=daname)
    ds.to_netcdf(ofile)
    print('[saved]:', ofile)
    
 
if __name__ == '__main__':
    from wyconfig import * #my plot settings
    da.plot()
    
    #savefig
    if len(sys.argv)>1 and 'savefig' in sys.argv[1:]:
        figname = __file__.replace('.py', f'.png')
        if 'overwritefig' in sys.argv[1:]:
            wysavefig(figname, overwritefig=True)
        else:
            wysavefig(figname)
    tt.check(f'**Done**')
    print()
    if 'notshowfig' in sys.argv:
        pass
    else:
        plt.show()