#! /usr/bin/env python
import os, sys, glob, subprocess, fileinput, time, datetime, getpass

import nmmbUsh
import hwrf.hycom
import produtil.setup,  produtil.run, produtil.fileop, produtil.cd, produtil.batchsystem, produtil.config
from produtil.fileop import remove_file, make_symlink, isnonempty, chdir, fortlink
from produtil.run import run, exe, checkrun, mpi, mpirun
import nmmbUsh
import logging
from nmmbUsh import logger

USER=getpass.getuser()

produtil.setup.setup()

logger.info("forecast started")

workdir=os.environ.get('WORKhmon')
parmdir=os.environ.get('PARMhmon')

conf=produtil.config.from_file(parmdir + "/hmon.conf")
conf.cycle=os.environ.get('CYCLE')
hmondir=conf.get('dir','hmondir')
fixdir=conf.get('dir','FIXhmon')
exedir=conf.get('dir','EXEChmon')
comdir=conf.get('dir','COM')

with open(workdir + "/get_storm_info/storm_info", 'r') as f:
    for line in f:
        if "STORM_NAME=" in line:
           STORM_NAME = line.replace("STORM_NAME=", "")
           STORM_NAME = STORM_NAME.strip('\n')
        if "START_DATE=" in line:
           START_DATE = line.replace("START_DATE=", "")
           START_DATE = START_DATE.strip('\n')
        if "STORM_ID=" in line:
           STORM_ID = line.replace("STORM_ID=","")
           STORM_ID = STORM_ID.strip('\n')

astr = STORM_NAME
storm_name = astr.lower()
astr = STORM_ID
storm_id = astr.lower()

RUNDIR=workdir

chdir(RUNDIR + "/forecast")

### ocean stage data dir:
atmos_d = comdir 

datadir = atmos_d

hycomfix=fixdir + "/hycom"
hycomparm=parmdir + "/hycom"

STORM_BASIN=str(STORM_ID)[2:3]

if STORM_BASIN == "L":
    hycombasin = "rtofs_hat10"
elif STORM_BASIN == "E":
    hycombasin = "rtofs_hep20"
elif STORM_BASIN == "C":
    hycombasin = "rtofs_hcp70"
else:
    logger.info("Not in 3 hurricane basin")


######## prepair the files for NMMB forecast #########
nmmbUsh.make_symlinks(atmos_d + "/GWD_bin_*",".")

nmmbUsh.make_symlinks(atmos_d + "/*.configure_nest_details_0*",".")

for d in ['01', '02', '03']:
    f = open(storm_name + storm_id + "." + START_DATE + ".configure_nest_details_"+d, "r")
    text = f.read()
    f.close()
    with open("configure_file_" + d + ".in",'a') as f: f.write(text)

for d in ['01', '02', '03']:
    cf = "configure_file_"+d
    produtil.fileop.deliver_file("configure_file_"+d+".in",cf,logger=logger)
    for line in fileinput.FileInput(cf,inplace=1):
        line = line.replace("_COUPLE_","true")
        print(line, end=' ')

for filename in glob.glob(os.path.join(fixdir + "/nmmb/", '*')):
    produtil.fileop.deliver_file(filename, ".",logger=logger)


produtil.fileop.deliver_file("IGBP_VEGPARM.TBL", "VEGPARM.TBL",logger=logger)
produtil.fileop.deliver_file("IGBP_LANDUSE.TBL", "LANDUSE.TBL",logger=logger)
produtil.fileop.deliver_file(parmdir + "/solver_state.txt", "solver_state.txt",logger=logger)
produtil.fileop.deliver_file(parmdir + "/nests.txt", "nests.txt",logger=logger)

nmmbUsh.rm_glob('core*','PET*','nmmb_hst_*','nmmb_rst_*')

make_symlink("configure_file_01", "configure_file", True)
make_symlink("configure_file_01", "model_configure", True)
make_symlink("global_o3prdlos.f77", "fort.28", True)
make_symlink("global_o3clim.txt", "fort.48", True)

######## prepair the files for HYCOM forecast ##########
nestdir= RUNDIR + "/forecast/nest"
produtil.fileop.makedirs(nestdir)

nmmbUsh.rm_glob('archv.*','archs.*','restart_*','regional.*','forcing.*','relax.*','tbaric.*','veldf*','iso.sigma.*')

produtil.fileop.remove_file("ovrtn_out",logger=logger)

# ----------------------- make copy and link files ---------------------
for filename in ['ALBASE', 'FIS', 'ISLTYP', 'IVGTYP', 'MXSNAL', 'SM', 'TG', 'VEGFRC']:
    make_symlink(atmos_d + "/" + storm_name + storm_id + "." + START_DATE + "." + filename + "_3.nc", filename + "_3.nc", True)
    make_symlink(atmos_d + "/" + storm_name + storm_id + "." + START_DATE + "." + filename + "_9.nc", filename + "_9.nc", True)

for d in ['01', '02', '03']:
    filename = 'input_domain_' + d + '_nemsio'
    make_symlink(atmos_d + "/" + storm_name + storm_id + "." + START_DATE + "." + filename, filename, True)

