#!/usr/bin/env python3
"""
Cross-implementation benchmark: readcon-core vs ASE vs eOn Python reader.

Generates a test trajectory, times each reader, and produces comparison
bar charts saved to docs/orgmode/img/.

Usage:
    uv run --with matplotlib --with numpy --with ase benches/compare_readers.py
    # or: pixi r -e python python benches/compare_readers.py

Requires: matplotlib, numpy.
Optional: readcon (Python bindings), ase (uv run --with ase).
eOn Python reader: uses ~/Git/Github/TheochemUI/eOn if present.
"""

import os
import sys
import time
import tempfile
from pathlib import Path

# ---------------------------------------------------------------------------
# Configuration
# ---------------------------------------------------------------------------

NUM_FRAMES = 100
REPEAT = 5  # number of timing repetitions

REPO_ROOT = Path(__file__).resolve().parent.parent
IMG_DIR = REPO_ROOT / "docs" / "orgmode" / "img"
IMG_DIR.mkdir(parents=True, exist_ok=True)

# eOn Python reader: use the actual eOn checkout if available
EON_REPO = Path.home() / "Git" / "Github" / "TheochemUI" / "eOn"
if EON_REPO.exists():
    sys.path.insert(0, str(EON_REPO))

# ---------------------------------------------------------------------------
# Generate test trajectory
# ---------------------------------------------------------------------------

SINGLE_FRAME = """\
Generated by benchmark
{"con_spec_version":2}
15.345600\t21.702000\t100.000000
90.000000\t90.000000\t90.000000
0 0
218 0 1
2
216 2
63.546000 1.007930
"""

