#!/usr/bin/env python3
#
# parm_tker.py --- Turbulent Kinetic Energy options
#                  See HYSPLIT Tutorial Sec. 9.6
#
# Change history:
# 15 Oct 2025 (BB) - standard-aligned, cross-platform version
# 12 Nov 2025 (BB) - fix met_dir to TTR/captex; trailing seps; -q path; minor cleanups
# 20 Nov 2025 (SZ) - create image files when DSP=NO
#                  - show HYSPLIT runtime messages
# 18 Jun 2026 (SZ) - standardize code for run_cmd(), convert_html_to_img(),
#                    etc. across Python scripts
#                  - ensure encoding is specified when opening files

import os
import sys
import platform
from pathlib import Path
import subprocess
import webbrowser

if sys.version_info < (3, 8):
    print("Please use Python 3.8 or higher.")
    sys.exit(1)

# ---------------- OS setup ----------------
system = platform.system().lower()
if system == "windows":
    PGM = Path(sys.argv[1]) if len(sys.argv) > 1 else Path("C:\\hysplit")
    OUT = Path(sys.argv[2]) if len(sys.argv) > 2 else PGM / "working"
    DSP = sys.argv[3] if len(sys.argv) > 3 else "YES"
    TTR = Path(os.environ.get("TTR", "C:\\Tutorial"))
    exe_suffix = ".exe"
else:
    PGM = Path(sys.argv[1]) if len(sys.argv) > 1 else Path.home() / "hysplit"
    OUT = Path(sys.argv[2]) if len(sys.argv) > 2 else PGM / "working"
    DSP = sys.argv[3] if len(sys.argv) > 3 else "YES"
    TTR = Path(os.environ.get("TTR", str(Path.home() / "tutorial")))
    exe_suffix = ""

DSP = str(DSP).upper()

if not PGM.exists():
    script = Path(__file__).name if "__file__" in globals() else "parm_tker.py"
    print(f"ERROR: HYSPLIT not found at {PGM}")
    print(f"""
Please specify the path to your HYSPLIT installation using a command-line argument:

python {script} YOUR_HYSPLIT_DIRECTORY [OUT] [DSP (YES|NO)]

Example:
    python {script} C:\\hysplit C:\\hysplit\\working YES
""")
    sys.exit(1)

if not TTR.exists():
    print(f"ERROR: HYSPLIT tutorial not found at {TTR}")
    print("""
Please set the TTR environment variable to point to your HYSPLIT tutorial directory.

For Windows users (Command Prompt):
   SET TTR=YOUR_HYSPLIT_TUTORIAL_DIRECTORY

For Linux users (bash shell):
   export TTR='YOUR_HYSPLIT_TUTORIAL_DIRECTORY'

For Linux users (csh or tcsh shell):
   setenv TTR 'YOUR_HYSPLIT_TUTORIAL_DIRECTORY'
""")
    sys.exit(1)

os.chdir(OUT)

# ------------------------------------------------------------------
# Helper functions
# ------------------------------------------------------------------

def run_cmd(cmd, desc="", allow_fail=False, stdout_file=None, append=False):
    if desc:
        print(f"{desc}...")

    exitcode = 0
    with subprocess.Popen(cmd,
                         stdout=subprocess.PIPE,
                         stderr=subprocess.STDOUT,
                         text=True) as p:
        if stdout_file:
            mode="a" if append else "w"
            with open(stdout_file, mode, encoding="utf-8") as f_out:
                for ln in p.stdout:
                    print(ln, end="")
                    f_out.write(ln)
        else:
            with p.stdout as f:
                for ln in iter(f.readline, b'\n'):
                    if len(ln) == 0:
                        break
                    print(ln.rstrip())

        done = False
        while not done:
            try:
                exitcode = p.wait(timeout=1)
                done = True
            except subprocess.TimeoutExpired:
                print("WARN: subprocess timed out.")

    if exitcode != 0:
        if allow_fail:
            print(f"WARN: command failed with exit code {exitcode} -> {' '.join(cmd)}")
        else:
            print(f"ERROR: command failed with exit code {exitcode} -> {' '.join(cmd)}")
            sys.exit(1)


