#!/usr/bin/env python3
#
# cust_dust.py ---  Dust Storm Example with Custom Landuse
#                   See HYSPLIT Tutorial Sec. 14.2
#
# Change history:
# 24 Feb 2026 (BB) - standard-aligned, cross-platform version
# 15 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
#                  - correct convert_html_to_anim() for ImageMagick 7.x and later

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

OUT.mkdir(parents=True, exist_ok=True)
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 convert_html_to_anim(input_html, output_gif="output.gif", img_prefix="temp"):
    # Delete files matching the "F??-{img_prefix}.svg" pattern
    for file in OUT.glob(f"F??-{img_prefix}.svg"):
        file.unlink()

    run_cmd([str(PGM / "exec" / f"splitsvg{exe_suffix}"),
             f"-i{input_html}", f"-o{img_prefix}.svg"])

    # Locate ImageMagick
    convert_exe = None
    v7_or_later = False
    if system == "windows":
        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 earler versions.
            for cmd in ('magick', 'convert'):
                candidate = os.path.join(imgk_dir, f'{cmd}{exe_suffix}')
                if os.path.exists(candidate):
                    convert_exe = candidate
                    if cmd == 'magick':
                        v7_or_later = True
                    break
    else:
        convert_exe = "convert"

    if convert_exe:
        if v7_or_later:
            # place -trim and +repage after the input files for ImageMagick 7.x and later
            cmd = [
                convert_exe,
                "-density", "150",
                "-loop", "100",
                "-delay", "50",
                "-background", "white",
                "-dispose", "Background",
                f"F??-{img_prefix}.svg",
                "-trim",
                "+repage",
                output_gif
            ]
        else:
            cmd = [
                convert_exe,
                "-trim",
                "+repage",
                "-density", "150",
                "-loop", "100",
                "-delay", "50",
                "-background", "white",
                "-dispose", "Background",
                f"F??-{img_prefix}.svg",
                output_gif
            ]
        run_cmd(cmd)
        print(f"{output_gif} created successfully")
    else:
        print("ImageMagick not found — GIF not generated")


# constants
wrk = TTR
met = wrk / "dust"
dat = "WRF201003.bin"

syr = "10"
smo = "03"
sda = "30"
shr = "00"

run = "30"
top = "10000.0"

# input file checks
if not met.exists():
    print(f"ERROR: Dust tutorial directory not found: {met}")
    sys.exit(1)

if not (met / dat).exists():
    print(f"ERROR: Meteorology file not found: {met / dat}")
    sys.exit(1)

if not (met / "AirNow_dust.txt").exists():
    print(f"ERROR: AirNow comparison file not found: {met / 'AirNow_dust.txt'}")
    sys.exit(1)

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

# CONTROL
with open("CONTROL", "w", encoding="ascii") as f:
    f.write(f"{syr} {smo} {sda} {shr}\n")
    f.write("3\n")
    f.write("38.00 -115.00 10.0\n")
    f.write("41.00 -110.00 10.0\n")
    f.write("38.25 -114.75 10.0\n")
    f.write(f"{run}\n")
    f.write("0\n")
    f.write(f"{top}\n")
    f.write("1\n")
    f.write(f"{met}{os.sep}\n")
    f.write(f"{dat}\n")
    f.write("1\n")
    f.write("PM10\n")
    f.write("0.0\n")
    f.write(f"{run}\n")
    f.write("00 00 00 00 00\n")

    f.write("2\n")
    f.write("0.0 0.0\n")
    f.write("0.05 0.05\n")
    f.write("30.0 30.0\n")
    f.write(f".{os.sep}\n")
    f.write("dust0.bin\n")
    f.write("1\n")
    f.write("100\n")
    f.write("00 00 00 06 00\n")
    f.write("00 00 00 00 00\n")
    f.write("00 24 00\n")

    f.write("0.0 0.0\n")
    f.write("0.05 0.05\n")
    f.write("30.0 30.0\n")
    f.write(f".{os.sep}\n")
    f.write("dust0_byhr.bin\n")
    f.write("1\n")
    f.write("100\n")
    f.write("00 00 00 06 00\n")
    f.write("00 00 00 00 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")

# SETUP.CFG
with open("SETUP.CFG", "w", encoding="ascii") as f:
    f.write("&SETUP\n")
    f.write(" ichem = 3,\n")
    f.write(" numpar = 100000,\n")
    f.write(" maxpar = 50000,\n")
    f.write("/\n")

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

# run model
Path("dust0.bin").unlink(missing_ok=True)
Path("dust0_byhr.bin").unlink(missing_ok=True)

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

# 24-hour concentration plot
Path("concplot.html").unlink(missing_ok=True)

with open("LABELS.CFG", "w", encoding="ascii") as f:
    f.write("'TITLE&','### cust_dust: 0.5 degree landuse ### &'\n")

run_cmd([
    str(PGM / "exec" / f"concplot{exe_suffix}"),
    "+g1",
    "-idust0.bin",
    "-z90",
    f"-j{PGM / 'graphics' / 'arlmap'}",
    "-x1.0E+06",
    "-c4",
    "-uug",
    "-v500+200+100+50+20+10",
    "-g0:400",
    "-h40.00:-112.00"
], "Generating dust 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), [("dust009.png",1)])

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


# DATEM matching + stats
Path("hysplit.txt").unlink(missing_ok=True)

run_cmd([
    str(PGM / "exec" / f"c2datem{exe_suffix}"),
    "-xi",
    "-idust0.bin",
    "-c1.0E+06",
    f"-m{met / 'AirNow_dust.txt'}",
    "-ohysplit.txt"
], "Creating DATEM-format model file")

run_cmd([
    str(PGM / "exec" / f"statmain{exe_suffix}"),
    "-t0",
    "-rhysplit.txt",
    f"-d{met / 'AirNow_dust.txt'}",
    "-l10.0",
    "-o1"
], "Computing statistics")

stat_file = OUT / "statA.txt"
if stat_file.exists():
    print(stat_file.read_text(encoding="ascii", errors="replace"))

# scatter plot
Path("scatter.html").unlink(missing_ok=True)

run_cmd([
    str(PGM / "exec" / f"scatter{exe_suffix}"),
    "+g1",
    "-idataA.txt",
    "-p10.0"
], "Generating scatter plot")

html_file = OUT / "scatter.html"
if html_file.exists():
    if DSP == "YES":
        webbrowser.open(str(html_file.resolve()))
    else:
        convert_html_to_img(str(html_file), [("dust011.png",1)])

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

# hourly concentration plot
Path("dustbyhr.html").unlink(missing_ok=True)

run_cmd([
    str(PGM / "exec" / f"concplot{exe_suffix}"),
    "+g1",
    "-idust0_byhr.bin",
    "-z90",
    f"-j{PGM / 'graphics' / 'arlmap'}",
    "-x1.0E+06",
    "-c4",
    "-uug",
    "-v500+200+100+50+20+10",
    "-g0:400",
    "-h40.00:-112.00",
    "-odustbyhr"
], "Generating hourly dust concentration plot")

# display

html_file = OUT / "dustbyhr.html"
if html_file.exists():
    if DSP == "YES":
        webbrowser.open(str(html_file.resolve()))
    else:
        convert_html_to_anim(str(html_file), "dustbyhr.gif")