def generate_atom_lines():
    """Generate 218 atom lines (216 Cu + 2 H)."""
    lines = []
    lines.append("Cu\n")
    lines.append("Coordinates of Component 1\n")
    for i in range(216):
        x = (i % 6) * 2.557 + 0.6394
        y = ((i // 6) % 6) * 3.617 + 0.9045
        z = (i // 36) * 2.087
        lines.append(f"   {x:.17f}    {y:.17f}    {z:.17f} 7    {i}\n")
    lines.append("H\n")
    lines.append("Coordinates of Component 2\n")
    lines.append(f"   8.682300    9.947000   11.733000 0  216\n")
    lines.append(f"   7.942100    9.947000   11.733000 0  217\n")
    return "".join(lines)


def generate_trajectory(num_frames, path):
    """Write a multi-frame .con trajectory file."""
    atom_block = generate_atom_lines()
    with open(path, "w") as f:
        for _ in range(num_frames):
            f.write(SINGLE_FRAME)
            f.write(atom_block)


# ---------------------------------------------------------------------------
# Timing helpers
# ---------------------------------------------------------------------------

def time_fn(fn, repeat=REPEAT):
    """Time a function, return median time in ms."""
    times = []
    for _ in range(repeat):
        t0 = time.perf_counter()
        fn()
        t1 = time.perf_counter()
        times.append((t1 - t0) * 1000)
    times.sort()
    return times[len(times) // 2]


# ---------------------------------------------------------------------------
# Reader implementations
# ---------------------------------------------------------------------------

def bench_eon_python(path):
    """eOn's pure-Python loadcons()."""
    try:
        from eon.fileio import loadcons
        return time_fn(lambda: loadcons(path))
    except Exception as e:
        print(f"  eOn Python: skipped ({e})")
        return None


def bench_c_reader(path):
    """Standalone C reader (eOn-style sscanf parsing)."""
    import subprocess
    c_binary = REPO_ROOT / "addl" / "referenceImpls" / "eon_cpp" / "bench_con_reader"
    if not c_binary.exists():
        # Try to build it
        c_src = c_binary.parent / "bench_con_reader.c"
        if c_src.exists():
            try:
                subprocess.run(
                    ["cc", "-O2", "-o", str(c_binary), str(c_src), "-lm"],
                    check=True, capture_output=True,
                )
            except Exception as e:
                print(f"  C reader: build failed ({e})")
                return None
        else:
            print(f"  C reader: source not found")
            return None

    try:
        result = subprocess.run(
            [str(c_binary), path, str(REPEAT)],
            capture_output=True, text=True, check=True,
        )
        # Output: "10.55 ms (best of 10 runs)"
        ms = float(result.stdout.strip().split()[0])
        return ms
    except Exception as e:
        print(f"  C reader: failed ({e})")
        return None


def bench_ase(path):
    """ASE's read_eon()."""
    try:
        from ase.io import read
        return time_fn(lambda: read(path, index=":", format="eon"))
    except Exception as e:
        print(f"  ASE: skipped ({e})")
        return None


def bench_readcon(path):
    """readcon-core Python bindings."""
    try:
        import readcon
        return time_fn(lambda: readcon.read_con(path))
    except Exception as e:
        print(f"  readcon: skipped ({e})")
        return None


def bench_readcon_string(path):
    """readcon-core from pre-loaded string (no file I/O)."""
    try:
        import readcon
        with open(path) as f:
            contents = f.read()
        return time_fn(lambda: readcon.read_con_string(contents))
    except Exception as e:
        print(f"  readcon (string): skipped ({e})")
        return None


# ---------------------------------------------------------------------------
# Plotting
# ---------------------------------------------------------------------------

def make_bar_chart(results, title, filename):
    """Generate a bar chart from {label: time_ms} dict."""
    try:
        import matplotlib
        matplotlib.use("Agg")
        import matplotlib.pyplot as plt
    except ImportError:
        print("  matplotlib not available, skipping plot")
        return

    labels = list(results.keys())
    values = list(results.values())
    colors = ["#004D40", "#FF655D", "#F1DB4B", "#4DB6AC"]

    fig, ax = plt.subplots(figsize=(8, 4.5))
    bars = ax.barh(labels, values, color=colors[:len(labels)])
    ax.set_xlabel("Time (ms)")
    ax.set_title(title)
    ax.invert_yaxis()

    # Value labels on bars
    for bar, val in zip(bars, values):
        ax.text(bar.get_width() + max(values) * 0.02, bar.get_y() + bar.get_height() / 2,
                f"{val:.1f} ms", va="center", fontsize=10)

    plt.tight_layout()
    out = IMG_DIR / filename
    plt.savefig(out, dpi=150)
    plt.close()
    print(f"  Saved: {out}")


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------

def main():
    with tempfile.NamedTemporaryFile(suffix=".con", delete=False, mode="w") as f:
        traj_path = f.name

    try:
        print(f"Generating {NUM_FRAMES}-frame trajectory (218 atoms/frame)...")
        generate_trajectory(NUM_FRAMES, traj_path)
        file_size = os.path.getsize(traj_path)
        print(f"  File size: {file_size / 1024:.1f} KiB")

        print(f"\nBenchmarking ({REPEAT} repetitions, median)...")
        results = {}

        t = bench_c_reader(traj_path)
        if t is not None:
            results["C (sscanf)\n(eOn-style)"] = t
            print(f"  C (sscanf):        {t:.1f} ms")

        t = bench_eon_python(traj_path)
        if t is not None:
            results["eOn Python\n(fileio.py)"] = t
            print(f"  eOn Python:        {t:.1f} ms")

        t = bench_ase(traj_path)
        if t is not None:
            results["ASE\n(ase.io.eon)"] = t
            print(f"  ASE:               {t:.1f} ms")

        t = bench_readcon(traj_path)
        if t is not None:
            results["readcon-core\n(file path)"] = t
            print(f"  readcon (file):    {t:.1f} ms")

        t = bench_readcon_string(traj_path)
        if t is not None:
            results["readcon-core\n(from string)"] = t
            print(f"  readcon (string):  {t:.1f} ms")

        if len(results) >= 2:
            print("\nGenerating plots...")
            make_bar_chart(
                results,
                f"Parsing {NUM_FRAMES} frames x 218 atoms",
                "parsing_throughput.svg",
            )
        else:
            print("\nNot enough readers available for comparison plot.")

        # Summary table
        print("\n--- Results ---")
        print(f"{'Reader':<25} {'Time (ms)':>10} {'Speedup':>10}")
        print("-" * 47)
        baseline = max(results.values()) if results else 1.0
        for label, t in sorted(results.items(), key=lambda x: x[1], reverse=True):
            label_clean = label.replace("\n", " ")
            speedup = baseline / t
            print(f"{label_clean:<25} {t:>10.1f} {speedup:>9.1f}x")

    finally:
        os.unlink(traj_path)


if __name__ == "__main__":
    main()
