#!/usr/bin/env python3
#
# traj_clus.py --- Cluster Analysis
#                  See HYSPLIT Tutorial Sec. 6.2
#
# Change history:
# 18 Aug 2025 (BB) - standard-aligned, cross-platform version
# 21 Aug 2025 (SZ) - use backward slash as directory separator for Windows PC.
#                  - show help messages when HYSPLIT or Tutorial is not found.
#                  - replace convert_html_to_img() with the one used before.
#                  - rename image filenames to match those created by the Linux script.
# 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

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 = ""

DSP = str(DSP).upper()

if not PGM.exists():
    script = Path(__file__).name if "__file__" in globals() else "traj_clus.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 exe(name: str) -> str:
    return str(PGM / "exec" / f"{name}{exe_suffix}")


def remove_if_exists(path: Path):
    try:
        if path.exists():
            path.unlink()
    except Exception:
        pass


def handle_html(html_path: Path, stem: str, pause_after=True):
    if html_path.exists():
        if DSP == "YES":
            webbrowser.open(str(html_path.resolve()))
            if pause_after:
                input("Press Enter to continue...")
        else:
            convert_html_to_img(str(html_path), [(stem,1)])
    else:
        print(f"Warning: {html_path.name} not generated")


def make_k_palette(n: int) -> str:
    digits = "1234567890"
    reps = (n + len(digits) - 1) // len(digits)
    return (digits * reps)[:n]


# pre-check (mirrors BAT)
if not Path("fdump83090100").exists() and not Path("fdump83092818").exists():
    print("trajectory endpoint {fdump} files missing, rerun previous batch file")
    if DSP == "YES":
        try:
            input("Press Enter to exit...")
        except EOFError:
            pass
    sys.exit(1)

ilab1 = "traj_clus"
nc = 6
arlmap = PGM / "graphics" / "arlmap"
has_arlmap = arlmap.exists()

# section 1
with open("INFILE", "w", encoding="ascii") as f:
    for p in sorted(Path(".").glob("fdump????????")):
        f.write(p.name + "\n")

dur = 48
intv = 1
skip = 1
ODIR = ".\\" if system == "windows" else "./"
PROJ = 0

with open("CCONTROL", "w", encoding="ascii") as f:
    f.write(f"{dur}\n{intv}\n{skip}\n{ODIR}\n{PROJ}\n")

run_cmd([exe("cluster")])
remove_if_exists(Path("TNOCLUS"))
remove_if_exists(Path("TCLUS"))

# section 2
mint = 15
minc = 3
maxc = 10
pct = 30

args = [exe("clusplot"), "+g1", "-iDELPCT", "-oclusplot.html", f"-l{ilab1}"]
run_cmd(args)
handle_html(OUT / "clusplot.html", "clus004.png", pause_after=True)

run_cmd([exe("clusend"), "-iDELPCT", f"-n{minc}", f"-a{maxc}", f"-t{mint}",
         f"-p{pct}", "-oCLUSEND"])
run_cmd([exe("cluslist"), "-iCLUSTER", f"-n{nc}", f"-oCLUSLIST_{nc}"])

if Path("CLUSTERno").exists():
    if Path(f"CLUSLIST_{nc}").exists():
        with open("CLUSTERno", "a", encoding="ascii") as out_f, \
            open(f"CLUSLIST_{nc}", "r", encoding="ascii") as in_f:
            out_f.write(in_f.read())
        remove_if_exists(Path(f"CLUSLIST_{nc}"))
    Path("CLUSTERno").rename(f"CLUSLIST_{nc}")

run_cmd([exe("clusmem"), f"-iCLUSLIST_{nc}"])

# discover clusters produced
made_clusters = []
for fp in sorted(Path(".").glob(f"TRAJ.INP.C*_{nc}")):
    try:
        N = int(fp.name.split("_")[0].split("C")[-1])
        made_clusters.append(N)
    except ValueError:
        continue

