#! /usr/bin/env python

import os, sys, glob, fileinput, time

import produtil.setup,  produtil.run, produtil.fileop, produtil.cd, produtil.batchsystem, produtil.config, produtil.numerics
from produtil.fileop import remove_file, make_symlink, isnonempty, chdir, fortlink
from produtil.run import run, exe
from produtil.numerics import to_datetime_rel, to_datetime
import nmmbUsh

import logging

from nmmbUsh import logger

produtil.setup.setup()

logger.info("nps bc started")

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

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')
INPUT=conf.get('dir','COMgfs')

wgrib2=os.environ.get("WGRIB2")

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=","")
           stormid = STORM_ID.strip('\n')

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

SDATE=produtil.numerics.to_datetime(START_DATE)
START_YYYY=SDATE.strftime("%Y")
START_MM=SDATE.strftime("%m")
START_DD=SDATE.strftime("%d")
START_HH=SDATE.strftime("%H")

NPS_START_DATE=START_YYYY+"-"+START_MM+"-"+START_DD+"_"+START_HH+":00:00"

chdir(workdir + "/nps_bc")

print(os.getcwd())

#
# ungrib
#
h0=0
h1=126
dh=6
hh=h0

produtil.fileop.deliver_file(fixdir + "/nps/Vtable.GFS", "Vtable",logger=logger)
produtil.fileop.deliver_file(fixdir +"/nps/METGRID.TBL.NMB","METGRID.TBL",logger=logger)
nmmbUsh.rm_glob('met_nmb.d01*dio','metgrid.log*')

#
# geogrid
#
gfile="../nps_ic/geo_nmb.d03.nc"
while not(os.path.isfile(gfile)) or os.path.getsize(gfile)<=25000000 :
    logger.info(gfile +" not ready, sleep 15")
    time.sleep(15)

nmmbUsh.make_symlinks("../nps_ic/b_r??_*", ".")
nmmbUsh.make_symlinks("../nps_ic/nest_start_*", ".")
nmmbUsh.make_symlinks("../nps_ic/geo_nmb.*.nc", ".")
#


fahr = 0
n_gribfiles = 0
while (fahr < 127):
    hr='{0:0>3}'.format(fahr)

    CDATE=produtil.numerics.to_datetime_rel(fahr*3600, START_DATE)
    RUN_YYYY=CDATE.strftime("%Y")
    RUN_MM=CDATE.strftime("%m")
    RUN_DD=CDATE.strftime("%d")
    RUN_HH=CDATE.strftime("%H")
    NPS_RUN_DATE=RUN_YYYY+"-"+RUN_MM+"-"+RUN_DD+"_"+RUN_HH+":00:00"

# namelist.nps 
    produtil.fileop.deliver_file("namelist.nps.in","namelist.nps",logger=logger)
    for line in fileinput.FileInput("namelist.nps",inplace=1):
        line = line.replace("NPS_START_DATE",NPS_RUN_DATE)
        line = line.replace("NPS_END_DATE",NPS_RUN_DATE)
        print(line, end=' ')

