#!/usr/bin/env python3
#
# disp_python.py — Python Display Options
#                  See HYSPLIT Tutorial Sec. 10.7
#
# Change history:
# 05 Nov 2025 (BB) - standard-aligned, cross-platform version
# 21 Nov 2025 (SZ) - code refactoring
#                  - 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
import subprocess
import webbrowser
from pathlib import Path

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

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").upper()
    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").upper()
    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 "disp_python.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,
            env=None):
    if desc:
        print(f"{desc}...")

    exitcode = 0
    with subprocess.Popen(cmd,
                         stdout=subprocess.PIPE,
                         stderr=subprocess.STDOUT,
                         text=True,
                         env=env) 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):
    if not Path("ASCDATA.CFG").exists():
        with open("ASCDATA.CFG", "w", encoding="ascii") as f:
            f.write("-90.0  -180.0\n1.0     1.0\n180     360\n2\n0.2\n")
            f.write(f"'{pgm /'bdyfiles'}{os.sep}'  directory of files\n")


def move_first_frame(stem: str):
    f = Path(f"{stem}0001.png")
    if f.exists():
        f.rename(f"{stem}.png")


# Build env that includes PGM/exec on PYTHONPATH for plot scripts
PY = sys.executable
EXEC_DIR = str(PGM / "exec")
def with_plot_env():
    env = os.environ.copy()
    env["PYTHONPATH"] = (EXEC_DIR + os.pathsep + env.get("PYTHONPATH","")).strip(os.pathsep)
    return env

# ---------------- Python trajplot demos ----------------
cmtraj = Path("cmtraj.txt")
trajplot_py = PGM / "exec" / "trajplot.py"
concplot_py = PGM / "exec" / "concplot.py"
MAP = str(PGM / "graphics" / "arlmap")

if not cmtraj.exists():
    print("cmtraj.txt file from Sec. 7.6 is required.  Stopping.")
    sys.exit(1)
else:
    if not trajplot_py.exists():
        print(f"NOTE: {trajplot_py} not found; skipping Python trajplot demos.")
    else:
        env = with_plot_env()
        # no street map
        run_cmd([PY, str(trajplot_py), "-i"+str(cmtraj), "-j"+MAP, "-l1",
                "-o"+"cmtraj006.png"],
                "trajplot.py (no street-map) -> cmtraj006.png", env=env, allow_fail=True)
        move_first_frame("cmtraj006")
        # street-map=0
        run_cmd([PY, str(trajplot_py), "-i"+str(cmtraj), "-j"+MAP, "-l1",
                "--street-map=0", "-o"+"cmtraj010.png"],
                "trajplot.py (--street-map=0) -> cmtraj010.png", env=env, allow_fail=True)
        move_first_frame("cmtraj010")
        # street-map=1
        run_cmd([PY, str(trajplot_py), "-i"+str(cmtraj), "-j"+MAP, "-l1",
                "--street-map=1", "-o"+"cmtraj009.png"],
                "trajplot.py (--street-map=1) -> cmtraj009.png", env=env, allow_fail=True)
        move_first_frame("cmtraj009")

# ---------------- Baseline dispersion (HTML concplot) ----------------
write_ascdata(PGM)

syr,smo,sda,shr = 83,9,25,17
olat,olon,lvl1 = 39.90,-84.22,10.0
run_hours = 11
ztop = 10000.0
inp = "hysplit2.bin"

met_dir = TTR / "captex"
if not met_dir.exists():
    print(f"WARNING: met dir not found: {met_dir}")
dat_name = "captex2_wrf27uw.bin"
if not (met_dir / dat_name).exists():
    print(f"WARNING: met dir not found: {met_dir / dat_name}")

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"{olat:.2f} {olon:.2f} {lvl1:.1f}\n")
    f.write(f"{run_hours}\n")
    f.write("0\n")
    f.write(f"{ztop:.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.00\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}\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")

with open("SETUP.CFG", "w", encoding="ascii") as f:
    f.write("&SETUP\n")
    f.write("initd=3,\n")
    f.write("kmsl=1,\n")
    f.write("numpar=5000,\n")
    f.write("maxpar=100000,\n")
    f.write("ndump=1,\n")
    f.write("ncycl=3,\n")
    f.write("/\n")

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

clean(inp)
hycs = exe("hycs_std")
if not hycs.exists():
    print(f"ERROR: hycs_std not found at {hycs}")
    sys.exit(1)
run_cmd([str(hycs)], "Running hycs_std")

concplot_exe = exe("concplot")
if not concplot_exe.exists():
    print(f"ERROR: concplot not found at {concplot_exe}")
    sys.exit(1)

clean("concplot.html")
with open("LABELS.CFG", "w", encoding="ascii") as f:
    fn = Path(__file__).name if '__file__' in globals() else 'disp_python.py'
    f.write(f"'TITLE&','### {fn} ### at flight level &'\n")
run_cmd([
    str(concplot_exe), "+g1", "-i"+inp, "-j"+MAP, "-z100", "-x1.0E+12", "-upg",
    "-h41.0:-82.0", "-g0:200", "-c4",
    "-v50000+20000+10000+5000+2000+1000+500",
    "-b1000", "-t1000", "-q"+str(met_dir/"data_case.txt")
], "Create concplot (HTML)")

# PNG export if DSP != YES
html = Path("concplot.html")
if html.exists():
    if DSP == "YES":
        try:
            webbrowser.open(html.resolve().as_uri())
        except Exception:
            pass
    else:
        convert_html_to_img(str(html), [("pyth002.png", 1)])

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

# Python concplot street-map -> pyth001.png (optional)
if concplot_py.exists():
    env = with_plot_env()
    run_cmd([PY, str(concplot_py),
             "-i"+inp, "-j"+MAP, "-z100", "-x1.0E+12", "-upg",
             "-h41.0:-82.0", "-g0:200", "-c4",
             "-v50000+20000+10000+5000+2000+1000+500",
             "-b1000", "-t1000", "-q"+str(met_dir/"data_case.txt"),
             "--street-map=1", "-o"+"pyth001.png"],
            "concplot.py (street-map=1) -> pyth001.png", env=env, allow_fail=True)
    move_first_frame("pyth001")
else:
    print(f"NOTE: {concplot_py} not found; skipping Python concplot demo.")

print("All done.")