#! /usr/bin/env python
import os, sys, stat, shutil, fileinput, time, re

import produtil.setup,  produtil.run, produtil.fileop, produtil.cd, produtil.batchsystem, produtil.numerics, produtil.config, produtil.dbnalert
from produtil.fileop import remove_file, make_symlink, isnonempty, chdir, fortlink
from produtil.run import run, exe, runstr
from produtil.numerics import to_datetime_rel
import nmmbUsh
import logging
from nmmbUsh import logger
import hmon_vitals

produtil.setup.setup()

logger.info("post started")

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

produtil.fileop.makedirs(workdir + "/post1")
produtil.fileop.makedirs(workdir + "/post")

chdir(workdir + "/post1")
logger.info(os.getcwd())

nmmbUsh.rm_glob('NMMPRS_d*')
nmmbUsh.rm_glob('*.ix')

conf=produtil.config.from_file(parmdir + "/hmon.conf")
conf.cycle=os.environ.get('CYCLE')
stormlabel=conf.get('config','stormlabel')
label=stormlabel[5]
fixdir=conf.get('dir','FIXhmon')
exedir=conf.get('dir','EXEChmon')
scriptdir=conf.get('dir','EXhmon')
comdir=conf.get('dir','COM')
syndat=conf.get('dir','syndat')

with open(workdir + "/get_storm_info/storm_info", 'r') as f:
    for line in f:
        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')
        if "STORM_NAME=" in line:
           STORMNAME = line.replace("STORM_NAME=","")
           STORM_NAME = STORMNAME.strip('\n')
        if "STORM_LAT=" in line:
           STORMLAT = line.replace("STORM_LAT=","")
           STORM_LAT = STORMLAT.strip('\n')
        if "STORM_LON=" in line:
           STORMLON = line.replace("STORM_LON=","")
           STORM_LON = STORMLON.strip('\n')
        if "CEN_LAT=" in line:
           CENLAT = line.replace("CEN_LAT=","")
           CEN_LAT = CENLAT.strip('\n')
        if "CEN_LON=" in line:
           CENLON = line.replace("CEN_LON=","")
           CEN_LON = CENLON.strip('\n')

cgb=os.environ.get("COPYGB")
cnv=os.environ.get("CNVGRIB")
wgrib2=os.environ.get("WGRIB2")
STORM_ID=stormid
STORM_NUM=STORM_ID[0:2]
STORM_BASIN=STORM_ID[2:3]
START_DATE=cycle
FLENGTH=126
FINC=0o6
FLENGTH = int(FLENGTH)

### 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, 
#    COMINmsg='/gpfs/hps/nco/ops/com/hur/prod/inphwrf',
    COMINmsg=conf.get('dir','COMmsg'),
#    storm_num='1' # request storm in slot 1
    storm_num=label # request storm in slot 1
    )

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

produtil.fileop.deliver_file(fixdir + "/ETAMPNEW_DATA.expanded_rain_DBL","eta_micro_lookup.dat",logger=logger)
produtil.fileop.deliver_file(parmdir + "/nmmb_cntrl.parm","nmmb_cntrl.parm",logger=logger)

produtil.fileop.remove_file("NMMPRS_d??.ll",logger=logger)
produtil.fileop.remove_file("fcst_minutes",logger=logger)

IFHR=0
FHR=0

FHR2="%02d"%(FHR)
FHR3="%03d"%(FHR)


while FHR <= FLENGTH:

    DATE = produtil.numerics.to_datetime_rel(int(FHR3)*3600, int(START_DATE))
    YYYY=DATE.strftime("%Y")
    CC=DATE.strftime("%C")
    YY=DATE.strftime("%y")
    MM=DATE.strftime("%m")
    DD=DATE.strftime("%d")
    HH=DATE.strftime("%H")

    logger.info("DATE=",DATE)

    for domain in ['01', '02', '03']:
        
        if domain == '01':
            filesize=1696470000
        elif domain == '02':
            filesize=387340000
        else:
            filesize=1096350000

