#!/usr/bin/env python

"""POP namelist creator
"""

# Typically ignore this.
# pylint: disable=invalid-name

# Disable these because this is our standard setup
# pylint: disable=wildcard-import,unused-wildcard-import,wrong-import-position
# pylint: disable=multiple-imports
import os, shutil, sys, glob, stat, filecmp, imp

CIMEROOT = os.environ.get("CIMEROOT")
if CIMEROOT is None:
    raise SystemExit("ERROR: must set CIMEROOT environment variable")
sys.path.append(os.path.join(CIMEROOT, "scripts", "Tools"))

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

logger = logging.getLogger(__name__)

# pylint: disable=too-many-arguments,too-many-locals,too-many-branches,too-many-statements
###############################################################################
def buildnml(case, caseroot, compname):
###############################################################################
    """Build the pop namelist """

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

    cimeroot = case.get_value("CIMEROOT")
    srcroot = case.get_value("SRCROOT")
    din_loc_root = case.get_value("DIN_LOC_ROOT")
    continue_run = case.get_value("CONTINUE_RUN")
    get_refcase = case.get_value("GET_REF_CASE")
    ninst = case.get_value("NINST_OCN")
    ocn_grid = case.get_value("OCN_GRID")
    run_type = case.get_value("RUN_TYPE")
    ocn_bgc_config = case.get_value("OCN_BGC_CONFIG")
    run_refcase = case.get_value("RUN_REF CASE")
    run_refdate = case.get_value("RUN_REFDATE")
    rundir = case.get_value("RUNDIR")
    testcase = case.get_value("TESTCASE")
    ntasks = case.get_value("NTASKS_PER_INST_OCN")

    # Obtain correct pop rpointer files
    if testcase != 'SBN':
        if get_refcase and run_type == "startup" and not continue_run:
            # During prestage step, rpointer files are copied from refdir
            # Get rid of old rpointer files if they exist and copy them
            # independently of the prestage.  This is needed for rerunability
            # of cases from `refcase` data for first run

            filenames = glob.glob(rundir + '/rpointer.ocn*')
            for filename in filenames:
                os.remove(filename)

            refdir = os.path.join(din_loc_root, "ccsm4_init", run_refcase, run_refdate)
            filenames = glob.glob(refdir + '/rpointer.ocn*')
            for filename in filenames:
                shutil.copy(filename, rundir)
                destfile = os.path.join(rundir, filename)
                st = os.stat(destfile)
                os.chmod(destfile, st.st_mode | stat.S_IWUSR)

    # call buildcpp to obtain config_cache.xml file which is needed to set pop 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", "pop", "cime_config", "buildcpp"))
        logger.info("     ... buildnml: calling pop buildcpp to set build time options")
        try:
            mod = imp.load_source("buildcpp", cmd)
            mod.buildcpp(case)
        except:
            raise

    # Set pop configuration directory
    confdir = os.path.join(caseroot, "Buildconf", "popconf")
    if not os.path.isdir(confdir):
        os.makedirs(confdir)

    # Make sure that rundir exists, if not make it
    if not os.path.exists(rundir):
        os.makedirs(rundir)

    # Clean up old files
    input_data_list = os.path.join(caseroot, "Buildconf", "pop.input_data_list")
    if os.path.isfile(input_data_list):
        os.remove(input_data_list)

    # Loop over instances
    inst_counter = 1
    inst_string = ""
    while inst_counter <= ninst:
        # determine instance string
        if ninst > 1:
            inst_string = '_' + '%04d' % inst_counter

        marbl_nt_build = case.get_value("MARBL_NT")
        if marbl_nt_build > 0:
            # Generate marbl_in and ecosys_diagnostics (files POP needs to read for MARBL)
            # Note that they are only needed when running with the ecosystem tracer module
            _construct_marbl_in(caseroot, ocn_grid, run_type, continue_run, ocn_bgc_config,
                                srcroot, confdir, inst_string, marbl_nt_build)
            _construct_ecosys_diagnostics(caseroot, ocn_grid, run_type, continue_run,
                                          ocn_bgc_config, srcroot, confdir, inst_string)

        # If multi-instance = case does not have restart file,
        # use single-= case restart for each instance
        suffixes = ["ovf", "restart", "tavg"]
        for suffix in suffixes:
            # See if rpointer.ocn.`suffix` exists and rpointer.ocn`inst_string`.`suffix` does not
            file1 = os.path.join(rundir, "rpointer.ocn" + inst_string + "." + suffix)
            file2 = os.path.join(rundir, "rpointer.ocn." + suffix)
            if not os.path.isfile(file1) and os.path.isfile(file2):
                shutil.copy(file2, file1)

        # create namelist_infile using user_nl_file as input
        user_nl_file = os.path.join(caseroot, "user_nl_pop" + inst_string)
        expect(os.path.isfile(user_nl_file),
               "Missing required user_nl_file %s " %(user_nl_file))
        infile = os.path.join(confdir, "namelist_infile")
        create_namelist_infile(case, user_nl_file, infile)

        # set environment variable declaring type of pop restart file (nc vs bin)
        # RESTART_INPUT_TS_FMT is needed by pop's build-namelist and is not in any xml files;
        # it comes from rpointer.ocn.restart, which is in `rundir` for continued runs,
        # but is in `refdir` for hybrid / branch runs that are not continuations

        pointer_file = os.path.join(rundir, "rpointer.ocn" + inst_string + ".restart")
        if get_refcase and run_type != 'startup' and not continue_run:
            # During prestage step, rpointer files are copied from refdir
            for refdir_path in ("cesm2_init", "ccsm4_init"):
                refdir = os.path.join(din_loc_root, refdir_path, run_refcase, run_refdate)
                pointer_file = os.path.join(refdir, "rpointer.ocn" + inst_string + ".restart")
                if os.path.isfile(pointer_file):
                    break
            if not os.path.isfile(pointer_file):
                pointer_file = os.path.join(rundir, "rpointer.ocn" + inst_string + "restart")

        if run_type == 'startup' and not continue_run:
            check_pointer_file = False
        else:
            check_pointer_file = True

        _format = 'bin'
        if check_pointer_file:
            expect(os.path.isfile(pointer_file),
                   "Missing required pointer_file %s ---"
                   "has pop initial data been prestaged to %s?" %(pointer_file, rundir))
            if 'RESTART_FMT=nc' in open(pointer_file).read():
                _format = 'nc'
        os.environ["RESTART_INPUT_TS_FMT"] = _format

        # ------------------------------------------------------------------------------
        # call build-namelist - output will go in caseroot/Buildconf/popconf/pop_in
        # ------------------------------------------------------------------------------

        # check to see if "-preview" flag should be passed
        if os.environ.get("PREVIEW_NML") is not None:
            preview_flag = "-preview"
        else:
            preview_flag = ""

        # determine the directory containing build-namelist script
        # first check to see if build-namelist exists in SourceMods, if it exists use it
        bldnml_dir = os.path.join(srcroot, "components", "pop", "bld")
        cfg_flag = ""
        if os.path.isfile(os.path.join(caseroot, "SourceMods", "src.pop", "build-namelist")):
            bldnml_dir = os.path.join(caseroot, "SourceMods", "src.pop")
            cfg_flag = "-cfg_dir " + os.path.join(srcroot, "components", "pop", "bld")

        # now call build-namelist
        cmd = os.path.join(bldnml_dir, "build-namelist")
        command = "%s %s %s -infile %s -caseroot %s -cimeroot %s -ocn_grid %s -ntasks %s" \
                  % (cmd, cfg_flag, preview_flag, infile, caseroot, cimeroot, ocn_grid, ntasks)
        if inst_string:
            command += " -inst_string %s " % inst_string

        rc, out, err = run_cmd(command, from_dir=confdir)
        expect(rc == 0, "Command %s failed rc=%d\nout=%s\nerr=%s"%(cmd, rc, out, err))
        if out:
            logger.info("     %s", out)
        if err:
            logger.info("     %s", err)

        # Note that pop's build-namelist invokes xmlchange - so must re-read the xml files variables
        # back into case
        case.read_xml()

        # copy pop namelist files from confdir to rundir
        if os.path.isdir(rundir):
            # pop_in
            _copy_file('pop_in', confdir, rundir, inst_string)

            # marbl_in and marbl_diagnostics_operators
            if marbl_nt_build > 0:
                _copy_file('marbl_in'+inst_string, confdir, rundir)
                _copy_file('marbl_diagnostics_operators', confdir, rundir)

            # tavg_contents file
            _copy_file(ocn_grid + "_tavg_contents", confdir, rundir)

            # increment instance counter
            inst_counter = inst_counter + 1

