#!/usr/bin/env python3
#
# chem_linear.py --- Linear Mass Conversion (PMCH --> POL2)
#                    See HYSPLIT Tutorial Sec. 11.1
#
# Change history:
# 08 Nov 2025 (BB) - standard-aligned, cross-platform version
# 11 Dec 2025 (SZ) - add help text; port to Linux; create images 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

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

# ---------------- OS / args ----------------
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 "chem_linear.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")


def clean(p):
    Path(p).unlink(missing_ok=True)


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


def write_ascdata(pgm: Path):
    clean("ASCDATA.CFG")
    with open("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")


# Common CAPTEX setup
syr, smo, sda, shr = 83, 9, 25, 17
lat, lon, lvl = 39.90, -84.22, 10.0
run_hrs, top = 25, 10000.0

met_dir = TTR / "captex"
dat_name = "captex2_wrf27uw.bin"
map_path = PGM / "graphics" / "arlmap"
inp_name = "hysplit2.bin"

# ------------- Base run (single species PMCH) -------------
write_ascdata(PGM)

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"{lat:.2f} {lon:.2f} {lvl:.1f}\n")
    f.write(f"{run_hrs}\n")
    f.write("0\n")
    f.write(f"{top:.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.0\n")
    f.write("00 00 00 00 00\n")
    f.write("1\n")
    f.write("42.0 -78.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"{inp_name}\n")
    f.write("1\n")
    f.write("100\n")
    f.write("83 09 25 18 00\n")
    f.write("83 09 28 15 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")

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

clean(inp_name)
run_cmd([exe("hycs_std")], "Running hycs_std (base)")

# Safe rename (overwrite if exists)
msg = Path("MESSAGE")
msg_base = Path("MESSAGE_base.txt")
if msg.exists():
    if msg_base.exists():
        msg_base.unlink()
    msg.rename(msg_base)

# plot_base (species index 1)
clean("plot_base.html")
with open("LABELS.CFG", "w", encoding="ascii") as f:
    f.write("'TITLE&','### chem_linear.py ### &'\n")
run_cmd([exe("concplot"), "+g1", f"-i{inp_name}", "-oplot_base",
         "-z80", "-s1", "-x1.0E+12", "-upg", f"-j{map_path}"], "Create plot_base")
html_file = OUT / "plot_base.html"
if html_file.exists():
    if DSP == "YES":
        webbrowser.open(html_file.resolve().as_uri())
        input("Press Enter to continue...")
    else:
        convert_html_to_img("plot_base.html", [("dry000.png", 8)])
else:
    print('WARNING: plot_base.html not generated')

# ------------- Linear chemistry run (PMCH + POL2, ichem=2) -------------
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"{lat:.2f} {lon:.2f} {lvl:.1f}\n")
    f.write(f"{run_hrs}\n")
    f.write("0\n")
    f.write(f"{top:.1f}\n")
    f.write("1\n")
    f.write(f"{str(met_dir)}{os.sep}\n")
    f.write(f"{dat_name}\n")
    f.write("2\n")
    # species #1
    f.write("PMCH\n")
    f.write("67000.0\n")
    f.write("3.0\n")
    f.write("00 00 00 00 00\n")
    # species #2 (produced)
    f.write("POL2\n")
    f.write("0.0\n")
    f.write("0.0\n")
    f.write("00 00 00 00 00\n")
    # grid & output
    f.write("1\n")
    f.write("42.0 -78.0\n")
    f.write("0.25 0.25\n")
    f.write("15.0 25.0\n")
    f.write(f"{OUT}{os.sep}\n")  # correct output dir line
    f.write(f"{inp_name}\n")
    f.write("1\n")
    f.write("100\n")
    f.write("83 09 25 18 00\n")
    f.write("83 09 28 15 00\n")
    f.write("00 03 00\n")
    # deposition parameter blocks for BOTH species
    f.write("2\n")
    # species 1 dep params
    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")
    # species 2 dep params
    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("ichem = 2,\n")  # linear chemistry
    f.write("maxdim = 2,\n")  # 2 species
    f.write("numpar = 20000,\n")
    f.write("maxpar = 100000,\n")
    f.write("/\n")

# Optional: override conversion via CHEMRATE.TXT
# with open("CHEMRATE.TXT","w") as f:
#     f.write("1 2 0.10 1.0\n")  # 10%/hr from 1->2

clean(inp_name)
run_cmd([exe("hycs_std")], "Running hycs_std (linear chemistry)")

# plot species 1 and 2
for sidx, outname, imgname in [(1, "plot_chem1", 'chem005.png'),
                               (2, "plot_chem2", 'chem005b.png')]:
    clean(f"{outname}.html")
    with open("LABELS.CFG", "w", encoding="ascii") as f:
        f.write("'TITLE&','### chem_linear.py ### &'\n")
    run_cmd([exe("concplot"), "+g1", "-ihysplit2.bin", f"-o{outname}",
             "-z80", f"-s{sidx}", "-x1.0E+12", "-upg", f"-j{map_path}"],
            f"Create {outname}")
    html_file = OUT / f"{outname}.html"
    if html_file.exists():
        if DSP == "YES":
            webbrowser.open(html_file.resolve().as_uri())
            input("Press Enter to continue...")
        else:
            convert_html_to_img(str(html_file), [(imgname, 1)])
    else:
        print(f'WARNING: {outname}.html not generated')

print("chem_linear.py DONE.")