#! /usr/bin/env python
import os, sys, stat, glob, subprocess, shutil, getpass, fileinput, logging

import hmon_vitals
import nmmbUsh
import produtil.setup,  produtil.run, produtil.fileop, produtil.numerics, produtil.config
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
from nmmbUsh import logger

import logging

produtil.setup.setup()

logger.info("launch started")

USER=getpass.getuser()

if len(sys.argv)!=3:
    logger.info("need two inputs, e.g., run_zeus-my1.sh 03E 2014060918")
    sys.exit()

STORM_ID=sys.argv[1]
START_DATE=sys.argv[2]

STORM_NUM=str(STORM_ID)[0:2]
STORM_BASIN=str(STORM_ID)[2:3]

FLENGTH="126"
ILENGTH="6"

RELOCATE="YES"

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

conf=produtil.config.from_file(parmdir + "/hmon.conf")
conf.cycle=os.environ.get('CYCLE')
expname=conf.get('config','expname')
USE_MESSAGE=conf.get('config','use_message')
stormlabel=conf.get('config','stormlabel')
label=stormlabel[5]
fixdir=conf.get('dir','FIXhmon')
scriptdir=conf.get('dir','EXhmon')
SYNDAT=conf.get('dir','syndat')
comdir=conf.get('dir','COM')
#atcfdir=conf.get('dir','TRACK')
INPUT=conf.get('dir','COMgfs')
msgdir=conf.get('dir','COMmsg')

hwrfush=ushdir
nps_geog=fixdir + "/nps_geog"

###############################################################################
MYDIR=scriptdir

EXPT=expname

#produtil.fileop.makedirs(atcfdir)

#
# get storm info
#
chdir(scriptdir)

SDATE=produtil.numerics.to_datetime(START_DATE)
START_YMD=SDATE.strftime("%Y%m%d%H")
START_YMDHM=START_YMD + "00"
START_YEAR=SDATE.strftime("%Y")
START_YY=SDATE.strftime("%y")
START_CC=SDATE.strftime("%C")
START_MM=SDATE.strftime("%m")
START_DD=SDATE.strftime("%d")
START_HH=SDATE.strftime("%H")

CDATE_M06=produtil.numerics.to_datetime_rel(6*-3600, START_DATE)
START_DATE_M06=CDATE_M06.strftime("%Y%m%d%H")

RUNDIR=workdir 
produtil.fileop.makedirs(RUNDIR,logger=logger)
chdir(RUNDIR)

## move get_storm_info to running directory 
if os.path.isfile("get_strom_info"):
    os.chmod("get_strom_info",stat.S_IRWXO)
    produtil.fileop.remove_file("get_storm_info")
elif os.path.isdir("get_strom_info"):
    shutil.rmtree("get_storm_info", ignore_errors=True)
produtil.fileop.makedirs( "get_storm_info",logger=logger)

chdir("get_storm_info")

### run hmon_vitals.py instead of storm_location-vit.py to get storm information from tcvitals
os.environ.update(
#    COMINarch='/gpfs/tp1/nco/ops/com/arch/prod/syndat', # or gp1 on gyre/ tp1 on tide
#    COMINarch=SYNDAT, # or gp1 on gyre/ tp1 on tide
#    COMINmsg='/gpfs/hps/nco/ops/com/hur/prod/inphwrf',
    COMINmsg=msgdir,
#    storm_num='1' # request storm in slot 1
    storm_num=label # request storm in slot 1
    )

if USE_MESSAGE=='NO':
    storm = hmon_vitals.vitals_for_storm(STORM_NUM,STORM_BASIN,START_DATE)
    produtil.fileop.deliver_file(SYNDAT + "/syndat_tcvitals." + START_YEAR, workdir + "/syndat_tcvitals." + START_YEAR, logger=logger)
else:
    storm = hmon_vitals.vitals_for_slot()
    produtil.fileop.deliver_file(msgdir + "message" + label, workdir + "/tmpvit" + label, logger=logger) 
    with open(workdir + "/syndat_tcvitals." + START_YEAR,'w') as f:
        with open(msgdir + "message" + label, 'r') as rfile:
            for line in rfile.readlines():
                line = line[:19] + '20' + line[19:]
                f.write(line)

if storm is None:
    logger.warning("Cannot find Storm!!")
else:
    logger.info("Storm: storm %s%s %s %s at lat=%.1f and lon=%.1f"%(
        storm.STORM_NUMBER,storm.STORM_BASIN,storm.STORM_NAME,
        storm.YYYY, storm.STORM_LAT, storm.STORM_LON))

storm06 = hmon_vitals.vitals_for_storm(STORM_NUM,STORM_BASIN,START_DATE_M06)
if storm06 is None:
    logger.info("This is the first cycle!!")

if storm06 is None:
    STORM_NAME06=storm.STORM_NAME
else:
    STORM_NAME06=storm06.STORM_NAME