fahr = 0
while (fahr < 127):
    hr='{0:0>4}'.format(fahr)
    filename = 'boco.' + hr
    make_symlink(atmos_d + "/" + storm_name + storm_id + "." + START_DATE + "." + filename, filename, True)
    fahr=fahr+6

produtil.fileop.deliver_file(hycomparm + "/hmon_" + hycombasin + ".basin.fcst.blkdat.input","blkdat.input",logger=logger)
produtil.fileop.deliver_file(hycomparm + "/hmon_" + hycombasin + ".basin.patch.input.90","patch.input",logger=logger)
produtil.fileop.deliver_file(hycomparm + "/hmon_" + hycombasin + ".basin.ports.input","ports.input",logger=logger)

produtil.fileop.deliver_file(hycomfix + "/hmon_" + hycombasin + ".basin.regional.grid.a","regional.grid.a",logger=logger)
produtil.fileop.deliver_file(hycomfix + "/hmon_" + hycombasin + ".basin.regional.grid.b","regional.grid.b",logger=logger)
produtil.fileop.deliver_file(hycomfix + "/hmon_" + hycombasin + ".basin.regional.depth.a","regional.depth.a",logger=logger)
produtil.fileop.deliver_file(hycomfix + "/hmon_" + hycombasin + ".basin.regional.depth.b","regional.depth.b",logger=logger)

produtil.fileop.deliver_file(datadir + "/" + storm_name + storm_id + "." + START_DATE + "." + hycombasin + ".restart.a", "restart_in.a",logger=logger)
produtil.fileop.deliver_file(datadir + "/" + storm_name + storm_id + "." + START_DATE + "." + hycombasin + ".restart.b", "restart_in.b",logger=logger)

for frcname in ['tauewd', 'taunwd', 'presur', 'airtmp', 'precip', 'surtmp', 'radflx', 'shwflx', 'vapmix', 'wndspd']:
    produtil.fileop.deliver_file(datadir + "/" + storm_name + storm_id + "." + START_DATE + ".forcing." + frcname + ".a", "forcing." + frcname + ".a",logger=logger)
    produtil.fileop.deliver_file(datadir + "/" + storm_name + storm_id + "." + START_DATE + ".forcing." + frcname + ".b", "forcing." + frcname + ".b",logger=logger)

for clmfrc in ['relax.ssh', 'relax.rmu']:
    produtil.fileop.deliver_file(hycomfix + "/hmon_" + hycombasin + ".basin." + clmfrc + ".a", clmfrc + ".a",logger=logger)
    produtil.fileop.deliver_file(hycomfix + "/hmon_" + hycombasin + ".basin." + clmfrc + ".b", clmfrc + ".b",logger=logger)

for clmfrc in ['tbaric', 'veldf2', 'veldf4', 'forcing.chl', 'forcing.offlux', 'forcing.rivers', 'iso.sigma']:
    produtil.fileop.deliver_file(hycomfix + "/hmon_" + hycombasin + ".basin." + clmfrc + ".a", clmfrc + ".a",logger=logger)
    produtil.fileop.deliver_file(hycomfix + "/hmon_" + hycombasin + ".basin." + clmfrc + ".b", clmfrc + ".b",logger=logger)

# nest 
logger.info("basin files datadir %s"%(repr(datadir)))
logger.info("basin files storm_name=%s storm_id=%s"%(repr(storm_name),repr(storm_id)))
logger.info("basin files START_DATE %s"%(repr(START_DATE)))
for bfile in glob.glob(datadir + "/" + storm_name + storm_id + "." + START_DATE + ".hmon_basin.*.b"): 
    fileroot=bfile.rsplit(".",1)[0]
    afile=fileroot + ".a"

    with open(bfile) as f:
        for i, line in enumerate(f,1):
            if i == 11:
                day=line.split()
                modelday=day[3]
    mydh=nmmbUsh.date_hycom2normal(logger,modelday)
    myd=(mydh)[0:8]
    yd=nmmbUsh.date2jday(logger,myd)

    myyyy=(yd)[0:4]
    mddd=(yd)[4:7]
    mhh=(mydh)[8:10]
    make_symlink(afile, nestdir + "/archv." + myyyy + "_" + mddd + "_" + mhh + ".a", True)
    make_symlink(bfile, nestdir + "/archv." + myyyy + "_" + mddd + "_" + mhh + ".b", True)

# link fix files into nest
make_symlink( RUNDIR + "/forecast/relax.rmu.a", nestdir + "/rmu.a", True)
make_symlink( RUNDIR + "/forecast/relax.rmu.b", nestdir + "/rmu.b", True)

# -------------------------- generate limits -----------------------------
with open("../nps_bc/namelist.nps", 'r') as f:
    for line in f:
        if " start_date = " in line:
            start_date = line.replace(" start_date = ", "")
            start_date = start_date.split(",")[0]
            start_date = start_date.replace("_", " ")
            start_date = start_date.replace("'", "")
            start_date = start_date.strip()
        if " end_date   = " in line:
            end_date = line.replace(" end_date   = ", "")
            end_date = end_date.split(",")[0]
            end_date = end_date.replace("_", " ")
            end_date = end_date.replace("'", "")
            end_date = end_date.strip()

