#!/usr/bin/env python

##@namespace pom.master
#@Biju Thomas, GSO, University of Rhode Island on June 12, 2014.
#@ Please report bugs/questions/comments to bijuthomas(at)uri(dot)edu.

import logging
from .init import Oceanini, fbtr, g3, na, pget, prun, phase, psend, rt, sst_pert
from .util import dateplushours, ysplitter, logi2int, jn2r
from .nml import nml, kppnml
import hwrf.launcher
from os.path import exists, getsize, split
from .domain import choose_domain
from .read_pom_conf import read_pom_conf
from .track import get_vitals, track_shorten
from .exceptions import *
import produtil.fileop
from produtil.fileop import rmall
from produtil.ecflow import set_ecflow_label

# To use default initialization in a given basin, enter "GDEM"  after the : below
# To use NCODA   initialization in a given basin, enter "NCODA" after the : below

##@var odata_d
# Mapping from basin letter to the list of available initializations for the basin.
odata_d = {"transatl":("fbtr","natr","rttr"), "eastpac":("g3ep","naep","rtep"),
           "westpac":("g3wp","nawp","rtwp"),"northind":("g3ni","nani","rtni"),
           "southind":("g3si","nasi","rtsi"),"swpac":("g3sw","nasw","rtsw"),
           "sepac":("g3se","nase","rtse"),"wn":("rtwn",),"cpac":("rtcp",) }