def convert_html_to_img(html_file, img_name_frame_nums):
    # Split SVG into Linux-consistent naming: F00-sec71_001.svg, etc.
    temp_img_name = "temp"
    run_cmd([str(PGM / "exec" / f"splitsvg{exe_suffix}"),
             f"-i{html_file}", f"-o{temp_img_name}.svg"])

    # Locate ImageMagick for cross-platform conversion
    if system == "windows":
        convert_exe = None
        imgk_dir = next(
            (k for k in os.environ.get('PATH', '').split(os.pathsep)
             if "ImageMagick" in k),
            None
        )
        if imgk_dir:
            # Use magick.exe for ImageMagick 7.x and later. Use convert.exe for earlier versions.
            for cmd in ('magick', 'convert'):
                candidate = os.path.join(imgk_dir, f'{cmd}{exe_suffix}')
                if os.path.exists(candidate):
                    convert_exe = candidate
                    break
    else:
        convert_exe = "convert"

    if convert_exe:
        # Match Linux naming for DSP=NO mode exactly
        for img_name, frame_num in img_name_frame_nums:
            svg = Path(f"F{frame_num:02d}-{temp_img_name}.svg")
            run_cmd([convert_exe, str(svg), img_name])
            print(f"Converted {svg} -> {img_name}")
            svg.unlink()
    else:
        print("ImageMagick not found — skipping .svg conversion")


def clean(path):
    Path(path).unlink(missing_ok=True)


def write_ascdata(pgm: Path):
    """Re-create ASCDATA.CFG with a trailing separator on the bdyfiles line."""
    clean("ASCDATA.CFG")
    with open("ASCDATA.CFG", "w", encoding="ascii") as f:
        f.write("-90.0   -180.0  lat/lon of lower left corner\n")
        f.write("1.0     1.0     lat/lon spacing in degrees\n")
        f.write("180     360     lat/lon number of data points\n")
        f.write("2               default land use category\n")
        f.write("0.2             default roughness length (m)\n")
        f.write(f"'{pgm / 'bdyfiles'}{os.sep}'  directory of files\n")


def write_control_tker(
    syr, smo, sda, shr,
    lat, lon, lvl,
    run_hours, top,
    met_dir: Path, dat_name: str,
    inp_name: str
):
    """CONTROL for all TKE/anisotropy runs."""
    with open("CONTROL", "w", encoding="ascii") as f:
        f.write(f"{syr:02d} {smo:02d} {sda:02d} {shr:02d}\n")
        f.write("1\n")
        f.write(f"{lat:.2f} {lon:.2f} {lvl:.1f}\n")
        f.write(f"{run_hours}\n")
        f.write("0\n")
        f.write(f"{top:.1f}\n")
        f.write("1\n")
        f.write(f"{str(met_dir)}{os.sep}\n")  # met directory must end with sep
        f.write(f"{dat_name}\n")
        f.write("1\n")
        f.write("PMCH\n")
        f.write("67000.0\n")
        f.write("3.0\n")
        f.write("00 00 00 00 00\n")
        f.write("1\n")
        f.write("42.0 -78.0\n")
        f.write("0.05 0.05\n")
        f.write("15.0 25.0\n")
        f.write(f".{os.sep}\n")
        f.write(f"{inp_name}\n")
        f.write("2\n")
        f.write("800 1000\n")
        f.write("83 09 26 03 00\n")
        f.write("83 09 28 15 00\n")
        f.write("00 01 00\n")
        f.write("1\n")
        f.write("0.0 0.0 0.0\n")
        f.write("0.0 0.0 0.0 0.0 0.0\n")
        f.write("0.0 0.0 0.0\n")
        f.write("0.0\n")
        f.write("0.0\n")


def write_setup(lines):
    clean("SETUP.CFG")
    with open("SETUP.CFG", "w", encoding="ascii") as f:
        f.write("&SETUP\n")
        for line in lines:
            f.write(f"{line}\n")
        f.write("/\n")


def run_hycs_std(hycs_exec: Path, label: str, inp_name: str):
    clean(inp_name)
    run_cmd([str(hycs_exec)], f"Run hycs_std for {label}")


def write_labels(label: str, script_label: str):
    with open("LABELS.CFG", "w", encoding="ascii") as f:
        f.write(f"'TITLE&','### {script_label}: {label} ### &'\n")