nmmbUsh.create_limits(logger, start_date, end_date)
# --------------------------------------------------------------------------

############## copy link coupler files ####################
produtil.fileop.deliver_file( parmdir + "/cpl_nml", "cpl_nml", logger=logger)
produtil.fileop.deliver_file( parmdir + "/namelist.input", "namelist.input", logger=logger)
produtil.fileop.deliver_file( "../storm_center", ".", logger=logger)

############# prepare files for coupled run ####################
cmd = exe( hmondir +"/exec/hmon_preproc_cpl_ds_latlon_sm")
produtil.run.checkrun(cmd>"preproc_cpl.txt", logger=logger)
make_symlink("coupler_grid_data", "fort.65", True)

cmd = exe( hmondir +"/exec/hmon_preproc_parent_outer_nest_lats_lons_netcdf")
produtil.run.checkrun(cmd>"preproc_parent.txt", logger=logger)

############# enviromental set up #############################
os.environ['ATP_ENABLED']='1'
os.environ['FORT_BUFFERED']='TRUE'
os.environ['INSIDE_APRUN']='1'
#os.environ['KMP_AFFINITY']='disabled'
os.environ['OMP_NUM_THREADS']='1'
#os.environ['ESMF_RUNTIME_COMPLIANCECHECK']='OFF'
#os.environ['depth']='4'

os.environ['MALLOC_MMAP_MAX_']='0'
os.environ['MALLOC_TRIM_THRESHOLD_']='$(( 128 * 1024 * 1024 ))'

os.environ['MPICH_ABORT_ON_ERROR']='1'
os.environ['MPICH_ENV_DISPLAY']='1'
os.environ['MPICH_RANK_REORDER_DISPLAY']='1'
os.environ['MPICH_VERSION_DISPLAY']='1'

############# coupled run ##############################
NMM_exe=exedir + "/hmon_NMMB"
CPL_exe=exedir + "/hmon_wm3c"
HYC_exe=exedir + "/hmon_" + hycombasin + "_forecast"

date=datetime.datetime.utcnow()

logger.info("Model started: "+ date.strftime('%a %b %d %H:%M:%S UTC %Y'))
logger.info("CPL_exe:", CPL_exe)
logger.info("NMM_exe:", NMM_exe)
logger.info("HYC_exe:", HYC_exe)

checkrun(mpirun(mpi(CPL_exe).rpn(24)*4+mpi(NMM_exe).rpn(12)*434+mpi(HYC_exe).rpn(24)*90).out('out.forecast').err('err.forecast'),logger=logger)

########### copy data to COM directory #######################
produtil.fileop.deliver_file("./configure_file_01.in",atmos_d + "/" + storm_name + storm_id + "." + START_DATE + "." + "configure_file_01.in",logger=logger)
produtil.fileop.deliver_file("./configure_file_02.in",atmos_d + "/" + storm_name + storm_id + "." + START_DATE + "." + "configure_file_02.in",logger=logger)
produtil.fileop.deliver_file("./configure_file_03.in",atmos_d + "/" + storm_name + storm_id + "." + START_DATE + "." + "configure_file_03.in",logger=logger)
produtil.fileop.deliver_file("./nmmb_hst_01_nio_0006h_00m_00.00s",atmos_d + "/" + STORM_ID + "." + START_DATE + ".nmmb_hst_01_nio_0006h_00m_00.00s",logger=logger)
produtil.fileop.deliver_file("./nmmb_hst_02_nio_0006h_00m_00.00s",atmos_d + "/" + STORM_ID + "." + START_DATE + ".nmmb_hst_02_nio_0006h_00m_00.00s",logger=logger)
produtil.fileop.deliver_file("./nmmb_hst_03_nio_0006h_00m_00.00s",atmos_d + "/" + STORM_ID + "." + START_DATE + ".nmmb_hst_03_nio_0006h_00m_00.00s",logger=logger)
produtil.fileop.deliver_file("./swath_d02",atmos_d + "/" + storm_name + storm_id + "." + START_DATE + "." + "swath_d02",logger=logger)
produtil.fileop.deliver_file("./swath_d03",atmos_d + "/" + storm_name + storm_id + "." + START_DATE + "." + "swath_d03",logger=logger)

#abfiles=list()
#for globme in [ "*archv*.[ab]" ]:
#  globout=list(glob.glob(globme))
#  abfiles.extend(globout)

#produtil.fileop.deliver_file("./archv.*.a",atmos_d + "/" + storm_name + storm_id + "." + START_DATE + "." + ".",logger=logger)
#produtil.fileop.deliver_file("./archv.*.b",atmos_d + "/" + storm_name + storm_id + "." + START_DATE + "." + ".",logger=logger)


date=datetime.datetime.utcnow()
logger.info("Model ended: " + date.strftime('%a %b %d %H:%M:%S UTC %Y'))

logger.info("run_nmmb coupled finished")