#!/usr/bin/env python3

"""
Namelist creator for E3SM's CICE 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 == "cice", compname)

    os.chdir(caseroot)

    os.environ["CASEROOT"] = caseroot

    srcroot            = case.get_value("SRCROOT")
    objroot            = case.get_value("OBJROOT")
    cice_auto_decomp   = case.get_value("CICE_AUTO_DECOMP")
    cice_namelist_opts = case.get_value("CICE_NAMELIST_OPTS")
    cice_config_opts   = case.get_value("CICE_CONFIG_OPTS")
    cice_mode          = case.get_value("CICE_MODE")
    cice_blckx         = case.get_value("CICE_BLCKX")
    cice_blcky         = case.get_value("CICE_BLCKY")
    cice_mxblcks       = case.get_value("CICE_MXBLCKS")
    comp_interface     = case.get_value("COMP_INTERFACE")
    ice_grid           = case.get_value("ICE_GRID")
    ice_nx             = case.get_value("ICE_NX")
    ice_ny             = case.get_value("ICE_NY")
    nthrds_ice         = case.get_value("NTHRDS_ICE")
    ntasks_ice         = case.get_value("NTASKS_ICE")
    ninst_ice          = case.get_value("NINST_ICE")
    rundir             = case.get_value("RUNDIR")
    testcase           = case.get_value("TESTCASE")
    casebuild          = case.get_value("CASEBUILD")

    ciceconf_dir = os.path.join(casebuild, "ciceconf")

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

    #--------------------------------------------------------------------
    # Invoke cice configure - output will go in $CASEROOT/Buildconf/ciceconf
    #--------------------------------------------------------------------

    sysmod = os.path.join(srcroot, "components/cice/bld/configure")
    sysmod += " -hgrid {} -nx {} -ny {} ".format(ice_grid, ice_nx, ice_ny)
    sysmod += " -comp_intf {} -cice_mode {} -nodecomp {}".format(comp_interface, cice_mode, cice_config_opts)
    run_cmd_no_fail(sysmod, from_dir=ciceconf_dir)

    # update env_build.xml settings to reflect changes in the configuration
    # this will trigger whether an automatic build is set upon the job resubmission
    if cice_auto_decomp:
        ntasks = int(ntasks_ice / ninst_ice)
        hgrid  = ice_grid
        if ice_grid == 'ar9v2': hgrid = 'ar9v1'
        if ice_grid == 'ar9v4': hgrid = 'ar9v3'

        config = run_cmd_no_fail("./generate_cice_decomp.pl -ccsmroot {} -res {} -nx {} -ny {} -nproc {} -thrds {} -output all".format(srcroot, hgrid, ice_nx, ice_ny, ntasks, nthrds_ice),
                                 from_dir=os.path.join(srcroot, "components/cice/bld")).split()

        if int(config[0]) >= 0:
            cice_blckx, cice_blcky, cice_mxblcks = config[2:5]
            case.set_value("CICE_BLCKX",         cice_blckx)
            case.set_value("CICE_BLCKY",         cice_blcky)
            case.set_value("CICE_MXBLCKS",       cice_mxblcks)
            case.set_value("CICE_DECOMPTYPE",    config[5])
            case.set_value("CICE_DECOMPSETTING", config[6])
        else:
            expect(False, "ERROR cice.buildnml: cice decomp not set for {} on {} x {} procs".format(ice_grid, ntasks, nthrds_ice))

    #--------------------------------------------------------------
    # determine CPP definitions
    #--------------------------------------------------------------

    # determine the actual CPP definitions (these have been determined by the call to configure)
    if os.path.isfile(os.path.join(ciceconf_dir, "CIME_cppdefs")):
        cppdefs = open(os.path.join(ciceconf_dir, "CIME_cppdefs"), "r").read().strip()
    elif os.path.isfile(os.path.join(ciceconf_dir, "CCSM_cppdefs")):
        cppdefs = open(os.path.join(ciceconf_dir, "CCSM_cppdefs"), "r").read().strip()
    cppdefs += " -DBLCKX={} -DBLCKY={} -DMXBLCKS={}".format(cice_blckx, cice_blcky, cice_mxblcks)

    # write out cppdefs to env_build.xml
    # this will determine if need to rebuild at runtime - force user to call $CASE.buildexe

    case.set_value("CICE_CPPDEFS", "{}".format(cppdefs))

    case.flush()

    # write out cppdefs to CICE_cppdefs.new
    # this will force gmake to rebuild when $CASE.buildexe is called

    if testcase != "SBN":
        with open(os.path.join(objroot, "ice/obj/CIME_cppdefs.new"), "w") as fd:
            fd.write(cppdefs + "\n")

    #--------------------------------------------------------------------
    # Invoke cice build-namelist - output will go in $CASEROOT/Buidlconf/ciceconf
    #--------------------------------------------------------------------

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

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

        if ninst_ice > 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.ice{}".format(inst_string))) and \
                   os.path.exists(os.path.join(rundir, "rpointer.ice")):
                safe_copy(os.path.join(rundir, "rpointer.ice"),
                          os.path.join(rundir, "rpointer.ice{}".format(inst_string)))

        # -----------------------------------------------------
        # create ciceconf/namelist
        # -----------------------------------------------------

        create_namelist_infile(case,
                               "{}/user_nl_cice{}".format(caseroot, inst_string),
                               "{}/namelist".format(ciceconf_dir))

        # -----------------------------------------------------
        # call build-namelist- output will go in $CASEROOT/Buildconf/ciceconf/ice_in
        # -----------------------------------------------------

        sysmod  = os.path.join(srcroot, "components/cice/bld/build-namelist")
        sysmod += " -infile {}/Buildconf/ciceconf/namelist".format(caseroot)
        sysmod += " -inputdata {}/Buildconf/cice.input_data_list".format(caseroot)
        sysmod += " -rundir {}".format(rundir)
        sysmod += " -caseroot {}".format(caseroot)
        sysmod += " -cimeroot {}".format(_CIMEROOT)
        sysmod += ' -inst_string "{}"'.format(inst_string)
        sysmod += ' -namelist "&cice {}/" '.format(cice_namelist_opts)
        sysmod += " -config config_cache.xml"
        run_cmd_no_fail(sysmod, from_dir=ciceconf_dir)

        # -----------------------------------------------------
        # Copy resolved namelist to $RUNDIR
        # -----------------------------------------------------

        if os.path.exists(rundir):
            safe_copy(os.path.join(ciceconf_dir, "ice_in"), os.path.join(rundir, "ice_in{}".format(inst_string)))

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

if __name__ == "__main__":
    _main_func()
