#!/usr/bin/env python3
#
# parm_mixd.py --- Mixing Depth options
#                  See HYSPLIT Tutorial Sec. 9.5
#
# 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_mixd.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):
    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_mixd(
    syr, smo, sda, shr,
    lat, lon, lvl,
    run_hours, top,
    met_dir: Path, dat_name: str,
    inp_name: str
):
    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")   # trailing sep is critical
        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_mixd.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_mixd(
        syr, smo, sda, shr,
        lat, lon, lvl,
        run_hours, top,
        met_dir, dat_name,
        inp_name,
    )

    # 1) Mixing depth from WRF (default)
    write_setup([
        "initd = 3,",
        "kmixd = 0,",
        "kmsl = 1,",
        "numpar = 5000,",
        "maxpar = 100000,",
        "ndump = 1,",
        "ncycl = 3,",
    ])
    run_hycs_std(hycs_exec, "mixing depth from WRF (default)", inp_name)
    write_labels("mixing depth from WRF (default)", script_label)
    run_concplot(concplot_exec, map_path, inp_name, "mixing depth from WRF (default)",
                 data_txt, img_name_frame_nums=[("case011.png", 1)])

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

    # 2) Mixing depth from temperature profile
    write_setup([
        "initd = 3,",
        "kmixd = 1,",
        "kmsl = 1,",
        "numpar = 5000,",
        "maxpar = 100000,",
        "ndump = 1,",
        "ncycl = 3,",
    ])
    run_hycs_std(hycs_exec, "mixing depth from temperature profile", inp_name)
    write_labels("mixing depth from temperature profile", script_label)
    run_concplot(concplot_exec, map_path, inp_name, "mixing depth from temperature profile",
                 data_txt, img_name_frame_nums=[("pbl003.png", 1)])

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

    # 3) Mixing depth from Richardson number
    write_setup([
        "initd = 3,",
        "kmsl = 1,",
        "numpar = 5000,",
        "maxpar = 100000,",
        "ndump = 1,",
        "ncycl = 3,",
    ])
    run_hycs_std(hycs_exec, "mixing depth from Richardson number", inp_name)
    write_labels("mixing depth from Richardson number", script_label)
    run_concplot(concplot_exec, map_path, inp_name, "mixing depth from Richardson number",
                 data_txt, img_name_frame_nums=[("pbl009.png", 1)])

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

    # 4) Fixed mixing depth = 5000 m
    write_setup([
        "initd = 3,",
        "kmsl = 1,",
        "kmixd = 5000,",
        "numpar = 5000,",
        "maxpar = 100000,",
        "ndump = 1,",
        "ncycl = 3,",
    ])
    run_hycs_std(hycs_exec, "mixing depth = 5000 m", inp_name)
    write_labels("mixing depth = 5000 m", script_label)
    run_concplot(concplot_exec, map_path, inp_name, "mixing depth = 5000 m",
                 data_txt, img_name_frame_nums=[("pbl005.png", 1)])

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

    # 5) Fixed mixing depth = 1250 m, turbulence off
    write_setup([
        "initd = 3,",
        "kmsl = 1,",
        "kmixd = 1250,",
        "numpar = 5000,",
        "maxpar = 100000,",
        "ndump = 1,",
        "ncycl = 3,",
        "tkerd = 0.0,",
        "tkern = 0.0,",
    ])
    run_hycs_std(hycs_exec, "mixing depth = 1250 m", inp_name)
    write_labels("mixing depth = 1250 m", script_label)
    run_concplot(concplot_exec, map_path, inp_name, "mixing depth = 1250 m",
                 data_txt, img_name_frame_nums=[("pbl007.png", 1)])


if __name__ == "__main__":
    main()