#!/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
def get_modelout_files(daname, model, expname, dsname='atmos_month', years=None, ens=None, return_years=False):
    #daname = 'netrad_toa'
    #dsname = 'atmos_month'
    #model = 'AM2.5'
    #expname = 'CTL1990s_tigercpu_intelmpi_18_540PE'
    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 model in ('FLOR',):
        idir = idir.replace(f'{model}/', '')
    if not os.path.exists(idir): #data not in modelout dir; maybe in work dir
        idir = idir.replace('/tigress/wenchang/MODEL_OUT', '/home/wenchang/scratch') \
            .replace(f'{expname}', f'work/{expname}')
        if model in ('FLOR',):
            idir = idir.replace('home/wenchang/scratch', 'home/wenchang/scratch/FLOR')
        if ens is not None:
            idir = idir.replace(f'en{ens:02d}/', '') \
                .replace('_tigercpu_', f'_e{ens}_tigercpu_')
        if not os.path.exists(idir): #data not in work dir either
            return
    #if scratch_on:
    #    idir = os.path.join('/home/wenchang/scratch', model, 'work', expname, 'POSTP')
    #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])
        years = range(year0, year1+1)
    ifiles = [os.path.join(idir, f'{year:04d}0101.{dsname}.nc') for year in years]

    if return_years:
        return ifiles, years
    else:
        return ifiles

 
if __name__ == '__main__':
    #from wyconfig import * #my plot settings
    daname = 't_surf'
    model = 'AM2.5'
    expname = 'CTL1990s_tigercpu_intelmpi_18_540PE'
    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]
    ifiles, years = get_modelout_files(daname, model, expname, return_years=True)
    print('**ifiles**')
    for ifile in ifiles: 
        print(ifile)
    print('**years**')
    print(years)
    
    #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()
    #plt.show()