def run_init(STORMNAME,STORMID,STARTDATE,
             EXEChwrf,PARMhwrf,FIXhwrf,VITDIR,GFSDIR,
             LCDIR,CSTREAM,COMIN,init_data=None,logger=None, 
             fcstlen=None,outstep=None,**kwargs):
    """!Run the ocean initialization.

    This is a wrapper around the pom.init module that selects the
    right ocean initialization for the chosen basin and delivers the
    output to the specified location.

    @param STORMNAME Upper-case storm name for filenames (KATRINA)
    @param STORMID Three character storm number and basin, upper-case (12L)
    @param STARTDATE Simulation analysis time as a string, YYYYMMDDHH (2005082918)
    @param EXEChwrf Directory with HWRF executables.
    @param PARMhwrf Directory with HWRF parameter files.
    @param FIXhwrf Directory with HWRF fixed files.
    @param GFSDIR Directory with input files from GFS.
    @param LCDIR Directory with loop current files
    @param CSTREAM Directory to place files for the POM forecast to read.
    @param COMIN HWRF final output directory.

    @param init_data The string name of the initialization method to
    use.  If this is unspecified, a suitable default is chosen.  Allowed values: GDEM, NCODA.

    @param logger a logging.Logger for log messages

    @param fcstlen forecast length in hours
    @param outstep output frequency in seconds (an integer)

    @param kwargs Additional keyword arguments are passed to Oceanini subclass constructors."""
    assert(GFSDIR.find('pom/output')<0)
    if logger is None: logger=logging.getLogger('pom')
    if fcstlen is None: fcstlen=126 # forecast length in hours
    if outstep is None: outstep=86400
    if outstep<540: outstep=540
    printfrq = outstep/86400.0 

    (y4,mm,mm1,dd,hh)=ysplitter(STARTDATE)
    BASIN=STORMID[2].upper()
    conf = kwargs.get('conf', None)
    if conf is None:
       WORKDIR = '/'.join(list([x for x in CSTREAM.split("/") if x != '']))
       WORKDIR = split(split(WORKDIR)[0])[0]
       WORKDIR = split(split(WORKDIR)[0])[0]
       CONFhwrf = "".join(["/",WORKDIR, "/com/",STARTDATE,"/",STORMID,"/storm1.conf"])
       if exists(CONFhwrf) and getsize(CONFhwrf) > 0:
          conf = hwrf.launcher.load(CONFhwrf)
       else:
          logger.error("%s does not exists" %(CONFhwrf))
          msg=("%s Missing or zero size" %(CONFhwrf))
          raise POMConfigError(msg)
    DOMAIN = choose_domain(BASIN, conf, logger=None)
    if DOMAIN is None:
        logger.critical("Domain selection Failed")
        msg=("Domain selection Failed")
        raise POMConfigError(msg)

    #conf_file = "".join([PARMhwrf,"/pom/", "pom.conf_tmpl"])
    #pom_conf = read_pom_conf(conf_file, logger=None)
    # Passing options in hwrf.conf [pom] section in
    init_data =  conf.getstr('pom','ini_data','gdem').upper()
    INDATA = init_data
    GEOVEL = conf.getint('pom','geovflag',1)
    SSTASIM = conf.getint('pom','assi_sst',1)
    PHY1D = conf.getint('pom','oned_pom',0)
    kppflag = conf.getint('pom','kppflag',0)
    kpp_ric = conf.getfloat('pom','kpp_ric',0.36)
    kpp_lt_log = conf.getstr('pom','kpp_lt_log','.false.')
    multistorm = conf.getstr('config','run_multistorm','no')
    if multistorm == "yes" and init_data != "RTOF":
        init_data = "RTOF"

    if  DOMAIN == "transatl":
        nproc_i = 3
        nproc_j = 3
        CENTERID = "NHC"
        if init_data == 'GDEM' or init_data == 'GDEM3':
            NL = 33
            ODATA = odata_d[DOMAIN][0]
            prep = fbtr("OCEAN",DOMAIN,ODATA,str(SSTASIM),INDATA,str(PHY1D),multistorm,STORMNAME,STORMID,
                     STARTDATE,EXEChwrf,PARMhwrf,FIXhwrf,LCDIR,GFSDIR,CSTREAM,COMIN, **kwargs)
        elif init_data == 'NCODA':
            NL = 34
            ODATA = odata_d[DOMAIN][1]
            prep = na("OCEAN",DOMAIN,ODATA,str(SSTASIM),INDATA,str(PHY1D),multistorm,STORMNAME,STORMID,
                       STARTDATE,EXEChwrf,PARMhwrf,FIXhwrf,LCDIR,GFSDIR,CSTREAM,COMIN,
                       **kwargs)
        elif init_data == 'RTOF':
            NL = 40
            ODATA = odata_d[DOMAIN][2]
            prep = rt("OCEAN",DOMAIN,ODATA,str(SSTASIM),INDATA,str(PHY1D),multistorm,STORMNAME,STORMID,
                       STARTDATE,EXEChwrf,PARMhwrf,FIXhwrf,LCDIR,GFSDIR,CSTREAM,COMIN,
                       **kwargs)
        else:
            msg=("%s:OCEAN INIT DATA IS NOT Supported" %(init_data))
            logger.critical(msg)
            raise POMConfigError(msg)
        (DATA,DATA1) = prep.setpath()
        ymdh = STARTDATE
        infiles = pget()
        infiles.getinp(prep,DATA)
        vitfile = jn2r(VITDIR,"syndat_tcvitals." + y4)
        trackfile = jn2r(DATA1,"track.full")
        trackshort = jn2r(DATA1,"track_1") 
        rmall(trackfile,trackshort,logger=logger)
        nvit = get_vitals(vitfile, CENTERID, STORMID, y4,trackfile,logger=logger)
        nvit = track_shorten(trackfile,trackshort,STARTDATE)
        if nvit <= 0:
            msg=("Did not find track info for %s. Write an empty message." %(STARTDATE))
            logger.warning(msg)
            with open(trackshort,'wt') as fo:
                fo.write("%s\n"%("000  000 00000     000000 0000 000 0000 000 000 0000 0000"))
        runexes = prun()
        if not runexes.setrun(prep,DATA):
            msg='setrun failed for L domain prep'
            logger.warning(msg)
            raise POMInitFailed(msg)
    elif DOMAIN == "eastpac" :
        nproc_i = 3
        nproc_j = 3
        CENTERID = "NHC"
        if init_data == 'GDEM' or init_data == 'GDEM3':
            NL = 73
            ODATA = odata_d[DOMAIN][0]
            prep = g3("OCEAN",DOMAIN,ODATA,str(SSTASIM),INDATA,str(PHY1D),multistorm,STORMNAME,STORMID,
                       STARTDATE,EXEChwrf,PARMhwrf,FIXhwrf,LCDIR,GFSDIR,CSTREAM,COMIN,
                       **kwargs)
        elif init_data == 'NCODA':
            NL = 34
            ODATA = odata_d[DOMAIN][1]
            prep = na("OCEAN",DOMAIN,ODATA,str(SSTASIM),INDATA,str(PHY1D),multistorm,STORMNAME,STORMID,
                   STARTDATE,EXEChwrf,PARMhwrf,FIXhwrf,LCDIR,GFSDIR,CSTREAM,COMIN,
                   **kwargs)
        elif init_data == 'RTOF':
            NL = 40
            ODATA = odata_d[DOMAIN][2]
            prep = rt("OCEAN",DOMAIN,ODATA,str(SSTASIM),INDATA,str(PHY1D),multistorm,STORMNAME,STORMID,
                       STARTDATE,EXEChwrf,PARMhwrf,FIXhwrf,LCDIR,GFSDIR,CSTREAM,COMIN,
                       **kwargs)
        else:
            msg=("%s:OCEAN INIT DATA IS NOT Supported" %(init_data))
            logger.critical(msg)
            raise POMConfigError(msg)
        (DATA,DATA1) = prep.setpath()
        ymdh = STARTDATE
        infiles = pget()
        infiles.getinp(prep,DATA)
        vitfile = jn2r(VITDIR,"syndat_tcvitals." + y4)
        trackfile = jn2r(DATA1,"track.full")
        trackshort = jn2r(DATA1,"track_1") 
        rmall(trackfile,trackshort,logger=logger)
        nvit = get_vitals(vitfile, CENTERID, STORMID, y4,trackfile,logger=logger)
        nvit = track_shorten(trackfile,trackshort,STARTDATE)
        if nvit <= 0:
            msg=("Did not find track info for %s. Write an empty message." %(STARTDATE))
            logger.warning(msg)
            with open(trackshort,'wt') as fo:
                fo.write("%s\n"%("000  000 00000     000000 0000 000 0000 000 000 0000 0000"))
        runexes = prun()  
        if not runexes.setrun(prep,DATA):
            msg='setrun failed for E domain prep'
            logger.warning(msg)
            raise POMInitFailed(msg)