###############################################################################

def _copy_file(filename, src_dir, dest_dir, inst_string=None):
    """ Copy a given file from confdir ($CASEROOT/Buildconf/popconf/) to rundir ($RUNDIR)
    """
    file_src = os.path.join(src_dir, filename)
    file_dest = os.path.join(dest_dir, filename)
    if inst_string:
        file_dest += inst_string
    shutil.copy(file_src, file_dest)

###############################################################################

def _construct_marbl_in(caseroot, ocn_grid, run_type, continue_run, ocn_bgc_config, srcroot,
                        confdir, inst_string, marbl_nt_build):
    # import wrappers to some MARBL calls
    sys.path.append(os.path.join(srcroot, "components", "pop", "MARBL_scripts"))
    from MARBL_wrappers import MARBL_settings_for_POP

    # ------------------------------------------------------------------------------
    # call MARBL's tool to generate a settings file
    # output will go to CASEROOT/Buildconf/popconf/marbl_in
    # ------------------------------------------------------------------------------

    # (i) Generate MARBL_settings_for_POP object
    MARBL_dir = os.path.join(srcroot, "components", "pop", "externals", "MARBL")
    MARBL_settings = MARBL_settings_for_POP(MARBL_dir, "user_nl_marbl"+inst_string, caseroot,
                                            ocn_grid, run_type, continue_run, ocn_bgc_config)

    # (ii) compare MARBL_NT to value from buildcpp (in config_cache)
    #      Abort if values do not match
    MARBL_NT = MARBL_settings.get_MARBL_NT()
    if MARBL_NT != marbl_nt_build:
        error_msg = "POP was built expecting %d MARBL tracers, but current configuration has %d."
        error_msg = error_msg + "\nRun \"./case.build --clean ocn\" and then rebuild."
        logger.error(error_msg, marbl_nt_build, MARBL_NT)
        sys.exit(1)

    # (iii) write settings file
    MARBL_settings.write_settings_file(os.path.join(confdir, "marbl_in"+inst_string))

