#!/usr/bin/env python3
#
# conc_puff.py --- Particle Distributions Using Puffs
#                  See HYSPLIT Tutorial Sec. 7.7.
#
# Change history:
# 03 Sep 2025 (BB) - v1.3: os.sep; lowercase vars; Windows backslashes in CONTROL;
#                    rename concplot.* -> disp_plot.*
# 04 Sep 2025 (SZ) - change the image filename.
#                  - use uppercase for variable names specified by the command-line
#                    arguments.
#                  - simplify some lines.
#                  - use magick.exe if available for Windows PC.
# 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)

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 "conc_puff.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)

# ----------------------------
# Switch into OUT 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")


# ---------- utilities ----------

def detect_met_path(dat_name: str) -> Path:
    """
    Try OUT/captex, TTR/captex, PGM/working/captex.
    Return path_obj.
    """
    candidates = [
        OUT / "captex",
        TTR / "captex",
        PGM / "working" / "captex",
    ]
    for c in candidates:
        if (c / dat_name).exists():
            return c

    # Fallback to OUT/captex even if file not found
    fallback = OUT / "captex"
    return fallback

# ---------- CONTROL/SETUP writers ----------


def write_control_blocks(met_dir: Path, dat_name: str):
    # CONTROL1
    with open("CONTROL1", "w", encoding="ascii") as f:
        f.write("83 09 25 17\n")  # SYR SMO SDA SHR
        f.write("1\n")  # number of sources
        f.write("39.90 -84.22 10.0\n")  # LAT LON LVL
        f.write("12\n")  # RUN
        f.write("0\n")  # vertical motion method
        f.write("10000.0\n")  # TOP
        f.write("1\n")  # number of met files
        f.write(f"{met_dir}{os.sep}\n")  # MET path (trailing slash)
        f.write(f"{dat_name}\n")  # DAT filename
        f.write("1\n")  # number of pollutants
        f.write("TEST\n")  # pollutant name
        f.write("100.0\n")  # emission rate
        f.write("0.01\n")  # emission duration
        f.write("00 00 00 00 00\n")  # release start offset
        f.write("1\n")  # receptors/grid flag
        f.write("0.0 0.0\n")  # dummy grid origin placeholders

    # CONTROL2
    with open("CONTROL2", "w", encoding="ascii") as f:
        f.write("30.0 30.0\n")  # domain size (km)
        f.write(f".{os.sep}\n")
        f.write("cdump\n")
        f.write("1\n")
        f.write("5000\n")
        f.write("00 00 00 00 00\n")
        f.write("00 00 00 00 00\n")
        f.write("01 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_setup(initd: int, numpar: int):
    with open("SETUP.CFG", "w", encoding="ascii") as f:
        f.write("&SETUP\n")
        f.write(f" initd = {initd},\n")
        f.write(f" numpar = {numpar},\n")
        f.write("/\n")


def run_model(label: str, grid_res: str, initd: int, numpar: int,
              concplot_extra: list=None, img_name_frame_nums=None,
              pause=False):
    print(f"\n>>> {label} <<<")

    # CONTROL = CONTROL1 + grid line + CONTROL2
    control1 = Path("CONTROL1").read_text(encoding="ascii")
    control2 = Path("CONTROL2").read_text(encoding="ascii")
    Path("CONTROL").write_text(control1 + f"{grid_res}\n" + control2,
                               encoding="ascii")

    write_setup(initd, numpar)

    if Path("cdump").exists():
        Path("cdump").unlink()

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

    with open("LABELS.CFG", "w", encoding="ascii") as f:
        f.write(f"'TITLE&','### {label} ### &'\n")

    concplot_cmd = [
        str(PGM / "exec" / f"concplot{exe_suffix}"),
        "+g1", "-icdump", "-z25", "-n1:4", "-j" + str(PGM / "graphics" / "arlmap")
    ]
    if concplot_extra:
        concplot_cmd.extend(concplot_extra)

    run_cmd(concplot_cmd, "Generating concentration plot")

    html_file = OUT / "concplot.html"
    if html_file.exists():
        if DSP == "YES":
            webbrowser.open(str(html_file.resolve()))
        else:
            convert_html_to_img(str(html_file), img_name_frame_nums)
    else:
        print("Warning: concplot.html not generated")

    if pause:
        input("Enter to continue ...")

# ---------- main ----------

dat = "captex2_wrf27uw.bin"
met_path = detect_met_path(dat)
write_control_blocks(met_path, dat)

# --------- 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")

# 1) Top-Hat
run_model("Top-Hat", "0.005 0.005", initd=4, numpar=1,
          img_name_frame_nums=[("cpuff009.png", 4)],
          pause=True if DSP == "YES" else False)

# 2) Gaussian
run_model("Gaussian", "0.005 0.005", initd=3, numpar=1,
          img_name_frame_nums=[("cpuff003.png", 2), ("cpuff010.png", 4)],
          pause=True if DSP == "YES" else False)

# 3) Gaussian low-res grid
run_model("Gaussian low-res grid", "0.05 0.05", initd=3, numpar=1,
          img_name_frame_nums=[("cpuff005.png", 2)],
          pause=True if DSP == "YES" else False)

# 4) Gaussian low-res more puffs
run_model("Gaussian low-res more puffs", "0.05 0.05", initd=3, numpar=100,
          concplot_extra=["-c1"],
          img_name_frame_nums=[("cpuff006.png", 4)])