with open("storm_info",'w') as f:
    f.write("START_DATE=" + START_DATE + "\n")
    f.write("STORM_NAME=" + storm.STORM_NAME + "\n")
    f.write("STORM_NUMBER=" + storm.STORM_NUMBER + "\n")
    f.write("STORM_ID=" + STORM_ID + "\n")
    f.write("STORM_LAT=" + str(storm.STORM_LAT) + "\n")
    f.write("STORM_LON=" + str(storm.STORM_LON) + "\n")
    f.write("STORM_LON_PARENT=" + str(storm.STORM_LON_PARENT) + "\n")
    f.write("STORM_NAME06=" + STORM_NAME06 + "\n")
    f.write("START_DATE06=" + START_DATE_M06 + "\n")
    f.write("CEN_LAT=" + str(storm.CEN_LAT) + "\n")
    f.write("CEN_LON=" + str(storm.CEN_LON) + "\n")
    f.write("NEST1_LAT=" + str(storm.NEST1_LAT) + "\n")
    f.write("NEST1_LON=" + str(storm.NEST1_LON) + "\n")
    f.write("NEST2_LAT=" + str(storm.NEST2_LAT) + "\n")
    f.write("NEST2_LON=" + str(storm.NEST2_LON))

produtil.fileop.makedirs(comdir,logger=logger)
with open(comdir + "/storm" + label + "." + "storm_info",'w') as f:
    f.write("export CYCLE=" + START_DATE + "\n")
    f.write("export STORM_NAME=" + storm.STORM_NAME + "\n")
    f.write("export STORM_ID=" + STORM_ID + "\n")

with open("storm_info", 'r') as f:
    for line in f:
        if "CEN_LAT=" in line:
            CEN_LAT = line.replace("CEN_LAT=", "")
            CEN_LAT = float(CEN_LAT)
            if "CEN_LON=" in line:
                CEN_LON = line.replace("CEN_LON=", "")
                CEN_LON = float(CEN_LON)

        if "NEST1_LAT=" in line:
            NEST1_LAT = line.replace("NEST1_LAT=", "")
            NEST1_LAT = float(NEST1_LAT)

        if "NEST2_LAT=" in line:
            NEST2_LAT = line.replace("NEST2_LAT=", "")
            NEST2_LAT = float(NEST2_LAT)

        if "NEST1_LON=" in line:
            NEST1_LON = line.replace("NEST1_LON=", "")
            NEST1_LON = float(NEST1_LON)

        if "NEST2_LON=" in line:
            NEST2_LON = line.replace("NEST2_LON=", "")
            NEST2_LON = float(NEST2_LON)

CEN_LAT = str(storm.CEN_LAT)
CEN_LON = str(storm.CEN_LON)
NEST1_LAT = str(storm.NEST1_LAT)
NEST2_LAT = str(storm.NEST2_LAT)
NEST1_LON = str(storm.NEST1_LON)
NEST2_LON = str(storm.NEST2_LON)

chdir("../")

f=open('storm_center','w')
f.write(CEN_LAT)
f.write("\n")
f.write(CEN_LON)
f.close()

EDATE=produtil.numerics.to_datetime_rel(int(FLENGTH)*3600, int(START_DATE))
IDATE=produtil.numerics.to_datetime_rel(int(ILENGTH)*3600, int(START_DATE))

START_YYYYMMDD=SDATE.strftime("%Y%m%d")
START_YYYY=SDATE.strftime("%Y")

END_YYYY=EDATE.strftime("%Y")
END_MM=EDATE.strftime("%m")
END_DD=EDATE.strftime("%d")
END_HH=EDATE.strftime("%H")

IND_YYYY=IDATE.strftime("%Y")
IND_MM=IDATE.strftime("%m")
IND_DD=IDATE.strftime("%d")
IND_HH=IDATE.strftime("%H")

NPS_START_DATE=START_YYYY+"-"+START_MM+"-"+START_DD+"_"+START_HH+":00:00"
NPS_END_DATE=END_YYYY+"-"+END_MM+"-"+END_DD+"_"+END_HH+":00:00"
NPS_IND_DATE=IND_YYYY+"-"+IND_MM+"-"+IND_DD+"_"+IND_HH+":00:00"

#####------------------ NPS for IC --------------------------
produtil.fileop.makedirs(RUNDIR+"/nps_ic")
chdir(RUNDIR+"/nps_ic")

###### ------------ NPS namelist for IC ------------------------------

for f in ['namelist.nps']:
    if os.path.isfile(f):
        produtil.fileop.remove_file(f)
    if STORM_BASIN == "L":
        produtil.fileop.deliver_file(parmdir + "/" + f + "_AL.in", f, logger=logger)
    elif STORM_BASIN == "E":
        produtil.fileop.deliver_file(parmdir + "/" + f + "_EP.in", f, logger=logger)
    elif STORM_BASIN == "C":
        produtil.fileop.deliver_file(parmdir + "/" + f + "_EP.in", f, logger=logger)
    else:
        logger.info("Not in 3 hurricane basin")

    for line in fileinput.FileInput(f, inplace=1):
        line = line.replace("NPS_START_DATE",NPS_START_DATE)
        line = line.replace("NPS_END_DATE",NPS_IND_DATE)
        line = line.replace("_CEN_LAT_",CEN_LAT)
        line = line.replace("_CEN_LON_",CEN_LON)
        line = line.replace("_NEST1_LAT_",NEST1_LAT)
        line = line.replace("_NEST1_LON_",NEST1_LON)
        line = line.replace("_NEST2_LAT_",NEST2_LAT)
        line = line.replace("_NEST2_LON_",NEST2_LON)
        line = line.replace("NPS_GEOG",nps_geog)
        line = line.replace("MAX_DOMS","3")
        print(line, end=' ')

