#!/usr/bin/env python3

"""
Namelist creator for E3SM's MOSART component
"""

import os, sys

_CIMEROOT = os.path.join(os.path.dirname(os.path.abspath(__file__)), "..","..","..","cime")
sys.path.append(os.path.join(_CIMEROOT, "scripts", "Tools"))

from standard_script_setup import *
from CIME.case import Case
from CIME.utils import expect, run_cmd_no_fail, safe_copy
from CIME.buildnml import create_namelist_infile, parse_input

logger = logging.getLogger(__name__)

###############################################################################
def buildnml(case, caseroot, compname):
###############################################################################
    expect(compname == "mosart", compname)

    os.chdir(caseroot)

    casebuild            = case.get_value("CASEBUILD")
    caseroot             = case.get_value("CASEROOT")
    srcroot              = case.get_value("SRCROOT")
    din_loc_root         = case.get_value("DIN_LOC_ROOT")
    ninst_rof            = case.get_value("NINST_ROF")
    rof_grid             = case.get_value("ROF_GRID")
    mosart_bldnml_opts   = case.get_value("MOSART_BLDNML_OPTS")
    mosart_namelist_opts = case.get_value("MOSART_NAMELIST_OPTS")
    rundir               = case.get_value("RUNDIR")
    run_type             = case.get_value("RUN_TYPE")
    run_refcase          = case.get_value("RUN_REFCASE")
    run_refdate          = case.get_value("RUN_REFDATE")
    run_reftod           = case.get_value("RUN_REFTOD")
    scriptsroot          = case.get_value("SCRIPTSROOT")

    rofconf_dir = os.path.join(casebuild, "mosartconf")

    if not os.path.isdir(rofconf_dir): os.mkdir(rofconf_dir)

    #--------------------------------------------------------------------
    # Verify rof grid is supported
    #--------------------------------------------------------------------

    rof_grid_supported = ("null", "r2", "r05", "r0125", "r01", "NLDAS", "ELM_USRDAT")
    expect(rof_grid in rof_grid_supported, "ROF_GRID '{}' is not supported in mosart. Choose from: null, r2, r05, r0125, r01, NLDAS, ELM_USRDAT".format(rof_grid))

    #--------------------------------------------------------------------
    # Invoke mosart build-namelist - output will go in $CASEBUILD/mosartconf
    #--------------------------------------------------------------------

    inst_string = ""
    for inst_counter in range(1, ninst_rof + 1):

        # -----------------------------------------------------
        # determine instance string
        # -----------------------------------------------------

        inst_string = ""
        if ninst_rof > 1:
            inst_string = "_{0:04d}".format(inst_counter)

            # If multi-instance case does not have restart file, use single-case restart
            # for each instance
            if not os.path.exists(os.path.join(rundir, "rpointer.rof{}".format(inst_string))) and \
                   os.path.exists(os.path.join(rundir, "rpointer.rof")):
                safe_copy(os.path.join(rundir, "rpointer.rof"),
                          os.path.join(rundir, "rpointer.rof{}".format(inst_string)))

        # -----------------------------------------------------
        # create mosartconf/cesm_namelist
        # -----------------------------------------------------

        if os.path.exists(os.path.join(casebuild, "mosart.input_data_list")):
            os.remove(os.path.join(casebuild, "mosart.input_data_list"))

        # The following is for backwards compatibility when runoff restart data was on clm restart files
        infile_text = ""
        if rof_grid != "null":

            fncheck = ""
            if run_type == 'hybrid' or run_type == "branch":

    	        # search for clm or mosart files with instance or not
                fncheck = "{}/{}.mosart{}.r.{}-{}.nc".format(rundir, run_refcase, inst_string, run_refdate, run_reftod)
                if not os.path.exists(fncheck):
                    fncheck = "{}/{}.mosart.r.{}-{}.nc".format(rundir, run_refcase, run_refdate, run_reftod)
                    if not os.path.exists(fncheck):
                        fncheck = "{}/{}.clm2{}.r.{}-{}.nc".format(rundir, run_refcase, inst_string, run_refdate, run_reftod)
                        if not os.path.exists(fncheck):
                            fncheck = "{}/{}.clm2.r.{}-{}.nc".format(rundir, run_refcase, run_refdate, run_reftod)
                            expect(os.path.exists(fncheck), "ERROR mosart.buildnml: could not find restart file for branch or hybrid start")

            if run_type == "hybrid": infile_text += "finidat_rtm = '{}'\n".format(fncheck)
            if run_type == "branch": infile_text += "nrevsn_rtm  = '{}'\n".format(fncheck)

        create_namelist_infile(case,
                               "{}/user_nl_mosart{}".format(caseroot, inst_string),
                               "{}/cesm_namelist".format(rofconf_dir),
                               infile_text=infile_text)

        # -----------------------------------------------------
        # call build-namelist
        # -----------------------------------------------------

        sysmod  = os.path.join(srcroot, "components/mosart/bld/build-namelist")
        sysmod += " -infile {}/cesm_namelist".format(rofconf_dir)
        sysmod += " -caseroot {}".format(caseroot)
        sysmod += " -scriptsroot {}".format(scriptsroot)
        sysmod += ' -namelist " &mosartexp {} /" '.format(mosart_namelist_opts)
        sysmod += ' -inst_string "{}" {}'.format(inst_string, mosart_bldnml_opts)
        run_cmd_no_fail(sysmod, from_dir=rofconf_dir)

        # -----------------------------------------------------
        # move mosart_in to $RUNDIR
        # -----------------------------------------------------
        if os.path.exists(rundir):
            safe_copy(os.path.join(rofconf_dir, "mosart_in"), os.path.join(rundir, "mosart_in{}".format(inst_string)))

###############################################################################
def _main_func():
###############################################################################
    caseroot = parse_input(sys.argv)
    with Case(caseroot) as case:
        buildnml(case, caseroot, "mosart")

if __name__ == "__main__":
    _main_func()