#
    elif DOMAIN == "cpac" :
        nproc_i = 3
        nproc_j = 3
        CENTERID = ["NHC", "JTWC"]
        if init_data == 'RTOF':
            NL = 40
            ODATA = odata_d[DOMAIN][0]
            prep = rt("OCEAN",DOMAIN,ODATA,str(SSTASIM),INDATA,str(PHY1D),multistorm,STORMNAME,STORMID,
                       STARTDATE,EXEChwrf,PARMhwrf,FIXhwrf,LCDIR,GFSDIR,CSTREAM,COMIN,
                       **kwargs)
        else:
            msg=("%s:OCEAN INIT DATA IS NOT Supported" %(init_data))
            logger.critical(msg)
            raise POMConfigError(msg)
        (DATA,DATA1) = prep.setpath()
        ymdh = STARTDATE
        infiles = pget()
        infiles.getinp(prep,DATA)
        vitfile = jn2r(VITDIR,"syndat_tcvitals." + y4)
        trackfile = jn2r(DATA1,"track.full")
        trackshort = jn2r(DATA1,"track_1")
        rmall(trackfile,trackshort,logger=logger)
        nvit = get_vitals(vitfile, CENTERID, STORMID, y4,trackfile,logger=logger)
        nvit = track_shorten(trackfile,trackshort,STARTDATE)
        if nvit <= 0:
            msg=("Did not find track info for %s. Write an empty message." %(STARTDATE))
            logger.warning(msg)
            with open(trackshort,'wt') as fo:
                fo.write("%s\n"%("000  000 00000     000000 0000 000 0000 000 000 0000 0000"))
        runexes = prun()
        if not runexes.setrun(prep,DATA):
            msg='setrun failed for E domain prep'
            logger.warning(msg)
            raise POMInitFailed(msg)