def run_concplot(concplot_exec: Path, map_path: Path, inp_name: str, label: str,
                 data_txt: Path, img_name_frame_nums=None):
    clean("concplot.html")
    cmd = [
        str(concplot_exec),
        "+g1",
        f"-i{inp_name}",
        "-z100",
        f"-j{map_path}",
        "-x1.0E+12",
        "-upg",
        "-h41.0:-82.0",
        "-g0:200",
        "-c4",
        "-v50000+20000+10000+5000+2000+1000+500",
        "-b1000",
        "-t1000",
        f"-q{data_txt}",
    ]
    run_cmd(cmd, f"Create concplot for {label}")
    html = Path("concplot.html")
    if html.exists():
        if DSP == "YES":
            try:
                webbrowser.open(html.resolve().as_uri())
            except Exception as e:
                print(f"WARNING: could not open concplot.html in browser: {e}")
        else:
            convert_html_to_img(str(html), img_name_frame_nums)


# ---------------- Main ----------------
def main():
    script_label = Path(__file__).name if "__file__" in globals() else "parm_tker.py"

    map_path = PGM / "graphics" / "arlmap"
    inp_name = "hysplit2.bin"
    hycs_exec = PGM / "exec" / f"hycs_std{exe_suffix}"
    concplot_exec = PGM / "exec" / f"concplot{exe_suffix}"

    if not hycs_exec.exists():
        print(f"ERROR: hycs_std not found at {hycs_exec}")
        sys.exit(1)
    if not concplot_exec.exists():
        print(f"ERROR: concplot not found at {concplot_exec}")
        sys.exit(1)

    write_ascdata(PGM)

    # CAPTEX configuration (tutorial data under TTR/captex)
    syr, smo, sda, shr = 83, 9, 25, 17
    lat, lon, lvl = 39.90, -84.22, 10.0
    run_hours = 11
    top = 10000.0

    met_dir = TTR / "captex"
    dat_name = "captex2_wrf27uw.bin"
    data_txt = TTR / "captex" / "data_case.txt"

    if not met_dir.exists():
        print(f"WARNING: meteorology directory not found: {met_dir}")
    if not (met_dir / dat_name).exists():
        print(f"WARNING: met file not found: {met_dir / dat_name}")
    if not data_txt.exists():
        print(f"WARNING: measurement file not found: {data_txt}")

    write_control_tker(
        syr, smo, sda, shr,
        lat, lon, lvl,
        run_hours, top,
        met_dir, dat_name,
        inp_name,
    )

    # 1) TKE Turbulence
    write_setup([
        "initd = 3,",
        "kblt = 3,",
        "kbls = 2,",
        "kmsl = 1,",
        "numpar = 5000,",
        "maxpar = 100000,",
        "ndump = 1,",
        "ncycl = 3,",
    ])
    run_hycs_std(hycs_exec, "TKE Turbulence", inp_name)
    write_labels("TKE Turbulence", script_label)
    run_concplot(concplot_exec, map_path, inp_name, "TKE Turbulence",
                 data_txt, img_name_frame_nums=[("tker002.png", 1)])

    if DSP == "YES":
        input("Press Enter to continue...\n")

    # 2) Compute Anisotropy (diagnostic run: compute then disable)
    write_setup([
        "initd = 3,",
        "kblt = 3,",
        "kbls = 2,",
        "tkerd = 0.0,",
        "tkern = 0.0,",
        "kmsl = 1,",
        "numpar = 5000,",
        "maxpar = 100000,",
        "ndump = 1,",
        "ncycl = 3,",
    ])
    run_hycs_std(hycs_exec, "Compute Anisotropy", inp_name)
    write_labels("Compute Anisotropy", script_label)
    run_concplot(concplot_exec, map_path, inp_name, "Compute Anisotropy",
                 data_txt, img_name_frame_nums=[("tker004.png", 1)])

    if DSP == "YES":
        input("Press Enter to continue...\n")

    # 3) Anisotropy = 0.05
    write_setup([
        "initd = 3,",
        "kblt = 3,",
        "kbls = 2,",
        "tkerd = 0.05,",
        "tkern = 0.05,",
        "kmsl = 1,",
        "numpar = 5000,",
        "maxpar = 100000,",
        "ndump = 1,",
        "ncycl = 3,",
    ])
    run_hycs_std(hycs_exec, "Anisotropy = 0.05", inp_name)
    write_labels("Anisotropy = 0.05", script_label)
    run_concplot(concplot_exec, map_path, inp_name, "Anisotropy = 0.05",
                 data_txt, img_name_frame_nums=[("tker006.png", 1)])


if __name__ == "__main__":
    main()