## find if files are ready or not
        fileName = "nmmb_hst_"+domain+"_nio_0"+FHR3+"h_00m_00.00s"
        while not(os.path.isfile("../forecast/" + fileName)) or os.path.getsize("../forecast/" + fileName)<=filesize :
            logger.info(fileName +" not ready, sleep 15")
            time.sleep(15)

        if (os.path.isfile("../forecast/" + fileName)) or os.path.getsize("../forecast/" + fileName)>filesize :
            nmmbUsh.make_symlinks("../forecast/" + fileName,".")

        logger.info(fileName +" ready, do post")
        time.sleep(5)

        if os.path.isfile("fort.14") and os.path.getsize("fort.14")>0: 
            produtil.fileop.remove_file("fort.14",logger=logger)
        
        if os.path.isfile("fort.110") and os.path.getsize("fort.110")>0:
            produtil.fileop.remove_file("fort.110",logger=logger)

        make_symlink("nmmb_cntrl.parm","fort.14", True)
        make_symlink("griddef.out", "fort.110", True)

        text = "nmmb_hst_" + domain + "_nio_0" + FHR3 + "h_00m_00.00s" + os.linesep + "binarynemsio" + os.linesep +  YYYY + "-" + MM + "-" + DD + "_" + HH + ":00:00" + os.linesep +  "NMM"       
        with open("itag",'w') as f: f.write(text)

        logger.info("mpirun -np " + "hmon_post > out_nemsio_d" + domain + "_" +FHR3 + "2>&1")
        produtil.run.checkrun(produtil.run.mpirun(produtil.run.mpi(exedir + "/hmon_post"),allranks=True).out("out_nemsio_d" + domain + "_" + FHR3).err("out_nemsio_d" + domain + "_" + FHR3),logger=logger)

        hhh=FHR3
        if int(FHR3)>99:
            grbF = "NMMPRS.GrbF"+FHR3
            if os.path.isfile(grbF):
                produtil.fileop.deliver_file(grbF, "NMMPRS_d"+domain + ".GrbF"+FHR3,logger=logger)
        else:
            grbF = "NMMPRS.GrbF"+FHR2
            if os.path.isfile(grbF):
                produtil.fileop.deliver_file(grbF, "NMMPRS_d" + domain + ".GrbF" + FHR3,logger=logger)
#            hhh=FHR2
        
### creat domains for grib files
        gribfile = "NMMPRS_d" + domain + ".GrbF" + FHR3

        wgrib_V_out=produtil.run.runstr(exe(os.environ["WGRIB"])['-V',gribfile])
        wgrib_V=wgrib_V_out.splitlines()

        lat0=None
        lon0=None
        dlat=None
        varname='None' 
        for line in wgrib_V:
            m=re.search('^\s*(\S+=\S+)',line)
            if m: varname=m.group(1)
            if varname == 'NLAT=Latitude':
                if line.find('lat0')>=0:
                    lat0=line.split()[5]
                    lon0=line.split()[7]
                if line.find('dLat')>=0:
                    dlat=line.split()[1]

        logger.info("lat0="+str(lat0))
        logger.info("lon0="+str(lon0))

##out grid for three domains
## d01 dx=0.2,  size=110*90, grid=551*451
## d02 dx=0.06, size=12*14,  grid=234*201
## d03 dx=0.02, size=7.5*9,  grid=450*375
##
        if domain == '01':
            ddd="d1.0p20"
            nx=551
            ny=451
            d1000=0.2*1000
            xx1=(int(float(CEN_LON)*1000) - int(d1000*nx/2))
            yy1=(int(float(CEN_LAT)*1000) - int(d1000*ny/2))
        elif domain == '02':
            ddd="d2.0p06"	
            nx=234
            ny=201
            d1000=0.06*1000
            xx1=int(float(lon0)*1000) 
            yy1=int(float(lat0)*1000)
        elif domain == '03':
            ddd="d3.0p02"
            nx=450
            ny=375
            d1000=0.02*1000
            xx1=int(float(lon0)*1000)
            yy1=int(float(lat0)*1000)
            track_lat = float(lat0)
            track_lon = float(lon0)

        if xx1 < 0:
                xx1=(xx1 + 360000)  #### use 0-360, positive longitude
        xx2=(xx1 + nx*d1000 - d1000)

        yy2=(yy1 + ny*d1000 - d1000)

        
 
        outgrid="255 0   " + str(nx) + " " + str(ny) + "  " +  str(yy1) + "  " +  str(xx1) + "  128  " + str(yy2) + "  " + str(xx2) + "  " + str(d1000) + " " + str(d1000) + " 64"
        logger.info("outgrid = " + outgrid) 

        outgribfile=storm_name+storm_id + "." + START_DATE + ".hmonprs." + ddd + ".f" + str(hhh) + ".grb"
        logger.info(cgb + "-g " + outgrid + " -x " + gribfile + " " + outgribfile)
        cmd = exe(cgb)["-g", outgrid, "-x", gribfile, outgribfile]
        produtil.run.checkrun(cmd,logger=logger)

### convert grib1 to grib2
        grib2file=storm_name+storm_id + "." + START_DATE + ".hmonprs." + ddd + ".f" + str(hhh) +".grb2"
        grib2idx=storm_name+storm_id + "." + START_DATE + ".hmonprs." + ddd + ".f" + str(hhh) +".grb2.idx"
        logger.info(cnv + " -g12 -p0 " + outgribfile + " " + grib2file)
        cmd = exe(cnv)["-g12","-p0",outgribfile,grib2file].env(OMP_NUM_THREADS=1,MKL_NUM_THREADS=1)         
        produtil.run.checkrun(cmd,logger=logger)
### create index file for grib2 file
        produtil.run.checkrun(exe(wgrib2)['-s',grib2file] > grib2idx,logger=logger)


