<filename>wrftamer/initialize_wrf_namelist.py
# coding: utf-8
import os

"""
This script will initialize the namelist.input (WRF) and namelist.wps (WPS).
It is needed to provide a config-file which contains information about the desired setting,
and a namelist.template
"""


def initialize_wrf_namelist(namelist_vars: dict, namelistfile: str, templatefile=None):
    # read template and configuration
    if templatefile is None:
        mypath = (
            os.path.split(os.path.realpath(__file__))[0]
            + "/../wrftamer/resources/namelist.template"
        )
    else:
        mypath = templatefile

    max_keylen = 0
    dom_keys = (
        "start_year start_month start_day start_hour start_date "
        "end_year end_month end_day end_hour end_date"
    )

    with open(mypath, "r") as f:

        tpl = f.read()
        namelist = tpl.format(**namelist_vars).split("\n")

        section = None
        # loop over all lines in the namelist template and reconstruct each one
        # using proper formatting and the corresponding values from the config file
        for line_nb, line in enumerate(namelist):
            if "&" in line:
                section = line.strip()[1:]
            elif line.strip() == "/":
                section = None
            if "=" not in line:
                continue
            key_full, val = line.split("=")
            key = key_full.strip()
            if key in dom_keys:
                val = " ".join([val.strip()] * namelist_vars["max_dom"])
            elif key in namelist_vars:
                val = str(namelist_vars[key]) + ","

            if section == "geogrid" and key in "dx dy":
                val = str(str(namelist_vars[key]).split(",")[0] + ",")

            if key == "run_hours":
                val = (
                    namelist_vars["dtend"] - namelist_vars["dtbeg"]
                ).total_seconds() // 3600
                val = f"{val:02.0f},"
            if section == "domains" and key == "eta_levels":
                # we might need to insert custom eta levels at this point
                eta_keys = [
                    k for k in namelist_vars.keys() if k.startswith("eta_levels")
                ]
                if eta_keys:
                    key = eta_keys[0]
                    val = namelist_vars[eta_keys[0]]
                    # append the remaining lines to the "value" (a bit hacky)
                    for eta_key in eta_keys[1:]:
                        val += "\n" + " ".join(
                            [
                                eta_key,
                                " " * (100 - len(eta_key)),
                                "=",
                                namelist_vars[eta_key],
                            ]
                        )
                        max_keylen = max(len(eta_key), max_keylen)
            elem = str(val).split(",")
            if (len(elem) > namelist_vars["max_dom"]) and not key.startswith(
                "eta_levels"
            ):
                val = ",".join(elem[: namelist_vars["max_dom"]]) + ","

            # update this line of the namelist
            namelist[line_nb] = " ".join([key, " " * (100 - len(key)), "=", val])
            max_keylen = max(len(key), max_keylen)

    with open(namelistfile, "w") as f:
        f.write("\n".join(namelist).replace(" " * (100 - max_keylen), ""))
