#!/usr/bin/env python3
#
# traj_error.py --- run HYSPLIT computational trajectory error
#                   See HYSPLIT Tutorial Sec. 4.5
#
# Change history:
# 25 Jun 2025 (BB) - initial, cross-platform version
# 14 Jul 2025 (SZ) - create png image files when DSP is 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

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

# Detect OS and set paths
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 = ""

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


# Clean up old files
for f in ["ASCDATA.CFG", "tdump_fwrd", "tdump_back", "SETUP.CFG", "LABELS.CFG"]:
    f_path = OUT / f
    if f_path.exists():
        f_path.unlink()

# Create ASCDATA.CFG
ascdata = OUT / "ASCDATA.CFG"
with ascdata.open("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")

# Parameters
syr, smo, sda, shr = 83, 9, 25, 17
lat, lon, lvl = 39.90, -84.22, 600.0
run, top = 68, 10000.0
met = TTR / "captex"
dat = "captex2_wrf27uw.bin"
met_file = met / dat

if not met_file.exists():
    print(f"ERROR: Meteorological file not found: {met_file}")
    print("Please adjust the HYSPLIT Tutorial path for your system")
    sys.exit(1)

def write_control_and_run(lat, lon, lvl, run_hours, output_name, label_desc,
                          img_name):
    """Write CONTROL file, run trajectory, and generate plot"""
    ctl = OUT / "CONTROL"
    with ctl.open("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} {lon} {lvl}\n")
        f.write(f"{run_hours}\n")
        f.write("0\n")
        f.write(f"{top}\n")
        f.write("1\n")
        f.write(f"{met}{os.sep}\n")
        f.write(f"{dat}\n")
        f.write(f".{os.sep}\n")
        f.write(f"{output_name}\n")

    print(f"\nCONTROL file preview for {label_desc}:")
    print(ctl.read_text(encoding="ascii"))

    # Remove existing output file
    output_file = OUT / output_name
    if output_file.exists():
        output_file.unlink()

    # Run trajectory
    run_cmd([str(PGM / "exec" / f"hyts_std{exe_suffix}")], f"Running {label_desc}")

    # Create LABELS.CFG
    label_cfg = OUT / "LABELS.CFG"
    with label_cfg.open("w", encoding="ascii") as f:
        f.write(f"'TITLE&','### {label_desc} ### &'\n")

    # Run trajplot
    plot_args = ["-v0", "-z50"] if run_hours > 0 else ["-v1", "-z50"]
    run_cmd([
        str(PGM / "exec" / f"trajplot{exe_suffix}"),
        "+g1", f"-i{output_name}"] + plot_args + [
        f"-j{PGM / 'graphics' / 'arlmap'}"
    ], f"Running trajplot for {label_desc}")

    # Display plot
    html_file = OUT / "trajplot.html"
    if html_file.exists():
        print(f"{label_desc} plot generated: trajplot.html")
        if DSP == "YES":
            webbrowser.open(str(html_file))
        else:
            convert_html_to_img(str(html_file), [(img_name,1)])
    else:
        print(f"Warning: trajplot.html not generated for {label_desc}")

    return output_file

def read_final_position(tdump_file):
    """Read final position from trajectory dump file"""
    try:
        with open(tdump_file, 'r', encoding="ascii") as f:
            lines = f.readlines()
            if lines:
                # Get the last non-empty line and extract columns 10, 11, 12 (0-indexed: 9, 10, 11)
                last_line = lines[-1].strip()
                if last_line:
                    tokens = last_line.split()
                    if len(tokens) >= 12:
                        final_lat = float(tokens[9])   # Column 10 (0-indexed 9)
                        final_lon = float(tokens[10])  # Column 11 (0-indexed 10)
                        final_lvl = float(tokens[11])  # Column 12 (0-indexed 11)
                        print(f"Final position: lat={final_lat}, lon={final_lon}, lvl={final_lvl}")
                        return final_lat, final_lon, final_lvl
    except (FileNotFoundError, ValueError, IndexError) as e:
        print(f"ERROR: reading {tdump_file}: {e}")
        sys.exit(1)

    print(f"ERROR: Could not extract final position from {tdump_file}")
    sys.exit(1)

def main():
    # Forward trajectory
    print("=== FORWARD TRAJECTORY ===")
    write_control_and_run(lat, lon, lvl, run, "tdump_fwrd", "Forward trajectory",
                          "terr001.png")

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

    # Read final position from forward trajectory
    final_lat, final_lon, final_lvl = read_final_position("tdump_fwrd")

    # Update parameters for backward trajectory
    global sda, shr
    sda, shr = 28, 13

    # Backward trajectory from final position
    print("\n=== BACKWARD TRAJECTORY ===")
    write_control_and_run(final_lat, final_lon, final_lvl, -run, "tdump_back",
                          "Backward trajectory", "terr003.png")

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

    # Combined trajectory plot
    print("\n=== COMBINED TRAJECTORY PLOT ===")
    label_cfg = OUT / "LABELS.CFG"
    with label_cfg.open("w", encoding="ascii") as f:
        f.write("'TITLE&','### Combined trajectory ### &'\n")

    run_cmd([
        str(PGM / "exec" / f"trajplot{exe_suffix}"),
        "+g1", "-itdump_back+tdump_fwrd", "-v1", "-z50",
        f"-j{PGM / 'graphics' / 'arlmap'}"
    ], "Creating combined trajectory plot")

    html_file = OUT / "trajplot.html"
    if html_file.exists():
        print("Combined trajectory plot generated: trajplot.html")
        if DSP == "YES":
            webbrowser.open(str(html_file))
        else:
            convert_html_to_img(str(html_file), [("terr005.png",1)])
    else:
        print("Warning: combined trajplot.html not generated")

    print("\nScript completed successfully!")

if __name__ == "__main__":
    main()