#!/usr/bin/env python3
#
# src_geol.py --- Emissions from an Unknown Location
#                 See HYSPLIT Tutorial Sec. 13.3
#
# Change history:
# 18 Oct 2025 (BB) - standard-aligned, cross-platform version
# 12 Feb 2026 (SZ) - minor edits for style consistency
#                  - 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
# 01 Jul 2026 (SZ) - show html graphics after renaming the file

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


# -------------------------------------------------------------------
# paths and constants
# -------------------------------------------------------------------
wrk = TTR
met_dir = wrk / "captex"
met_file = "RP198309.gbl"

lat = 39.90
lon = -84.22
lvl = 10.0
top = 10000.0

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


def write_setup():
    with open("SETUP.CFG", "w", encoding="ascii") as f:
        f.write("&SETUP\n")
        f.write(" numpar = 5000,\n")
        f.write(" maxpar = 100000,\n")
        f.write("/\n")


def write_control(cdump_filename):
    with open("CONTROL", "w", encoding="ascii") as f:
        f.write("83 09 01 00\n")
        f.write("1\n")
        f.write(f"{lat} {lon} {lvl}\n")
        f.write("60\n")
        f.write("0\n")
        f.write(f"{top}\n")
        f.write("1\n")
        f.write(f"{met_dir}{os.sep}\n")
        f.write(f"{met_file}\n")
        f.write("1\n")
        f.write("PMCH\n")
        f.write("0.0\n")
        f.write("60.0\n")
        f.write("00 00 00 00 00\n")
        f.write("1\n")
        f.write("41.0 -73.0\n")
        f.write("0.25 0.25\n")
        f.write("15.0 25.0\n")
        f.write(f".{os.sep}\n")
        f.write(f"{cdump_filename}\n")
        f.write("1\n")
        f.write("100\n")
        f.write("83 09 01 00 00\n")
        f.write("83 09 03 12 00\n")
        f.write("00 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")


# -------------------------------------------------------------------
# loop over measurement (q="0") and inverse (q="1")
# -------------------------------------------------------------------
for q, inp, concplot_filename, img_name_frame_nums in [
        ("0", "cnum", "concplot_num.html", [("geol007.png",1)]),
        ("1", "cden", "concplot_inv.html", [("geol008.png",1)])]:

    # cleanup
    for f in Path(".").glob("CONTROL.*"):
        f.unlink()
    for f in Path(".").glob(f"{inp}.*"):
        f.unlink()

    write_control(inp)

    # -------------------------------------------------------------------
    # dat2cntl (creates CONTROL.###)
    # -------------------------------------------------------------------
    run_cmd([
        str(PGM / "exec" / "dat2cntl"),
        f"-i{met_dir}{os.sep}hypo_meas.txt",
        f"-d{q}",
        "-c1.0"
    ], "Running dat2cntl")

    # -------------------------------------------------------------------
    # loop over CONTROL.###
    # -------------------------------------------------------------------
    for run in range(1, 1000):
        num = f"{run:03d}"
        ctl = Path(f"CONTROL.{num}")

        write_setup()

        if ctl.exists():
            print(num)
            Path("SETUP.CFG").replace(f"SETUP.{num}")
            run_cmd([
                str(PGM / "exec" / "hycs_std"),
                num
            ])
            for pat in (
                f"MESSAGE.{num}",
                f"VMSDIST.{num}",
                f"WARNING.{num}",
                f"CONTROL.{num}",
                f"CONC.{num}",
                f"SETUP.{num}",
            ):
                Path(pat).unlink(missing_ok=True)

    # -------------------------------------------------------------------
    # conprob
    # -------------------------------------------------------------------
    Path("cmean").unlink(missing_ok=True)

    run_cmd([
        str(PGM / "exec" / "conprob"),
        f"-b{inp}",
        "-t20",
        "-v0.0"
    ], "Running conprob")

    # -------------------------------------------------------------------
    # LABELS.CFG
    # -------------------------------------------------------------------
    with open("LABELS.CFG", "w", encoding="ascii") as f:
        if q == "0":
            f.write("'TITLE&','### conc_captex_inverse: Source Sensitivity ### &'\n")
        else:
            f.write("'TITLE&','### conc_captex_inverse: Inverse Emissions ### &'\n")

        f.write("'UNITS&',' mass &'\n")
        f.write("'VOLUM&',' &'\n")
        f.write("'MAPID&',' Receptor &'\n")

    # -------------------------------------------------------------------
    # concplot
    # -------------------------------------------------------------------
    Path("concplot.html").unlink(missing_ok=True)

    run_cmd([
        str(PGM / "exec" / "concplot"),
        "+g1",
        "-m0",
        "-k1",
        "-z80",
        f"-j{PGM}{os.sep}graphics{os.sep}arlmap",
        "-icmean"
    ], "Generating concentration plot")

    html = Path("concplot.html")
    if html.exists():
        concplot_path = Path(concplot_filename)
        if concplot_path.exists():
            concplot_path.unlink()
        html.rename(concplot_path)
        html = concplot_path

        if DSP.upper() == "YES":
            try:
                webbrowser.open(html.resolve().as_uri())
            except Exception as e:
                print(f"WARNING: could not open {html} in web browser: {e}")
        else:
            convert_html_to_img(str(html), img_name_frame_nums)

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