#    
    elif DOMAIN == "wn" :
        nproc_i = 5
        nproc_j = 4
        CENTERID = "NHC"
        if init_data == 'RTOF':
            NL = 40
            ODATA = odata_d[DOMAIN][0]
            prep = rt("OCEAN",DOMAIN,ODATA,str(SSTASIM),INDATA,str(PHY1D),multistorm,STORMNAME,STORMID,
                       STARTDATE,EXEChwrf,PARMhwrf,FIXhwrf,LCDIR,GFSDIR,CSTREAM,COMIN,
                       **kwargs)
        else:
            msg=("%s:OCEAN INIT DATA IS NOT Supported" %(init_data))
            logger.critical(msg)
            raise POMConfigError(msg)
        (DATA,DATA1) = prep.setpath()
        ymdh = STARTDATE
        infiles = pget()
        infiles.getinp(prep,DATA)
        vitfile = jn2r(VITDIR,"syndat_tcvitals." + y4)
        trackfile = jn2r(DATA1,"track.full")
        trackshort = jn2r(DATA1,"track_1")
        rmall(trackfile,trackshort,logger=logger)
        nvit = get_vitals(vitfile, CENTERID, STORMID, y4,trackfile,logger=logger)
        nvit = track_shorten(trackfile,trackshort,STARTDATE)
        if nvit <= 0:
            msg=("Did not find track info for %s. Write an empty message." %(STARTDATE))
            logger.warning(msg)
            with open(trackshort,'wt') as fo:
                fo.write("%s\n"%("000  000 00000     000000 0000 000 0000 000 000 0000 0000"))
        runexes = prun()
        if not runexes.setrun(prep,DATA):
            msg='setrun failed for basin scale (EPAC & ATL)  domain prep'
            logger.warning(msg)
            raise POMInitFailed(msg)
    elif DOMAIN in ["westpac", "northind", "southind", "swpac", "sepac"]:
        nproc_i = 3
        nproc_j = 3 
        CENTERID = "JTWC"
        if init_data == 'GDEM'or init_data == 'GDEM3':
            NL = 73
            ODATA = odata_d[DOMAIN][0]
            prep = g3("OCEAN",DOMAIN,ODATA,str(SSTASIM),INDATA,str(PHY1D),multistorm,STORMNAME,STORMID,
                       STARTDATE,EXEChwrf,PARMhwrf,FIXhwrf,LCDIR,GFSDIR,CSTREAM,COMIN,
                       **kwargs)
        elif init_data == 'NCODA':
            NL = 34
            ODATA = odata_d[DOMAIN][1]
            prep = na("OCEAN",DOMAIN,ODATA,str(SSTASIM),INDATA,str(PHY1D),multistorm,STORMNAME,STORMID,
                   STARTDATE,EXEChwrf,PARMhwrf,FIXhwrf,LCDIR,GFSDIR,CSTREAM,COMIN,
                   **kwargs)
        elif init_data == 'RTOF':
            NL = 40
            ODATA = odata_d[DOMAIN][2]
            prep = rt("OCEAN",DOMAIN,ODATA,str(SSTASIM),INDATA,str(PHY1D),multistorm,STORMNAME,STORMID,
                       STARTDATE,EXEChwrf,PARMhwrf,FIXhwrf,LCDIR,GFSDIR,CSTREAM,COMIN,
                       **kwargs)
        else:
            msg=("%s:OCEAN INIT DATA IS NOT Supported" %(init_data))
            logger.critical(msg)
            raise POMConfigError(msg)
        (DATA,DATA1) = prep.setpath()
        ymdh = STARTDATE
        infiles = pget()
        infiles.getinp(prep,DATA)
        vitfile = jn2r(VITDIR,"syndat_tcvitals." + y4)
        trackfile = jn2r(DATA1,"track.full")
        trackshort = jn2r(DATA1,"track_1") 
        rmall(trackfile,trackshort,logger=logger)
        nvit = get_vitals(vitfile, CENTERID, STORMID, y4,trackfile,logger=logger)
        nvit = track_shorten(trackfile,trackshort,STARTDATE)
        if nvit <= 0:
            msg=("Did not find track info for %s. Write an empty message." %(STARTDATE))
            logger.warning(msg)
            with open(trackshort,'wt') as fo:
                fo.write("%s\n"%("000  000 00000     000000 0000 000 0000 000 000 0000 0000"))
        runexes = prun()
        if not runexes.setrun(prep,DATA):
            msg='setrun failed for W/A/S/P/X/O/B/U/T/Q domain prep'
            logger.warning(msg)
            raise POMUnsupportedBasin(msg)
    else:
        msg=("%s : Domain is NOT Supported" %(DOMAIN))
        logger.critical(msg)
        raise POMUnsupportedBasin(msg)
##########################################################################################
# SST perturbation start
    ens=conf.getint('config','ENS',99)
    ensize=conf.getint('ensemble','ensize',20)
    if ens>0 and ens<99 and ens<=ensize:
        sstpert=conf.getstr('ensemble','sst_pert')
        if sstpert == 'yes':
            sstpert = sst_pert("OCEAN",DOMAIN,ODATA,str(SSTASIM),INDATA,
	    str(PHY1D),STORMNAME,STORMID,STARTDATE,EXEChwrf,PARMhwrf,FIXhwrf,
	    LCDIR,GFSDIR,CSTREAM,COMIN,str(ens),**kwargs)
            infiles = pget()
            infiles.getinp(sstpert,DATA)
            runexes = prun()
            if not runexes.setrun(sstpert,DATA):
                msg='setrun failed for sst_pert'
                logger.warning(msg)
                raise POMInitFailed(msg)
