#!/usr/bin/env python3
#
# traj_vert.py --- run HYSPLIT vertical trajectory analysis
#                  See HYSPLIT Tutorial Sec. 5.1
#
# Change history:
# 18 Jul 2025 (BB) - standard-aligned, cross-platform version
# 06 Aug 2025 (BB) - added PNG fallback conversion
# 08 Aug 2025 (SZ) - update for Linux
# 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

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

# ----------------------------
# OS + command-line 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 = ""

if not PGM.exists():
    script = Path(__file__).name if "__file__" in globals() else "traj_vert.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_TOTURIAL_DIRECTORY'

For Linux users (csh or tcsh shell):
   setenv TTR 'YOUR_HYSPLIT_TOTURIAL_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")


# ----------------------------
# Parameters
# ----------------------------
syr, smo, sda, shr = 83, 9, 28, 13
lat, lon, lvl = 38.665, -71.369, 216.4
run, top = -68, 10000.0
met = TTR / "captex"
dat = "captex2_wrf27uw.bin"


def exe(name):
    return str(PGM / "exec" / f"{name}{exe_suffix}")


# ----------------------------
# CONTROL writing function
# ----------------------------
def write_control(control1_lines, mode, tdump_file):
    ctl = OUT / "CONTROL"
    control2_lines = [
        f"{top}",
        "1",
        f"{met}{os.sep}",
        dat,
        f".{os.sep}",
        tdump_file
    ]
    with ctl.open("w", encoding="ascii") as f:
        for line in control1_lines:
            f.write(f"{line}\n")
        f.write(f"{mode}\n")
        for line in control2_lines:
            f.write(f"{line}\n")

    print(f"\nCONTROL file ({tdump_file}):\n{ctl.read_text()}")


# ----------------------------
# Main
# ----------------------------
def main():
    # Clean up old files
    cleanup_files = [
        "ASCDATA.CFG", "SETUP.CFG", "LABELS.CFG", "trajplot.html",
        "tdump_data", "tdump_isob", "tdump_isen", "tdump_dens",
        "tdump_sigma", "tdump_divg", "traj_files.txt",
        "STATIONPLOT.CFG"
    ]
    for f in cleanup_files:
        p = OUT / f
        if p.exists():
            p.unlink()

    # ASCDATA.CFG
    asc = OUT / "ASCDATA.CFG"
    with asc.open("w", encoding="ascii") as f:
        f.write("-90.0   -180.0\n")
        f.write("1.0     1.0\n")
        f.write("180     360\n")
        f.write("2\n")
        f.write("0.2\n")
        f.write(f"'{PGM / 'bdyfiles'}{os.sep}'\n")

    # CONTROL1
    control1 = [
        f"{syr:02d} {smo:02d} {sda:02d} {shr:02d}",
        "1",
        f"{lat} {lon} {lvl}",
        f"{run}"
    ]

    # Modes
    modes = [
        (0, "tdump_data"),
        (1, "tdump_isob"),
        (2, "tdump_isen"),
        (3, "tdump_dens"),
        (4, "tdump_sigma"),
        (5, "tdump_divg")
    ]

    # Run HYTS_STD for each mode
    traj_files = []
    for mode, fname in modes:
        write_control(control1, mode, fname)
        traj_files.append(fname)
        run_cmd([exe("hyts_std")], f"Running hyts_std mode {mode}")

    # List output files
    (OUT / "traj_files.txt").write_text("\n".join(traj_files), encoding="ascii")

    # LABELS
    (OUT / "LABELS.CFG").write_text("'TITLE&','### traj_vert.py ### &'\n",
                                    encoding="ascii")

    # STATIONPLOT
    (OUT / "STATIONPLOT.CFG").write_text(
        "39.420 -84.155 divg\n"
        "38.810 -83.000 data\n"
        "39.179 -75.200 isen\n"
        "38.350 -74.877 dens\n"
        "39.800 -72.420 sigma\n"
        "38.900 -72.889 isob"
    )

    # TRajplot
    run_cmd([
        exe("trajplot"),
        "+g1",
        "-i+traj_files.txt",
        "-v0",
        "-z30",
        f"-j{PGM / 'graphics' / 'arlmap'}"
    ], "Generating trajplot")

    html = OUT / "trajplot.html"

    # ----------------------------
    # PNG Fallback Logic
    # ----------------------------
    if html.exists():
        print("trajplot.html generated")

        if DSP.upper() == "YES":
            webbrowser.open(str(html))
        else:
            convert_html_to_img(str(html), [("vert004.png", 1)])
    else:
        print("trajplot.html was not generated.")

    print("\nScript completed successfully!")


if __name__ == "__main__":
    main()