#----------------------------------------------------------
# Run dispersion at 43N 75W

syr=83
smo=9
sda=1
shr=0

olat=43.00
olon=-75.00
lvl1=10.0

run=60
ztop=10000.0

hypo_conc_file = Path("hypo_conc.bin")

with open("SETUP.CFG", "w", encoding="ascii") as f:
    f.write("&SETUP\n")
    f.write(" initd = 0,\n")
    f.write(" khmax = 24,\n")
    f.write(" numpar = 240000,\n")
    f.write(" maxpar = 120000,\n")
    f.write("/\n")

with open("CONTROL", "w", encoding="ascii") as f:
    f.write(f"{syr} {smo:02d} {sda:02d} {shr:02d}\n")
    f.write("1                      \n")
    f.write(f"{olat} {olon} {lvl1}  \n")
    f.write(f"{run}                 \n")
    f.write("0                      \n")
    f.write(f"{ztop}                \n")
    f.write("1                      \n")
    f.write(f"{met_dir}{os.sep}     \n")
    f.write(f"{met_file}            \n")
    f.write("1                      \n")
    f.write("PMCH                   \n")
    f.write("3000.0                 \n")
    f.write("60.0                   \n")
    f.write("00 00 00 00 00         \n")
    f.write("1                      \n")
    f.write("41.0 -73.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"{hypo_conc_file}      \n")
    f.write("1                      \n")
    f.write("100                    \n")
    f.write("00 00 00 00 00         \n")
    f.write("00 00 00 00 00         \n")
    f.write("00 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")

# # Delete the file hypo_conc.bin if it exists
if hypo_conc_file.exists():
    hypo_conc_file.unlink()

run_cmd([
    str(PGM / "exec" / "hycs_std")
], "Running HYSPLIT")

run_cmd([
    str(PGM / "exec" / "conavgpd"),
    f"-i{hypo_conc_file}",
    "-ocdump",
    "-h60",
], "Running CONAVGPD")

with open("LABELS.CFG", "w", encoding="ascii") as f:
    f.write("'TITLE&','### src_geol: dispersion at 43N 75W ### &'\n")
    f.write("'UNITS&',' pg &'\n")

Path("concplot.html").unlink(missing_ok=True)

run_cmd([
    str(PGM / "exec" / "concplot"),
    "+g1",
    "-icdump",
    f"-j{PGM}{os.sep}graphics{os.sep}arlmap",
    "-z100",
    "-x1.0E+12",
    "-H43.0:-70.0",
    "-c4",
    "-v20000+10000+5000+2000+1000+500+200+100",
], "Generating concentration plot")

html = Path("concplot.html")
if html.exists():
    if DSP.upper() == "YES":
        try:
            webbrowser.open(html.resolve().as_uri())
        except Exception as e:
            print(f"WARNING: could not open {html} in web browser: {e}")
    else:
        convert_html_to_img(str(html), [("geol009.png",1)])