# check if gfs grib2 file is ready
    if fahr == 0:
        file = INPUT + "/gfs.t"+START_HH+"z.pgrb2.0p25.anl"
        f_file = INPUT + "/gfs.t"+START_HH+"z.pgrb2.0p25.f000"
        ingribfile_tmp = "gfs.t"+START_HH+"z.pgrb2.0p25.f000"
        ingribfile_anl = "gfs.t"+START_HH+"z.pgrb2.0p25.anl"
        timeout= time.time() + 60*15
        while not(os.path.isfile(file)) or os.path.getsize(file)<=50000000 :
            if time.time() > timeout:
                sys.exit("FATAL ERROR: " + file + " is missing.")
            logger.info(file +" not ready, sleep 15")
            time.sleep(15)
        logger.info(file)
        produtil.fileop.deliver_file(f_file, ingribfile_tmp, logger=logger)
        produtil.fileop.deliver_file(file, ingribfile_anl, logger=logger)
        outgribfile = "f2anl.grb2"
        cmd = wgrib2 + ' ' + ingribfile_tmp + ' | egrep ":TMP:surface:|SOIL|SNOD|LAND|ICEC|PBL|:UGRD:10 m above|:VGRD:10 m above|:TMP:2 m above|:RH:2 m above|WEASD" | ' + wgrib2 + ' -i ' + ingribfile_tmp + ' -grib ' + outgribfile
        os.system(cmd)
        cmd = 'cat ' + ingribfile_anl + ' ' + outgribfile + ' > final.grb2'
        os.system(cmd)
        ingribfile = "final.grb2"
     
    else:
        file = INPUT + "/gfs.t"+START_HH+"z.pgrb2.0p25.f" + hr
        ingribfile = "gfs.t"+START_HH+"z.pgrb2.0p25.f" + hr
        timeout = time.time() + 60*15    # 15 minutes from now
        while not(os.path.isfile(file)) or os.path.getsize(file)<=50000000 :
            if time.time() > timeout:
                sys.exit("FATAL ERROR: " + file + " is missing.")
            logger.info(file +" not ready, sleep 15")
            time.sleep(15)

        logger.info(file)

        produtil.fileop.deliver_file(file, ingribfile, logger=logger)

# link fixed grib file to GRIBFILE.AAA   
    make_symlink(ingribfile,"GRIBFILE.AAA",True)

#
# ungrib
#
    logger.info("mpirun" + " " + "ungrib" + str(fahr))
    produtil.run.checkrun(exe(exedir + "/hmon_ungrib").out('out.ungrib').err('err.ungrib'),logger=logger)

#
# metgrid
#

    logger.info("mpirun" + " " + "metgrid")
    produtil.run.checkrun(produtil.run.mpirun(produtil.run.mpi(exedir + "/hmon_metgrid"),allranks=True).out("out.metgrid").err("err.metgrid"),logger=logger)

    os.remove("GRIBFILE.AAA")
    fahr=fahr+6

#
# nemsinterp
#
nmmbUsh.make_symlinks("../nps_ic/met_nmb.d02.*.nc", ".")
nmmbUsh.make_symlinks("../nps_ic/met_nmb.d03.*.nc", ".")
produtil.fileop.deliver_file("namelist.nps.in","namelist.nps",logger=logger)
for line in fileinput.FileInput("namelist.nps",inplace=1):
    line = line.replace("NPS_START_DATE",NPS_START_DATE)
    line = line.replace("NPS_END_DATE",NPS_RUN_DATE)
    line = line.replace("max_dom = 1,","max_dom = 3,")
    print(line, end=' ')

nmmbUsh.rm_glob('boco.*','GWD_bin*','input_domain_*','configure_nest_details*','domain_details*')

logger.info("mpirun" + " " + "nemsinterp")
produtil.run.checkrun(produtil.run.mpirun(produtil.run.mpi(exedir + "/hmon_nemsinterp"),allranks=True).out('out.nemsinterp').err('err.nemsinterp'),logger=logger)

##
## copy input files to com directory
##

atmos_d = comdir

produtil.fileop.makedirs(atmos_d,logger=logger)

produtil.fileop.deliver_file("./configure_nest_details_01",atmos_d + "/" + storm_name + storm_id + "." + START_DATE + "." + "configure_nest_details_01",logger=logger)
produtil.fileop.deliver_file("./configure_nest_details_02",atmos_d + "/" + storm_name + storm_id + "." + START_DATE + "." + "configure_nest_details_02",logger=logger)
produtil.fileop.deliver_file("./configure_nest_details_03",atmos_d + "/" + storm_name + storm_id + "." + START_DATE + "." + "configure_nest_details_03",logger=logger)
produtil.fileop.deliver_file("./namelist.nps",atmos_d + "/" + storm_name + storm_id + "." + START_DATE + "." + "namelist.nps_bc",logger=logger)

for ifile in glob.glob("boco.*"):
    produtil.fileop.deliver_file(ifile,atmos_d + "/" + storm_name + storm_id + "." + START_DATE + "." + ifile,logger=logger)

logger.info("run_nps_bc finished")