### extract grib files for tracker 
        fileName = "NMMPRS_d" + domain+".GrbF" + FHR3
        while not(os.path.isfile(fileName)) or os.path.getsize(fileName)<=0 :
            logger.info(fileName +" not ready, sleep 60")
            time.sleep(60)
  
        logger.info(fileName +" ready, do index file for tracker")
        time.sleep(15)
   
        wgrib_s=runstr(exe(os.environ["WGRIB"])['-s', "NMMPRS_d" + domain + ".GrbF" + FHR3]) 
        keep=''
        for line in wgrib_s.splitlines():
            if re.search('(:PRMSL:MSL:|:UGRD:500 mb:|:VGRD:500 mb:|:HGT:700 mb:|:UGRD:700 mb:|:VGRD:700 mb:|:ABSV:700 mb:|:HGT:850 mb:|:UGRD:850 mb:|:VGRD:850 mb:|:ABSV:850 mb:|:UGRD:10 m above gnd:|:VGRD:10 m above gnd:)',line):
                keep+=line+'\n'
        
        produtil.run.checkrun(exe(os.environ["WGRIB"])['-i',
           "NMMPRS_d" + domain + ".GrbF" + FHR3,
           '-grib', '-o', "NMMPRS_d" + domain + ".GrbF" + FHR3 + "_filt"]
           << keep)

## copy grib file to com directory
        atmos_d = comdir
#    produtil.fileop.deliver_file(outgribfile,atmos_d + "/" + outgribfile,logger=logger)
        produtil.fileop.deliver_file(grib2file,atmos_d + "/" + grib2file,logger=logger)
        produtil.fileop.deliver_file(grib2idx,atmos_d + "/" + grib2idx,logger=logger)

## link the file to post directory for tracker
        make_symlink(workdir + "/post1/" + grib2file,workdir + "/post/" + grib2file, True)

# The line below makes a DBNAlert object that will alert the DBNet.
# We will reuse this for all later alerts.
        alerter=produtil.dbnalert.DBNAlert(
            ['MODEL','{type}','{job}','{location}']) # arguments to dbn_alert

# These "alerter()" lines send a DBN alert using the alerter object we just created.
        alerter(location=atmos_d+"/"+grib2file,type='HMON_GB2')
        alerter(location=atmos_d+"/"+grib2idx,type='HMON_GB2_WIDX')

# The type='NULL' is a placeholder during parallel tests and code
# hand-off.  The NCO Dataflow Team will replace "NULL" with the
# correct value after code handoff.


### initial track domain center is at d3 center, size is 20x20
  
    track_lat2=int(1000*(track_lat+10))
    if track_lat2>=80000:
        track_lat2=80000

    track_lat1=(track_lat2 - 20000)
    track_lon1=int(track_lon*1000)-10*1000
    track_lon2=int(track_lon*1000)+10*1000

    trackgrid="255 0 1001 1001 " +  str(track_lat1) + " " + str(track_lon1) + " 128 " + str(track_lat2) + " " +  str(track_lon2) + " 020 020 64" 
    logger.info("trackgrid=" + trackgrid)

### combine d1, d2, d3 for tracker
    produtil.run.checkrun(exe(cgb)["-g" + trackgrid, '-xM',
                                   "NMMPRS_d01.GrbF" + FHR3 + "_filt",
                                   "NMMPRS_d02.GrbF" + FHR3 + "_filt",
                                   "NMMPRS_d12.GrbF" + FHR3 + "_ll"],
                          logger=logger)

    produtil.run.checkrun(exe(cgb)["-g" + trackgrid, '-xM',
                                   "NMMPRS_d12.GrbF" + FHR3 + "_ll",
                                   "NMMPRS_d03.GrbF" + FHR3 + "_filt",
                                   "NMMPRS_d123.GrbF" + FHR3 + "_ll"],
                          logger=logger)

### creat index files for tracker
    fmin = (FHR*60)
    minstr = "%5.5d" %(fmin)
    trakfile = "hmon.trak." + storm_name + storm_id + "." + START_DATE + ".f" + minstr 
    trakfile_ix = trakfile + ".ix"
    fileName = "NMMPRS_d123.GrbF" + FHR3 + "_ll"
    logger.info("fileName=" + fileName)
    make_symlink(fileName, trakfile,True)

    gbi=os.environ.get("GRBINDEX")
    cmd = exe(gbi)[trakfile,trakfile+".ix"]
    produtil.run.run(cmd, logger=logger)

    make_symlink(workdir + "/post1/"+ trakfile_ix, workdir + "/post/" + trakfile_ix, True)
    make_symlink(workdir + "/post1/"+ fileName, workdir + "/post/" + trakfile, True)

## finish writing grib files
    IFHR=IFHR + 1
    FHR=FHR + FINC
    FHR2="%02d" % (FHR)
    FHR3="%03d" % (FHR)

logger.info("run_post finished")