#!/usr/bin/env python

"""
CLM namelist creator
"""
import sys, os, shutil

_CIMEROOT = os.environ.get("CIMEROOT")
if _CIMEROOT is None:
    raise SystemExit("ERROR: must set CIMEROOT environment variable")

_LIBDIR = os.path.join(_CIMEROOT, "scripts", "Tools")
sys.path.append(_LIBDIR)

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

logger = logging.getLogger(__name__)

_config_cache_template = """
<?xml version="1.0"?>
<config_definition>
<commandline></commandline>
<entry id="phys" value="{clm_phys}" list="" valid_values="clm4_5,clm5_0">Specifies clm physics</entry>
</config_definition>
"""

###############################################################################
def buildnml(case, caseroot, compname):
###############################################################################
    """Build the clm namelist """

    # Build the component namelist 
    if compname != "clm":
        raise AttributeError

    lnd_root = case.get_value("COMP_ROOT_DIR_LND")
    din_loc_root = case.get_value("DIN_LOC_ROOT")
    configuration = case.get_value("CLM_CONFIGURATION")
    structure = case.get_value("CLM_STRUCTURE")
    ccsm_co2_ppmv = case.get_value("CCSM_CO2_PPMV")
    clm_co2_type = case.get_value("CLM_CO2_TYPE")
    clm_namelist_opts = case.get_value("CLM_NAMELIST_OPTS")
    clm_bldnml_opts = case.get_value("CLM_BLDNML_OPTS")
    clm_nml_use_case = case.get_value("CLM_NML_USE_CASE")
    clm_force_coldstart = case.get_value("CLM_FORCE_COLDSTART")
    lnd_tuning_mode = case.get_value("LND_TUNING_MODE")
    clm_accelerated_spinup = case.get_value("CLM_ACCELERATED_SPINUP")
    comp_atm = case.get_value("COMP_ATM")
    lnd_grid = case.get_value("LND_GRID")
    lnd_domain_path = case.get_value("LND_DOMAIN_PATH")
    lnd_domain_file = case.get_value("LND_DOMAIN_FILE")
    ninst_lnd = case.get_value("NINST_LND")
    rundir = case.get_value("RUNDIR")
    run_type = case.get_value("RUN_TYPE")
    run_startdate = case.get_value("RUN_STARTDATE")
    run_refcase = case.get_value("RUN_REFCASE")
    run_refdate = case.get_value("RUN_REFDATE")
    run_reftod = case.get_value("RUN_REFTOD")
    glc_nec = case.get_value("GLC_NEC")
    mask = case.get_value("MASK_GRID")

    # -----------------------------------------------------
    # Set clmconf
    # -----------------------------------------------------

    clmconf = os.path.join(caseroot, "Buildconf", "clmconf")
    if not os.path.isdir(clmconf):
        os.makedirs(clmconf)

    # -----------------------------------------------------
    # Create config_cache.xml file 
    # -----------------------------------------------------

    # Note that build-namelist utilizes the contents of the config_cache.xml file in 
    # the namelist_defaults.xml file to obtain namelist variables

    clm_phys = case.get_value("CLM_PHYSICS_VERSION")

    config_cache_text = _config_cache_template.format(clm_phys=clm_phys)
    config_cache_path = os.path.join(caseroot, "Buildconf", "clmconf", "config_cache.xml")
    with open(config_cache_path, 'w') as config_cache_file:
        config_cache_file.write(config_cache_text)

    # -----------------------------------------------------
    # Determine input arguments into build-namelist
    # -----------------------------------------------------

    startfile_type = "finidat"
    start_type = "default"
    if run_type == "hybrid":
        start_type = "startup"
    elif run_type != "startup":
        start_type = run_type

    if run_type == "branch":
        startfile_type = "nrevsn"
        if clm_force_coldstart == "on":
            clm_force_coldstart = "off"
            logger.warning( "WARNING: You've turned on CLM_FORCE_COLDSTART for a branch run_type, which is a contradiction, the coldstart will be ignored\n" +
                            "  turn off CLM_FORCE_COLDSTART, or set RUN_TYPE=hybrid to get rid of this warning")


    if (clm_force_coldstart == "on"):
        logger.warning( "WARNING: CLM is starting up from a cold state" )
        start_type = "cold"

    if lnd_grid == 'T31':
        lnd_grid = '48x96'
    if lnd_grid == 'T42':
        lnd_grid = '64x128'
    if lnd_grid == 'T85':
        lnd_grid = '128x256'
    if lnd_grid == 'T341':
        lnd_grid = '512x1024'

    clmusr = ""
    if lnd_grid == "CLM_USRDAT":
        clm_usrdat_name = case.get_value("CLM_USRDAT_NAME")
        lnd_grid = clm_usrdat_name
        clmusr = " -clm_usr_name %s "%clm_usrdat_name

    if comp_atm != "datm":
        nomeg = "-no-megan"
    else:
        nomeg = ""

    if clm_nml_use_case != "UNSET":
        usecase = "-use_case %s" %clm_nml_use_case
    else:
        usecase = ""

    if ( (mask != "null") and (mask != "UNSET") ):
        gridmask = "-mask %s" %mask
    else:
        gridmask = ""

    start_ymd = run_startdate.replace('-','')

    if ('-01-01' in run_startdate) or ('-09-01' in run_startdate):
        ignore = "-ignore_ic_year"
    else:
        ignore = "-ignore_ic_date"

    tuning = "-lnd_tuning_mode %s "%lnd_tuning_mode
    
    spinup = "-clm_accelerated_spinup %s "%clm_accelerated_spinup
    
    infile = os.path.join(clmconf, "namelist")
    
    inputdata_file = os.path.join(caseroot,"Buildconf","clm.input_data_list")
    
    lndfrac_file = os.path.join(lnd_domain_path,lnd_domain_file)
    
    config_cache_file = os.path.join(caseroot,"Buildconf","clmconf","config_cache.xml")

    # -----------------------------------------------------
    # Clear out old data
    # -----------------------------------------------------

    if os.path.exists(inputdata_file):
        os.remove(inputdata_file)

    # -----------------------------------------------------
    # loop over instances
    # -----------------------------------------------------

    ninst = int(ninst_lnd)
    for inst_counter in range(1, ninst+1):

        # determine instance string
        inst_string = ""
        if ninst > 1:
            inst_string = '_' + '%04d' % inst_counter

        # If multi-instance case does not have restart file, use
        # single-case restart for each instance
        rpointer = "rpointer.lnd" 
        if (os.path.isfile(os.path.join(rundir,rpointer)) and
            (not os.path.isfile(os.path.join(rundir,rpointer + inst_string)))):
            shutil.copy(os.path.join(rundir, rpointer),
                        os.path.join(rundir, rpointer + inst_string))

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

        if run_type == "hybrid" or run_type == "branch":
            clm_startfile = "%s.clm2%s.r.%s-%s.nc"%(run_refcase,inst_string,run_refdate,run_reftod)
            if not os.path.exists(os.path.join(rundir, clm_startfile)):
                clm_startfile = "%s.clm2.r.%s-%s.nc"%(run_refcase,run_refdate,run_reftod)
                logger.warning( "WARNING: the start file being used for a multi-instance case is a single instance: "+clm_startfile )

            clm_icfile = "%s = \'%s\'"%(startfile_type, clm_startfile)
        else:
            clm_icfile = ""

        infile_lines = []
        infile_lines.append(clm_icfile)

        user_nl_file = os.path.join(caseroot, "user_nl_clm" + inst_string)
        namelist_infile = os.path.join(clmconf, "namelist")

        create_namelist_infile(case, user_nl_file, namelist_infile, "\n".join(infile_lines))

        cmd = os.path.join(lnd_root,"bld","build-namelist")

        command = ("%s -cimeroot %s -infile %s -csmdata %s -inputdata %s %s -namelist \"&clm_inparm  start_ymd=%s %s/ \" "
                   "%s %s -res %s %s -clm_start_type %s -envxml_dir %s "
                   "-configuration %s -structure %s "
                   "-lnd_frac %s -glc_nec %s -co2_ppmv %s -co2_type %s -config %s "
                   "%s %s %s %s"
                   %(cmd, _CIMEROOT, infile, din_loc_root, inputdata_file, ignore, start_ymd, clm_namelist_opts, 
                     nomeg, usecase, lnd_grid, clmusr, start_type, caseroot,
                     configuration, structure,
                     lndfrac_file, glc_nec, ccsm_co2_ppmv, clm_co2_type, config_cache_file,
                     clm_bldnml_opts, spinup, tuning, gridmask))

        rc, out, err = run_cmd(command, from_dir=clmconf)
        expect(rc==0,"Command %s failed rc=%d\nout=%s\nerr=%s"%(cmd,rc,out,err))
        if out is not None:
            logger.debug("     %s"%out)
        if err is not None:
            logger.debug("     %s"%err)

        # -----------------------------------------------------
        # copy resolved namelist to rundir
        # -----------------------------------------------------

        if os.path.isdir(rundir):
            file1 = os.path.join(clmconf, "lnd_in")
            file2 = os.path.join(rundir, "lnd_in")
            if ninst > 1:
                file2 += inst_string
            logger.debug("CLM namelist copy: file1 %s file2 %s " %(file1, file2))
            shutil.copy(file1,file2)

###############################################################################
def _main_func():

    caseroot = parse_input(sys.argv)
    with Case(caseroot) as case:
        buildnml(case, caseroot, "clm")

if __name__ == "__main__":
    _main_func()
