#!/usr/bin/env python3
#
# disp_dyn.py --- Dynamic or Lagrangian Sampling
#                 See HYSPLIT Tutorial Sec. 17.2
#
# Change history:
# 08 May 2026 (BB) - standard-aligned, cross-platform version
# 16 Jun 2026 (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

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"
    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", Path.home() / "tutorial"))
    exe_suffix = ""

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

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 find_met_dir(dat_name: str):
    env_override = os.environ.get("METDIR")

    candidates = []

    if env_override:
        candidates.append(Path(env_override))

    candidates += [
        wrk / "captex",
        TTR / "captex",
        TTR,
        PGM / "tutorial" / "captex",
        PGM / "tutorial",
        OUT / "captex",
        PGM / "met",
        PGM / "meteorology",
    ]

    for base in candidates:
        if base and (base / dat_name).exists():
            return base

    for root in [TTR, wrk, PGM]:
        try:
            for p in Path(root).rglob(dat_name):
                return p.parent
        except Exception:
            pass

    return None


# ---------- Mirror BAT variables ----------
wrk = OUT.parent

inp = "hysplit2.bin"

syr, smo, sda, shr = 83, 9, 25, 17
lat, lon, lvl = 39.90, -84.22, 10.0
run, top = 11, 10000.0

dat = "captex2_wrf27uw.bin"

met = find_met_dir(dat)
if met is None:
    print(f"ERROR: Meteorological file not found: {dat}")
    print("Set METDIR to the directory containing the file.")
    print("Example Windows: SET METDIR=C:\\Tutorial\\captex")
    sys.exit(1)

# Clean files
for fname in [
    "ASCDATA.CFG",
    "CONTROL",
    "SETUP.CFG",
    "LAGSET.CFG",
    inp,
    "LAGOUT.TXT",
]:
    p = OUT / fname
    if p.exists():
        p.unlink()

# ASCDATA.CFG
with open(OUT / "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")

# CONTROL
odir = f".{os.sep}"

with open(OUT / "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}\n")
    f.write("0\n")
    f.write(f"{top:.1f}\n")
    f.write("1\n")
    f.write(f"{met}{os.sep}\n")
    f.write(f"{dat}\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"{odir}\n")
    f.write(f"{inp}\n")
    f.write("2\n")
    f.write("800 1000\n")
    f.write("83 09 26 00 00\n")
    f.write("83 09 28 15 00\n")
    f.write("01 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")

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

# LAGSET.CFG
with open(OUT / "LAGSET.CFG", "w", encoding="ascii") as f:
    f.write("1\n")
    f.write("40.86 -81.81 914.0\n")
    f.write("290.0 50.0\n")
    f.write("83 09 26 02 48\n")
    f.write("83 09 26 02 48\n")
    f.write("06\n")
    f.write("06\n")
    f.write("'LAGOUT.TXT'\n")

# Run model
run_cmd(
    [str(PGM / "exec" / f"hycs_std{exe_suffix}")],
    "Running hycs_std"
)

# Display lagrangian output
lagout = OUT / "LAGOUT.TXT"

if lagout.exists():
    print("\n--- LAGOUT.TXT ---\n")
    print(lagout.read_text(encoding="ascii"))
else:
    print("Warning: LAGOUT.TXT not generated")

# Remove LAGSET.CFG
lagset = OUT / "LAGSET.CFG"

if lagset.exists():
    lagset.unlink()
    print("\nDeleted LAGSET.CFG")

print("\nScript completed successfully!")