#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Wed Nov 16 16:18:04 EST 2022
# get data as input for Kerry Emanuel's downscaling TC model
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
from misc.modelout import get_modelout_files
from misc.shell import run_shell
#
if __name__ == '__main__':
    tt.check('end import')
#
#start from here
#model = 'AM2.1p1'
#expname = 'amipHadISSTchancorr_rmWarming_ens_tigercpu_intelmpi_18_30PE'
#get model,expname based on cwd
cwd = os.getcwd()
model, expname = cwd.split('/')[-2:]
if model == 'MODEL_OUT':
    model = 'FLOR'
print(model, expname)#; sys.exit()

danames = ['temp', 'sphum','ps', 't_surf', 'ucomp', 'vcomp']#temp/sphum: monthly and plev convert; ps/t_surf: monthly and direct retrieve; ucomp/vcomp: daily and plev convert
dsname_default = 'atmos_month'
#ens = range(1,1+1)
ens = range(2,5+1)
for ii in ens:
    ifiles,years = get_modelout_files(model=model, expname=expname, dsname=dsname_default, ens=ii)
    print(ifiles[0], '\n->\n'+ifiles[-1]) 
    odir = os.path.dirname(ifiles[0]).replace('POSTP', 'analysis_wy/TCmodelInput')
    if not os.path.exists(odir):
        #os.mkdir(odir)
        os.makedirs(odir)
        print('[dir created]:', odir)
    for ifile in ifiles:
        for daname in danames:
            if daname in ('ucomp', 'vcomp'):
                ifile = ifile.replace(dsname_default, 'atmos_daily') #daily ucomp and vcomp are needed
            if daname in ('temp', 'sphum', 'ucomp', 'vcomp'): #plev convert
                ofile = os.path.join(odir, os.path.basename(ifile).replace('.nc', f'.plev.{daname}.nc'))
            else:
                ofile = os.path.join(odir, os.path.basename(ifile).replace('.nc', f'.{daname}.nc'))
            if os.path.exists(ofile):
                print('[exists]:', ofile)
                continue
            #print(ifile)
            #print(ofile)
            if daname in ('temp', 'sphum'):
                #plev convert
                cmd = f'/tigress/wenchang/bin/wyplev2p1.sh {ifile} {daname} {ofile}'
                run_shell(cmd)
                #deflate nc file
                cmd = f'ncks -O -7 -L 1 {ofile} {ofile}'
                run_shell(cmd)
            elif daname in ('ps', 't_surf'):
                cmd = f'ncks -O -7 -L 1 -v {daname} {ifile} {ofile}'
                run_shell(cmd)
            elif daname in ('ucomp', 'vcomp'):
                #plev convert
                #cmd = f'/tigress/wenchang/bin/wyplev2p1.sh {ifile} {daname} {ofile} "85000 20000" '
                cmd = f'/tigress/wenchang/bin/wyplev2p1.sh {ifile} {daname} {ofile} "85000 25000" '
                run_shell(cmd)
                #deflate nc file
                cmd = f'ncks -O -7 -L 1 {ofile} {ofile}'
                run_shell(cmd)
            

 
if __name__ == '__main__':
    #from wyconfig import * #my plot settings
    
    #savefig
    if 'savefig' in sys.argv or 's' in sys.argv:
        figname = __file__.replace('.py', f'.png')
        if 'overwritefig' in sys.argv or 'o' in sys.argv:
            wysavefig(figname, overwritefig=True)
        else:
            wysavefig(figname)
    tt.check(f'**Done**')
    print()
    if 'notshowfig' in sys.argv:
        pass
    else:
        pass#plt.show()