#!/usr/bin/env python3
#
# disp_maps.py --- County Map Boundaries
#                  See HYSPLIT Tutorial Sec. 10.2
#
# Change history:
# 23 Oct 2025 (BB) - standard-aligned, cross-platform version
# 19 Nov 2025 (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
import webbrowser

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

DSP = str(DSP).upper()

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

For Linux users (csh or tcsh shell):
   setenv TTR 'YOUR_HYSPLIT_TOTURIAL_DIRECTORY'""")
    sys.exit(1)

# Operate from working
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(path: Path):
    path.unlink(missing_ok=True)


def write_ascdata(pgm: Path):
    clean(Path("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")  # trailing sep matters


def write_control(
    syr, smo, sda, shr,
    lat, lon, lvl,
    run_hours, top,
    met_dir: Path, dat_name: str,
    inp_name: str
):
    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_hours}\n")
        f.write("0\n")
        f.write(f"{top:.1f}\n")
        f.write("1\n")
        f.write(f"{str(met_dir)}{os.sep}\n")  # trailing sep critical
        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.05 0.05\n")
        f.write("15.0 25.0\n")
        f.write(f".{os.sep}\n")
        f.write(f"{inp_name}\n")
        f.write("2\n")
        f.write("800 1000\n")
        f.write("83 09 26 03 00\n")
        f.write("83 09 28 15 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")


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


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


def run_concplot(map_arg: str, inp_name: str, qfile: Path, label: str,
                 img_name_frame_nums=None):
    clean(Path("concplot.html"))
    write_labels(Path(__file__).name if "__file__" in globals() else "disp_maps.py")
    concplot = PGM / "exec" / f"concplot{exe_suffix}"
    if not concplot.exists():
        print(f"ERROR: concplot not found at {concplot}")
        sys.exit(1)
    cmd = [
        str(concplot),
        "+g1",
        f"-i{inp_name}",
        "-z100",
        f"-j{map_arg}",
        "-x1.0E+12",
        "-upg",
        "-h41.0:-82.0",
        "-g0:200",
        "-c4",
        "-v50000+20000+10000+5000+2000+1000+500",
        "-b1000",
        "-t1000",
        f"-q{qfile}",
    ]
    run_cmd(cmd, f"Create concplot ({label})")
    html = Path("concplot.html")
    if html.exists():
        if DSP == "YES":
            try:
                webbrowser.open(html.resolve().as_uri())
            except Exception as e:
                print(f"WARNING: could not open concplot.html in browser: {e}")
        else:
            convert_html_to_img(str(html), img_name_frame_nums)

# ---------------- Main flow ----------------
def main():
    # Batch constants / defaults
    MAP = str(PGM / "graphics" / "arlmap")     # used as -j argument
    INP = "hysplit2.bin"

    # ASCDATA
    write_ascdata(PGM)

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

    # MET defaults to TTR\captex (batch)
    met_dir = TTR / "captex"
    if not met_dir.exists():
        print(f"WARNING: meteorology directory not found at TTR: {met_dir}")
    dat_name = "captex2_wrf27uw.bin"
    if not (met_dir / dat_name).exists():
        print(f"WARNING: met file not found: {met_dir / dat_name}")

    # Write CONTROL & SETUP exactly like the batch
    write_control(syr, smo, sda, shr, lat, lon, lvl, run_hours, top,
                  met_dir, dat_name, INP)
    write_setup()

    # Fresh run (delete INP then hycs_std)
    clean(Path(INP))
    hycs = PGM / "exec" / f"hycs_std{exe_suffix}"
    if not hycs.exists():
        print(f"ERROR: hycs_std not found at {hycs}")
        sys.exit(1)
    run_cmd([str(hycs)], "Running hycs_std")

    # First concplot with ARLMAP
    qfile = met_dir / "data_case.txt" \
        if (met_dir / "data_case.txt").exists() \
        else (TTR / "captex" / "data_case.txt")
    if not qfile.exists():
        print(f"WARNING: data_case.txt not found in {met_dir} or {TTR / 'captex'}; "
              "concplot will likely miss obs.")
    run_concplot(MAP, INP, qfile, "arlmap",
                 img_name_frame_nums=[("case011.png", 1)])

    # ---- PAUSE equivalent ----
    if DSP == "YES":
        input("Press Enter to continue to shapefiles overlay...\n")

    # Build shapefiles.txt from TTR\maps
    maps_dir = TTR / "maps"
    roads = maps_dir / "roads.shp"
    county = maps_dir / "county.shp"
    states = maps_dir / "states.shp"

    with open("shapefiles.txt", "w", encoding="ascii") as f:
        f.write(f"'{roads}'  0 0.005 0.8 0.7 0.5\n")
        f.write(f"'{county}' 0 0.010 0.4 0.6 0.8\n")
        f.write(f"'{states}' 0 0.015 0.2 0.2 0.5\n")

    for shp in [roads, county, states]:
        if not shp.exists():
            print(f"WARNING: missing shapefile: {shp}")

    # Second concplot with shapefiles.txt as -j
    run_concplot("shapefiles.txt", INP, qfile, "shapefiles",
                 img_name_frame_nums=[("maps002.png", 1)])

if __name__ == "__main__":
    main()