#! /usr/bin/env python

import os, sys, glob, shutil, fileinput

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

produtil.setup.setup()

logger.info("forecast started")

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

chdir(workdir + "/forecast")

conf=produtil.config.from_file(parmdir + "/hmon.conf")
conf.cycle=os.environ.get('CYCLE')
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()

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

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)
produtil.fileop.deliver_file(parmdir + "/cpl_nml_uncoupled", "cpl_nml",logger=logger)

atmos_d = comdir 

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

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)

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

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("_RESTART_","false")
        line = line.replace("_COUPLE_","false")
        print(line, end=' ')

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)

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'


logger.info("mpirun" + " " + "/NMMB.x 1>out.nmmb 2>err.nmmb")
checkrun(mpirun(mpi(exedir+"/hmon_NMMB").rpn(12)*434).out('out.nmmb').err('err.nmmb'),logger=logger)
#produtil.run.checkrun(produtil.run.mpirun(produtil.run.mpi(exedir+"/hmon_NMMB").rpn(12)*412).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)

logger.info("run_nmmb finished")