#!/usr/bin/env python3
#
# parm_data.py --- Base Configuration Optimization
#                  See HYSPLIT Tutorial Sec. 9.2
#
# Change history:
# 15 Oct 2025 (BB) - standard-aligned, cross-platform version
# 12 Nov 2025 (BB) - fix met_dir to TTR/captex; robust paths; minor cleanups
# 19 Nov 2025 (SZ) - create image files when DSP=NO
# 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 (certified header)
# ----------------------------
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_data.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)

# Change to working directory
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 as in the batch script (with trailing sep)."""
    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_base(
    syr, smo, sda, shr, lat, lon, lvl, run_hours, top, met_dir: Path, dat_name: str, inp_name: str
):
    """CONTROL for the first three runs (3D particles, Top-Hat, Gaussian)."""
    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")
        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("1\n")
        f.write("100\n")
        f.write("83 09 26 03 00\n")
        f.write("83 09 28 15 00\n")
        f.write("00 03 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_control_flight_level(
    syr, smo, sda, shr, lat, lon, lvl, run_hours, top, met_dir: Path, dat_name: str, inp_name: str
):
    """CONTROL for the final 'at flight level' run."""
    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")
        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 26 04 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):
    """Write SETUP.CFG from a list of body 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):
    clean("hysplit2.bin")
    run_cmd([str(hycs_exec)], f"Run hycs_std for {label}")


def write_labels(label: str, script_label: str):
    """Create LABELS.CFG similar to the batch version."""
    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, extra_args=None,
                 img_name_frame_nums=None):
    if extra_args is None:
        extra_args = []

    clean("concplot.html")

    base_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",
        f"-q{data_txt}",
    ]
    base_cmd.extend(extra_args)

    run_cmd(base_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)


def main():
    script_label = Path(__file__).name if "__file__" in globals() else "parm_data.py"

    # Paths / executables
    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)

    # Recreate ASCDATA.CFG
    write_ascdata(PGM)

    # CAPTEX configuration (NOTE: met_dir must be under TTR/captex)
    syr, smo, sda, shr = 83, 9, 25, 17
    lat, lon, lvl = 39.90, -84.22, 10.0
    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}")

    # -------- Case 1: 3D particles --------
    write_control_base(syr, smo, sda, shr, lat, lon, lvl, 13, top, met_dir, dat_name, inp_name)
    write_setup([
        "numpar = 50000,",
        "maxpar = 100000,",
        "ndump = 1,",
        "ncycl = 3,",
    ])
    run_hycs_std(hycs_exec, "3D particles")
    write_labels("3D particles", script_label)
    run_concplot(concplot_exec, map_path, inp_name, "3D particles", data_txt=data_txt,
                 img_name_frame_nums=[("case007.png", 1)])

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

    # -------- Case 2: Top-Hat Puff --------
    write_control_base(syr, smo, sda, shr, lat, lon, lvl, 13, top, met_dir, dat_name, inp_name)
    write_setup([
        "initd = 4,",
        "numpar = 5000,",
        "maxpar = 100000,",
        "ndump = 1,",
        "ncycl = 3,",
    ])
    run_hycs_std(hycs_exec, "Top-Hat Puff")
    write_labels("Top-Hat Puff", script_label)
    run_concplot(concplot_exec, map_path, inp_name, "Top-Hat Puff", data_txt=data_txt,
                 img_name_frame_nums=[("case008.png", 1)])

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

    # -------- Case 3: Gaussian Puff --------
    write_control_base(syr, smo, sda, shr, lat, lon, lvl, 13, top, met_dir, dat_name, inp_name)
    write_setup([
        "initd = 3,",
        "numpar = 5000,",
        "maxpar = 100000,",
        "ndump = 1,",
        "ncycl = 3,",
    ])
    run_hycs_std(hycs_exec, "Gaussian Puff")
    write_labels("Gaussian Puff", script_label)
    run_concplot(concplot_exec, map_path, inp_name, "Gaussian Puff", data_txt=data_txt,
                 img_name_frame_nums=[("case009.png", 1)])

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

    # -------- Case 4: At flight level --------
    write_control_flight_level(syr, smo, sda, shr, lat, lon, lvl, 11, top,
                               met_dir, dat_name, inp_name)
    write_setup([
        "initd = 3,",
        "kmsl  = 1,",
        "numpar = 5000,",
        "maxpar = 100000,",
        "ndump = 1,",
        "ncycl = 3,",
    ])
    run_hycs_std(hycs_exec, "at flight level")
    write_labels("at flight level", script_label)
    run_concplot(
        concplot_exec,
        map_path,
        inp_name,
        "at flight level",
        data_txt=data_txt,
        extra_args=["-b1000", "-t1000"],
        img_name_frame_nums=[("case011.png", 1)]
    )


if __name__ == "__main__":
    main()