for f in ['testb.nml']:
    if os.path.isfile(f):
        produtil.fileop.remove_file(f)
    produtil.fileop.deliver_file(parmdir + "/" + f + ".in", f, logger=logger)

    for line in fileinput.FileInput(f, inplace=1):
        line = line.replace("NPS_GEOG",nps_geog)
        print(line, end=' ')

chdir ("../")

##### ----------------- NPS for BC -----------------------------

produtil.fileop.makedirs(RUNDIR+"/nps_bc")
chdir(RUNDIR+"/nps_bc")

###### ----------------- NPS namelist for BC -----------------------
if STORM_BASIN == "L":
    produtil.fileop.deliver_file(parmdir+"/namelist.nps_AL.in", "namelist.nps.in", logger=logger)
elif STORM_BASIN == "E":
    produtil.fileop.deliver_file(parmdir+"/namelist.nps_EP.in", "namelist.nps.in", logger=logger)
elif STORM_BASIN == "C":
    produtil.fileop.deliver_file(parmdir+"/namelist.nps_EP.in", "namelist.nps.in", logger=logger)
else:
    logger.info("Not in 3 hurricane basin")

#produtil.fileop.deliver_file(parmdir+"/namelist.nps.in", "namelist.nps.in", logger=logger)
for line in fileinput.FileInput("namelist.nps.in",inplace=1):
    line = line.replace("_CEN_LAT_",CEN_LAT)
    line = line.replace("_CEN_LON_",CEN_LON)
    line = line.replace("_NEST1_LAT_",NEST1_LAT)
    line = line.replace("_NEST1_LON_",NEST1_LON)
    line = line.replace("_NEST2_LAT_",NEST2_LAT)
    line = line.replace("_NEST2_LON_",NEST2_LON)
    line = line.replace("NPS_GEOG",nps_geog)
    line = line.replace("MAX_DOMS","1")
    print(line, end=' ')

chdir ("../")

#####------------------------------ hycom init ----------------------------------------
for f in ['storm.conf']:
    if os.path.isfile(f):
        produtil.fileop.remove_file(f)
    produtil.fileop.deliver_file(parmdir + "/" + f + ".in", stormlabel + ".conf")
    for line in fileinput.FileInput(stormlabel + ".conf",inplace=1):
        line = line.replace("_YMDHM_", START_YMDHM)
        line = line.replace("_YMDH_", START_DATE)
        line = line.replace("_YMD_", START_YMD)
        line = line.replace("_YYYY_", START_YYYY)
        line = line.replace("_YY_", START_YY)
        line = line.replace("_CC_", START_CC)
        line = line.replace("_MM_", START_MM)
        line = line.replace("_DD_", START_DD)
        line = line.replace("_HH_", START_HH)
        line = line.replace("_STORM_ID_", STORM_ID)
        line = line.replace("_STORM_BASIN_", STORM_BASIN)
        line = line.replace("_STORM_NUM_", STORM_NUM)
        line = line.replace("_DOMLON_", CEN_LON)
        line = line.replace("_DOMLAT_", CEN_LAT)
        line = line.replace("_EXPT_", EXPT)
        line = line.replace("_SUBEXPT_", expname)
        print(line, end=' ')

#############------------------------forecast configure files---------------------
produtil.fileop.makedirs(RUNDIR+"/forecast")
chdir(RUNDIR+"/forecast")

for d in ['01', '02', '03']:
    if STORM_BASIN == "L":
        produtil.fileop.deliver_file(parmdir + "/configure_file_"+d+"_AL.in", "configure_file_"+d+".in")
    elif STORM_BASIN == "E":
        produtil.fileop.deliver_file(parmdir + "/configure_file_"+d+"_EP.in", "configure_file_"+d+".in")
    elif STORM_BASIN == "C":
        produtil.fileop.deliver_file(parmdir + "/configure_file_"+d+"_EP.in", "configure_file_"+d+".in")
    else:
        logger.info("Not in 3 hurricane basin")

    for line in fileinput.FileInput("configure_file_"+d+".in",inplace=1):
        line = line.replace("_CEN_LAT_",CEN_LAT)
        line = line.replace("_CEN_LON_",CEN_LON)
        line = line.replace("_START_YYYY_",START_YYYY)
        line = line.replace("_START_MM_",START_MM)
        line = line.replace("_START_DD_",START_DD)
        line = line.replace("_START_HH_",START_HH)
        line = line.replace("_FLENGTH_",FLENGTH)
        print(line, end=' ')


##############------------------------uncoupled forecast----------------------------------------
chdir("../")


logger.info("launched")