#! /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 ic started")

workdir=os.environ.get('WORKhmon')
parmdir=os.environ.get('PARMhmon')
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=", "")
           cycle = 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(str(cycle))
START_YYYY=SDATE.strftime("%Y")
START_MM=SDATE.strftime("%m")
START_DD=SDATE.strftime("%d")
START_HH=SDATE.strftime("%H")

chdir(workdir+"/nps_ic",logger=logger)

print(os.getcwd())

#
# geogrid
#
produtil.fileop.deliver_file(fixdir +"/nps/GEOGRID.TBL.NMB", "GEOGRID.TBL",logger=logger)

make_symlink("testb.nml", "fort.81", True)

nmmbUsh.rm_glob('geo_nmb.*','geogrid.log*','b_d??_*','GWD*','nest_start_*')

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

### link GFS grib data to NPS IC -------

file = INPUT + "/gfs.t"+START_HH+"z.pgrb2.0p25.f006"
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)

for file in glob.glob(INPUT + "/gfs.t"+START_HH+"z.pgrb2.0p25.anl"):
    if os.path.isfile(file):
        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"
        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"
        make_symlink(ingribfile,"GRIBFILE.AAA",True)

for file in glob.glob(INPUT + "/gfs.t"+START_HH+"z.pgrb2.0p25.f006"):
    if os.path.isfile(file):
        ingribfile = "gfs.t"+START_HH+"z.pgrb2.0p25.f006"
        produtil.fileop.deliver_file(file, ingribfile, logger=logger)
        make_symlink(ingribfile,"GRIBFILE.AAB",True)

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

#
# ungrib 
#
produtil.fileop.deliver_file(fixdir + "/nps/Vtable.GFS", "Vtable",logger=logger)


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

#
# metgrid
#
produtil.fileop.deliver_file(fixdir +"/nps/METGRID.TBL.NMB","METGRID.TBL",logger=logger)

nmmbUsh.rm_glob('met_nmb.d01*dio','metgrid.log*')

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)

#
# nemsinterp
#
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)

atmos_d = comdir 

produtil.fileop.makedirs(atmos_d,logger=logger)
produtil.fileop.deliver_file("./input_domain_01_nemsio",atmos_d + "/" + storm_name + storm_id + "." + cycle + "." + "input_domain_01_nemsio_nps",logger=logger)
produtil.fileop.deliver_file("./input_domain_02_nemsio",atmos_d + "/" + storm_name + storm_id + "." + cycle + "." + "input_domain_02_nemsio_nps",logger=logger)
produtil.fileop.deliver_file("./input_domain_03_nemsio",atmos_d + "/" + storm_name + storm_id + "." + cycle + "." + "input_domain_03_nemsio_nps",logger=logger)
produtil.fileop.deliver_file("./namelist.nps",atmos_d + "/" + storm_name + storm_id + "." + cycle + "." + "namelist.nps_ic",logger=logger)
produtil.fileop.deliver_file("./testb.nml",atmos_d + "/" + storm_name + storm_id + "." + cycle + "." + "testb.nml",logger=logger)
produtil.fileop.deliver_file("./GWD_bin_01",atmos_d + "/GWD_bin_01",logger=logger)
produtil.fileop.deliver_file("./GWD_bin_02",atmos_d + "/GWD_bin_02",logger=logger)
produtil.fileop.deliver_file("./GWD_bin_03",atmos_d + "/GWD_bin_03",logger=logger)

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

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

logger.info("run_nps_ic finished")