#!/usr/bin/env python

"""
CAM namelist creator
"""
import sys, os, time, shutil, re, filecmp, imp

_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.XML.standard_module_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__)

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

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

    srcroot             = case.get_value("SRCROOT")
    rundir              = case.get_value("RUNDIR")
    din_loc_root        = case.get_value("DIN_LOC_ROOT")
    atm_ncpl            = case.get_value("ATM_NCPL")
    CAM_NAMELIST_OPTS   = case.get_value("CAM_NAMELIST_OPTS")
    CAM_NML_USE_CASE    = case.get_value("CAM_NML_USE_CASE")
    DEBUG               = case.get_value("DEBUG")
    NTASKS_ATM          = case.get_value("NTASKS_ATM")
    NINST_ATM           = case.get_value("NINST_ATM")
    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")

    #--------------------------------------------------------------------
    # call buildcpp to set both cppdefs and config_cache.xml file for generating namelist
    #--------------------------------------------------------------------
    call_buildcpp = False
    if not os.path.exists(os.path.join(caseroot,"LockedFiles","env_build.xml")):
        call_buildcpp = True
    else:
        file1 = os.path.join(caseroot,"env_build.xml")  
        file2 = os.path.join(caseroot,"LockedFiles","env_build.xml")
        if not filecmp.cmp(file1, file2):
            call_buildcpp = True
    if call_buildcpp:
        cmd = os.path.join(os.path.join(srcroot,"components","cam","cime_config","buildcpp"))
        logger.info("     ...calling cam buildcpp to set build time options")
        try:
            mod = imp.load_source("buildcpp", cmd)
            mod.buildcpp(case)
        except:
            raise

    # Verify that we have a config_cache file (generated by the call to buildcpp) 
    camconf = os.path.join(caseroot, "Buildconf", "camconf")
    filename = os.path.join(camconf, "config_cache.xml")
    expect(os.path.isfile(filename),
           " Missing config_cache.xml - cannot run build-namelist")

    #--------------------------------------------------------------------
    # Invoke cam build-namelist - output will go in $CASEROOT/Buildconf/camconf
    #--------------------------------------------------------------------

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

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

        inst_string = ""
        if ninst > 1:
            single_case_rpointer = os.path.join(rundir, "rpointer.atm")
            inst_string = '_%04d' % inst_counter
            instance_rpointer = os.path.join(rundir,"rpointer.atm"+inst_string)

            # If multi-instance case does not have restart file, use
            # single-case restart for each instance

            if os.path.isfile(single_case_rpointer) and \
               not os.path.isfile(instance_rpointer):
                shutil.copy(single_case_rpointer, instance_rpointer)

        # -----------------------------------------------------
        # create camconf/namelist
        # -----------------------------------------------------

        infile_lines = []

        # This simplifies the filename mangling for different cases.
        def create_ic_filename(inst_string, i_or_r):
            return "%s.cam%s.%s.%s-%s.nc" % \
                (RUN_REFCASE, inst_string, i_or_r, RUN_REFDATE, RUN_REFTOD)

        if RUN_TYPE == 'hybrid':
            ncdata = create_ic_filename(inst_string, 'i')
            # Fallback if no instance-specific file is found.
            if not os.path.isfile(os.path.join(rundir, ncdata)):
                ncdata = create_ic_filename('', 'i')
            infile_lines.append(" ncdata = '" + ncdata + "'")
            if ninst > 1:
                logger.debug("%s is being used for ncdata" % ncdata)

        if RUN_TYPE == 'branch':
            cam_branch_file = create_ic_filename(inst_string, 'r')
            # Fallback if no instance-specific file is found.
            if not os.path.isfile(os.path.join(rundir, cam_branch_file)):
                cam_branch_file = create_ic_filename('', 'r')
            infile_lines.append(" cam_branch_file = '" + cam_branch_file + "'")
            if ninst > 1:
                logger.debug("%s is being used for cam_branch_file" % cam_branch_file)

        dtime = ( 3600 * 24 ) / int(atm_ncpl)
        infile_lines.append(" dtime = " + str(dtime))
        start_ymd = RUN_STARTDATE.replace('-','')
        infile_lines.append(" start_ymd = " + start_ymd)

        if DEBUG:
            infile_lines.append(" state_debug_checks = .true.")

        user_nl_file = os.path.join(caseroot, "user_nl_cam" + inst_string)
        namelist_infile = os.path.join(camconf, "namelist")

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

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

        ntasks = int(NTASKS_ATM) / ninst
# waiting for cime to implement this
#        ntasks = int(case.get_value("NTASKS_PER_INST_ATM"))

        buildnl_opts = ["-ntasks", str(ntasks), "-csmdata", din_loc_root,
                        "-infile", os.path.join(camconf, "namelist")]

        if ('-01-01' in RUN_STARTDATE) or ('-09-01' in RUN_STARTDATE):
            buildnl_opts.append("-ignore_ic_year")
        else:
            buildnl_opts.append("-ignore_ic_date")

        if CAM_NML_USE_CASE != 'UNSET':
            buildnl_opts += ["-use_case", CAM_NML_USE_CASE]

        input_data_list = os.path.join(caseroot, "Buildconf", "cam.input_data_list")
        if os.path.isfile(input_data_list):
            os.remove(input_data_list)
        buildnl_opts += ["-inputdata", input_data_list]

        buildnl_opts += ["-namelist",
                         '" &atmexp ' + CAM_NAMELIST_OPTS + '/" ']

        cmd = os.path.join(srcroot, "components", "cam", "bld", "build-namelist")
        cmd += " " + " ".join(buildnl_opts)

        rc, out, err = run_cmd(cmd, from_dir=camconf)
        expect(rc==0,"Command %s failed rc=%d\nout=%s\nerr=%s"%(cmd,rc,out,err))

        # -----------------------------------------------------
        # copy resolved namelist, atm_in, to rundir
        # -----------------------------------------------------

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

        # -----------------------------------------------------
        # copy drv_flds_in to rundir if it does not exist
        # -----------------------------------------------------

        file1 = os.path.join(camconf, "drv_flds_in")
        file2 = os.path.join(rundir, "drv_flds_in")
        if (os.path.isfile(file1)) and (not os.path.isfile(file2)):
            shutil.copy(file1,file2)

        # -----------------------------------------------------
        # copy aquap_in to rundir if it exists
        # -----------------------------------------------------

        file1 = os.path.join(camconf, "aquap_in")
        file2 = os.path.join(rundir, "aquap_in")
        if (os.path.isfile(file1)):
            shutil.copy(file1,file2)

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

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

if __name__ == "__main__":
    _main_func()