###############################################################################

def _construct_ecosys_diagnostics(caseroot, ocn_grid, run_type, continue_run, ocn_bgc_config,
                                  srcroot, confdir, inst_string):
    # import wrappers to some MARBL calls
    sys.path.append(os.path.join(srcroot, "components", "pop", "MARBL_scripts"))
    from MARBL_wrappers import MARBL_settings_for_POP
    from MARBL_wrappers import write_ecosys_diagnostics_file
    from MARBL_wrappers import MARBL_diagnostics_for_POP

    # ------------------------------------------------------------------------------
    # call MARBL's tool to generate a generic diagnostics file;
    # output will go to CASEROOT/Buildconf/popconf/marbl_diagnostics
    # (which will then be used by ocn.tavg.ecosys.csh to generate ecosys_tavg_contents)
    #
    # Also pass all MARBL tracer short names to another tool to generate
    # a similar generic diagnostics file for the POP-owned ecosys-related
    # diagnostics (tracer diagnostics MARBL expects the GCM to compute).
    #
    # In: MARBL_settings (or user_nl_marbl)
    # Out: Buildconf/popconf/ecosys_diagnostics
    # Out: Buildconf/popconf/marbl_diagnostics_list
    # ------------------------------------------------------------------------------

    # (i) Generate MARBL_settings_for_POP object
    MARBL_dir = os.path.join(srcroot, "components", "pop", "externals", "MARBL")
    MARBL_settings = MARBL_settings_for_POP(MARBL_dir, "user_nl_marbl"+inst_string, caseroot,
                                            ocn_grid, run_type, continue_run, ocn_bgc_config)

    # (ii) Create ecosys_diagnostics with POP-generated diagnostics
    write_ecosys_diagnostics_file(MARBL_settings.get_tracer_names(),
                                  MARBL_settings.get_autotroph_names(),
                                  MARBL_settings.get_zooplankton_names(),
                                  MARBL_settings.get_autotroph_names(calcifier_only=True),
                                  MARBL_settings.ladjust_bury_coeff(),
                                  os.path.join(confdir, "ecosys_diagnostics"))

    # (iii) Import MARBL_wrappers and generate MARBL_diagnostics_for_POP object
    MARBL_diagnostics = MARBL_diagnostics_for_POP(MARBL_dir, caseroot, MARBL_settings)

    # (iv) append MARBL diagnostics to ecosys_diagnostics
    #      This also creates marbl_diagnostics_list to help track which
    #      diagnostics came from MARBL
    MARBL_diagnostics.write_diagnostics_file(os.path.join(confdir, "ecosys_diagnostics"),
                                             os.path.join(confdir, "marbl_diagnostics_list"),
                                             append=True)

###############################################################################

def _main_func():

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

if __name__ == "__main__":
    _main_func()