# SST perturbation end 
###########################################################################################
    diag = phase(nproc_i,nproc_j,"OCEAN",DOMAIN,ODATA,str(SSTASIM),INDATA,str(PHY1D),
                 multistorm,STORMNAME,STORMID,STARTDATE,EXEChwrf,PARMhwrf,FIXhwrf,LCDIR,
                 GFSDIR,CSTREAM,COMIN,**kwargs)
    (DATA,DATA1) = diag.setpath("PHASE1")
    namelst=nml()
    #Comment out and try to use the outstep to control the output time step.
    #printfrq = pom_conf['pom.prtd1val']

    namelst("'MPIPOM-TC:"+STORMNAME+STORMID+"'", STORMNAME,
            6, 45, "'"+y4+"-"+mm+"-"+dd+" "+hh+":00:00 +00:00'",0,
            "'restart.phase0.nc'",2.0,2.0,printfrq,3,
            0, 22.4, 0.,GEOVEL,NL, PHY1D, 0, 1, kppflag)

    namelst.make(DATA1)
    infiles.getinp(diag,DATA1)
    if not  runexes.setrun(diag,DATA1):
        msg='setrun failed for diag'
        logger.warning(msg)
        raise POMInitFailed(msg)
    outfiles = psend()
    set_ecflow_label('wake',"Starting...",logger)
    prog = phase(nproc_i,nproc_j,"OCEAN",DOMAIN,ODATA,str(SSTASIM),INDATA,str(PHY1D),
                 multistorm,STORMNAME,STORMID,STARTDATE,EXEChwrf,PARMhwrf,FIXhwrf,LCDIR,
                 GFSDIR,CSTREAM,COMIN,**kwargs)
    ymdh = dateplushours(ymdh, -72)
    (DATA,DATA2) = prog.setpath("PHASE2")
    (y4,mm,mm1,dd,hh)=ysplitter(ymdh)
    namelst("'MPIPOM-TC:"+STORMNAME+STORMID+"'", STORMNAME,
            6, 45, "'"+y4+"-"+mm+"-"+dd+" "+hh+":00:00 +00:00'",1,
            "'restart.phase1.nc'",3.0,3.0,printfrq,1,
            0, 22.4, 9999.,GEOVEL,NL, PHY1D, 0, 1, kppflag)
 
    namelst.make(DATA2) 
    if outfiles.sendout(diag,DATA1,"restart.0001.nc",DATA2,"restart.phase1.nc"):
        outfiles.sendout(diag,DATA1,"pom_proc_count",DATA2,"pom_proc_count")
        vitfile = jn2r(VITDIR,"syndat_tcvitals." + y4)
        trackfile = jn2r(DATA2,"track.full")
        trackshort = jn2r(DATA2,"track_1") 
        rmall(trackfile,trackshort,logger=logger)
        nvit = get_vitals(vitfile, CENTERID, STORMID, y4,trackfile,logger=logger)
        nvit = track_shorten(trackfile,trackshort,STARTDATE)
        infiles.getinp(prog,DATA2)
        if nvit >= 2 and init_data != 'RTOF':
            if runexes.setrun(prog,DATA2):
                set_ecflow_label('wake',"Cold wake added successfully.",logger)
                outfiles.sendout(prog,DATA2,"restart.0001.nc",DATA,"restart.phase2.nc")
                outfiles.sendout(prog,DATA2,"pom_proc_count",DATA,"pom_proc_count")
            else:
                set_ecflow_label('wake',"Cold wake init failed, will run without.",logger)
                outfiles.sendout(prog,DATA2,"restart.phase1.nc",DATA,"restart.phase2.nc")
                outfiles.sendout(prog,DATA2,"pom_proc_count",DATA,"pom_proc_count")
        else:
            if init_data=='RTOF':
                set_ecflow_label('wake',"Using RTOFS cold wake.",logger)
            else:
                set_ecflow_label('wake',"No cold wake (too few vitals times).",logger)
            outfiles.sendout(prog,DATA2,"restart.phase1.nc",DATA,"restart.phase2.nc")
            outfiles.sendout(prog,DATA2,"pom_proc_count",DATA,"pom_proc_count")
    else:
        msg='sendout failed'
        logger.warning(msg)
        raise POMInitFailed(msg)

    # Make a namelist for the forecast model as pom.nml in the top-level directory:
    namelst=nml()
    (y4,mm,mm1,dd,hh)=ysplitter(STARTDATE)
    namelst("'MPIPOM-TC:"+STORMNAME+STORMID+"'", STORMNAME,
            6, 30, "'"+y4+"-"+mm+"-"+dd+" "+hh+":00:00 +00:00'",1,
            "'restart.phase2.nc'",9999.,fcstlen/24.,printfrq,1,
            0, 22.4, 9999.,GEOVEL,NL, PHY1D, 0, 1, kppflag)
    namelst.make('.')
    # Make a namelist for the forecast model as kpp.nml in the top-level directory:
    kpplst=kppnml()
    kpplst(kpp_lt_log,kpp_ric)
    kpplst.make('.')
    return True