for N in made_clusters:
    run_cmd([exe("merglist"), f"-i+TRAJ.INP.C{N}_{nc}", "-omdump", "-ptdump"])
    if Path("mdump.tdump").exists():
        remove_if_exists(OUT / f"C{N}.tdump")
        Path("mdump.tdump").rename(OUT / f"C{N}.tdump")
    run_cmd([exe("trajmean"), f"-i+TRAJ.INP.C{N}_{nc}"])
    if Path("tmean").exists():
        remove_if_exists(OUT / f"C{N}mean.tdump")
        Path("tmean").rename(OUT / f"C{N}mean.tdump")

remove_if_exists(OUT / "MEAN.LIST")
with open(OUT / "MEAN.LIST", "w", encoding="ascii") as f:
    for N in made_clusters:
        if (OUT / f"C{N}mean.tdump").exists():
            f.write(f"C{N}mean.tdump\n")

run_cmd([exe("merglist"), "-i+MEAN.LIST", "-omdump", "-ptdump"])
if Path("mdump.tdump").exists():
    remove_if_exists(OUT / "Cmean.tdump")
    Path("mdump.tdump").rename(OUT / "Cmean.tdump")
remove_if_exists(OUT / "MEAN.LIST")

# section 3 (plots)
for N in range(1, nc + 1):
    tdump = OUT / f"C{N}.tdump"
    if not tdump.exists():
        print(f"Skipping cluster {N}: {tdump.name} not found")
        continue
    try:
        if tdump.stat().st_size == 0:
            print(f"Skipping cluster {N}: {tdump.name} is empty")
            continue
    except OSError:
        print(f"Skipping cluster {N}: cannot stat {tdump.name}")
        continue

    (OUT / "LABELS.CFG").write_text(f"'TITLE&','TRAJECTORIES IN CLUSTER {N} of {nc} Example&'\n",
                                    encoding="ascii")

    cmd = [exe("trajplot"), "+g1", f"-iC{N}.tdump", f"-e{dur}", "-z80", "-k0", "-l0", "-v4"]
    if has_arlmap:
        cmd.append(f"-j{arlmap}")
    run_cmd(cmd, allow_fail=True)

    if (OUT / "trajplot.html").exists():
        remove_if_exists(OUT / f"trajplot_C{N}.html")
        (OUT / "trajplot.html").rename(OUT / f"trajplot_C{N}.html")
        handle_html(OUT / f"trajplot_C{N}.html", f"trajplot_C{N:02d}.png", pause_after=True)
    else:
        print(f"WARN: trajplot.html missing after plotting cluster {N}; skipping open/convert.")

# mean plot with matching -k n:m
tmout = OUT / "Cmean.tdump"
k_n = sum(1 for N in made_clusters if (OUT / f"C{N}mean.tdump").exists())

if not tmout.exists() or k_n == 0:
    print("WARN: No cluster means available to plot; skipping mean plot.")
else:
    (OUT / "LABELS.CFG").write_text(f"'TITLE&','CLUSTER MEAN TRAJECTORIES {ilab1}&'\n",
                                    encoding="ascii")
    palette = make_k_palette(k_n)
    cmd = [exe("trajplot"), "+g1", f"-i{tmout.name}", f"-e{dur}", "-z80",
           f"-k{k_n}:{palette}", "-l-12", "-v4"]
    if has_arlmap:
        cmd.append(f"-j{arlmap}")
    run_cmd(cmd, allow_fail=True)

    if (OUT / "trajplot.html").exists():
        remove_if_exists(OUT / "trajplot_MEAN.html")
        (OUT / "trajplot.html").rename(OUT / "trajplot_MEAN.html")
        handle_html(OUT / "trajplot_MEAN.html", "clus006.png", pause_after=False)
    else:
        print("WARN: trajplot.html missing after mean plot; skipping open/convert.")

print("\nScript completed successfully!")