#!/usr/bin/env python3
#
# disp_merg.py --- Merge dispersion plume and trajectories
#                  See HYSPLIT Tutorial Sec. 10.6
#
# Change history:
# 03 Nov 2025 (BB) - initial cross-platform port
# 12 Nov 2025 (BB) - add safe shapefile handling to prevent concplot 899
# 21 Nov 2025 (SZ) - code refactoring
#                  - 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 glob
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"  # not used here; kept for parity
    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 = ""

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

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(path).unlink(missing_ok=True)

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")
        # trailing separator required
        f.write(f"'{pgm /'bdyfiles'}{os.sep}'  directory of files\n")

def write_setup_disp():
    clean("SETUP.CFG")
    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("initd = 3,\n")
        f.write("kmsl = 1,\n")
        f.write("/\n")

def write_control_disp(
    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_control_traj(
    syr, smo, sda, shr,
    lat, lon, lvl_m_agl,
    run_hours, top,
    met_dir: Path, dat_name: str,
    tdump_name: str
):
    # Minimal CONTROL for trajectories (as in batch)
    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_m_agl:.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(f".{os.sep}\n")
        f.write(f"{tdump_name}\n")

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

def create_shapefiles_txt(paths: list) -> str:
    """
    Create shapefiles.txt with only the shapefiles that actually exist.
    If none exist, return the built-in arlmap path so -j points there instead.
    """
    entries = []
    for path, color_info in paths:
        if path.exists():
            entries.append(f"'{path}' {color_info}")
    if entries:
        with open("shapefiles.txt", "w", encoding="ascii") as f:
            f.write("\n".join(entries) + "\n")
        return "shapefiles.txt"

    print("NOTE: No background shapefiles found; falling back to arlmap.")
    return str(PGM / "graphics" / "arlmap")


# ------------------- Main flow -------------------
def main():
    inp = "hysplit2.bin"
    tdump = "tdump_fwrd"

    # Background shapefiles as in the batch (roads/county/states)
    shp_roads  = TTR / "maps" / "roads.shp"
    shp_county = TTR / "maps" / "county.shp"
    shp_states = TTR / "maps" / "states.shp"

    # 0) ASCDATA
    write_ascdata(PGM)

    # 1) Decide background for first plots
    j_bg = create_shapefiles_txt([
        (shp_roads,  '0 0.005 0.8 0.7 0.5'),
        (shp_county, '0 0.010 0.4 0.6 0.8'),
        (shp_states, '0 0.015 0.2 0.2 0.5')
    ])

    # 2) Dispersion run CONTROL/SETUP (batch values)
    syr, smo, sda, shr = 83, 9, 25, 17
    lat, lon, lvl = 39.90, -84.22, 10.0
    run_hours = 11
    top = 10000.0

    # MET = TTR\captex
    met_dir = TTR / "captex"
    if not met_dir.exists():
        print(f"WARNING: meteorology directory not found: {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}")

    # CONTROL & SETUP for dispersion
    write_setup_disp()
    write_control_disp(syr, smo, sda, shr, lat, lon, lvl, run_hours, top,
                       met_dir, dat_name, inp)

    # 3) Run hycs_std
    # Reuse the previous run output file
    if Path(inp).exists():
        print("Using previous simulation output file")
    else:
        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")

    # 4) concplot (+g1 -a1 -z10 -j<bg> ... ) base plume result (with GIS dump)
    concplot = PGM / "exec" / f"concplot{exe_suffix}"
    if not concplot.exists():
        print(f"ERROR: concplot not found at {concplot}")
        sys.exit(1)

    with open("LABELS.CFG", "w", encoding="ascii") as f:
        fn = Path(__file__).name if '__file__' in globals() else 'disp_merg.py'
        f.write(f"'TITLE&','### {fn}: base plume result ### &'\n")

    clean("concplot.html")
    run_cmd([
        str(concplot),
        "+g1",
        f"-i{inp}",
        "-a1",
        "-z100",
        f"-j{j_bg}",
        "-x1.0E+12",
        "-upg",
        "-h41.0:-82.0",
        "-g0:200",
        "-c4",
        "-b1000",
        "-t1000",
        "-v50000+20000+10000+5000+2000+1000",
    ], "Create base plume concplot (GIS dump)")

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

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

    # 5) If GIS_01000_html_01.txt exists -> build concplot shapefile
    gis_conc_txt = Path("GIS_01000_html_01.txt")
    if gis_conc_txt.exists():
        for leftover in glob.glob("concplot.???"):
            try:
                Path(leftover).unlink()
            except OSError:
                pass

        ascii2shp = PGM / "exec" / f"ascii2shp{exe_suffix}"
        txt2dbf  = PGM / "exec" / f"txt2dbf{exe_suffix}"
        if not ascii2shp.exists():
            print(f"ERROR: ascii2shp not found at {ascii2shp}")
            sys.exit(1)
        if not txt2dbf.exists():
            print(f"ERROR: txt2dbf not found at {txt2dbf}")
            sys.exit(1)

        with open(gis_conc_txt, "r", encoding="ascii") as fin:
            print("ascii2shp: creating concplot.shp")
            r = subprocess.run([str(ascii2shp), "-d", "concplot", "polygons"],
                               stdin=fin, capture_output=True, text=True)
            if r.returncode != 0:
                msg = r.stderr.strip() or r.stdout.strip() or "ascii2shp failed"
                print(f"ERROR: {msg}")
                sys.exit(1)

        run_cmd([
            str(txt2dbf),
            "-C11", "-C5", "-C9", "-C5", "-C6../", "-C6", "-d,",
            "GIS_01000_html_01.att", "concplot.dbf",
        ], "txt2dbf for concplot.dbf")

        # Clean batch leftovers
        for leftover in glob.glob("GIS_01000_html_01.*"):
            try:
                Path(leftover).unlink()
            except OSError:
                pass
    else:
        print("ERROR: Failed to create shapefiles ... GIS_01000_html_01.txt is missing!")

    # 6) Build trajectory CONTROL, run hyts_std
    write_control_traj(syr, smo, sda, shr, lat, lon, 750.0, run_hours, top,
                       met_dir, dat_name, tdump)
    if Path(tdump).exists():
        for f in glob.glob("tdmp_????"):
            try:
                Path(f).unlink()
            except OSError:
                pass
    clean("SETUP.CFG")

    hyts = PGM / "exec" / f"hyts_std{exe_suffix}"
    if not hyts.exists():
        print(f"ERROR: hyts_std not found at {hyts}")
        sys.exit(1)
    run_cmd([str(hyts)], "Running hyts_std")

    # 7) trajplot (+g1 -j<bg> -itdump_fwrd -a5 -k1 -z10 -l1 -v0)
    trajplot = PGM / "exec" / f"trajplot{exe_suffix}"
    if not trajplot.exists():
        print(f"ERROR: trajplot not found at {trajplot}")
        sys.exit(1)

    with open("LABELS.CFG", "w", encoding="ascii") as f:
        fn = Path(__file__).name if '__file__' in globals() else 'disp_merg.py'
        f.write(f"'TITLE&','### {fn}: base trajectory result ### &'\n")

    clean("trajplot.html")
    run_cmd([
        str(trajplot),
        "+g1",
        f"-j{j_bg}",
        f"-i{tdump}",
        "-a5",
        "-k1",
        "-z10",
        "-l1",
        "-v0",
    ], "Create base trajectory trajplot")

    html = Path("trajplot.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), [("smrg002.png", 1)])

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

    # 8) If GIS_traj_html_01.txt exists -> build trajplot shapefile
    gis_traj_txt = Path("GIS_traj_html_01.txt")
    if gis_traj_txt.exists():
        for leftover in glob.glob("trajplot.*"):
            try:
                Path(leftover).unlink()
            except OSError:
                pass

        ascii2shp = PGM / "exec" / f"ascii2shp{exe_suffix}"
        txt2dbf  = PGM / "exec" / f"txt2dbf{exe_suffix}"
        if not ascii2shp.exists():
            print(f"ERROR: ascii2shp not found at {ascii2shp}")
            sys.exit(1)
        if not txt2dbf.exists():
            print(f"ERROR: txt2dbf not found at {txt2dbf}")
            sys.exit(1)

        with open(gis_traj_txt, "r", encoding="ascii") as fin:
            print("ascii2shp: creating trajplot.shp")
            r = subprocess.run([str(ascii2shp), "-d", "trajplot", "lines"],
                               stdin=fin, capture_output=True, text=True)
            if r.returncode != 0:
                msg = r.stderr.strip() or r.stdout.strip() or "ascii2shp failed"
                print(f"ERROR: {msg}")
                sys.exit(1)

        run_cmd([
            str(txt2dbf),
            "-C7", "-C9", "-C5", "-C9", "-d,", "-d,", "-d,",
            "GIS_traj_html_01.att", "trajplot.dbf",
        ], "txt2dbf for trajplot.dbf")
    else:
        print("ERROR: GIS_traj_html_01.txt not found; trajplot shapefile not created.")

    # 9) Merged plot 1: TRAJPLOT with STATES + CONCPLOT.SHP (only include overlays that exist)
    j_merge1 = create_shapefiles_txt([
        (shp_states, '0 0.015 0.2 0.2 0.5'),
        (Path("concplot.shp"), '0 0.010 0.4 0.6 0.8')
    ])

    with open("LABELS.CFG", "w", encoding="ascii") as f:
        fn = Path(__file__).name if '__file__' in globals() else 'disp_merg.py'
        f.write(f"'TITLE&','### {fn}: base trajectory with plume ### &'\n")
    run_cmd([
        str(trajplot),
        "+g1",
        f"-j{j_merge1}",
        f"-i{tdump}",
        "-k1",
        "-z10",
        "-l1",
        "-v0",
    ], "TRAJPLOT merged with plume shapefile")

    html = Path("trajplot.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), [("smrg004.png", 1)])

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

    # 10) Merged plot 2: CONCPLOT with STATES + TRAJPLOT.SHP
    j_merge2 = create_shapefiles_txt([
        (shp_states, '0 0.015 0.2 0.2 0.5'),
        (Path("trajplot.shp"), '0 0.010 0.0 0.0 0.0')
    ])
    with open("LABELS.CFG", "w", encoding="ascii") as f:
        fn = Path(__file__).name if '__file__' in globals() else 'disp_merg.py'
        f.write(f"'TITLE&','### {fn}: base plume with trajectory ### &'\n")
    run_cmd([
        str(concplot),
        "+g1",
        f"-j{j_merge2}",
        f"-i{inp}",
        "-k1",
        "-z100",
        "-x1.0E+12",
        "-upg",
        "-h41.0:-82.0",
        "-g0:200",
        "-c4",
        "-b1000",
        "-t1000",
        "-v50000+20000+10000+5000+2000+1000",
        f"-q{met_dir / 'data_case.txt'}"
    ], "CONCPLOT merged with trajectory shapefile")

    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), [("smrg010.png", 1)])

    # Cleanup final temp label/shapefiles list (batch removes these)
    clean("LABELS.CFG")
    clean("shapefiles.txt")

if __name__ == "__main__":
    main()