"""
Performance testbed for Dr Eureka solvers.

Runs many chains (n states each) m times, measures moves, states visited, and wall time
per solver, then prints correctness and leaderboards (by algorithm, by AI, and extras).
"""

import argparse
import random
import time
from collections import defaultdict
from dataclasses import dataclass, field

from tqdm import trange

# Import chain generation and solvers from the main script
from dr_eureka.__main__ import SOLVERS, generate_random_states
from dr_eureka.core import GameState
from dr_eureka.solver import DrEurekaSolver, SearchResult


@dataclass
class RunResult:
    """Result of one solver on one (start, goal) pair."""

    solver_name: str
    moves: int | None  # None if no solution
    visited_states: int | None
    wall_time_seconds: float


@dataclass
class TestbedResult:
    """Aggregated results from the full testbed."""

    chain_length: int
    iterations: int
    seed: int | None
    # Per (iteration, pair_index, solver_name) -> RunResult
    runs: list[tuple[int, int, str, RunResult]] = field(default_factory=list)

    def _group_by_solver(self) -> dict[str, list[RunResult]]:
        by_solver: dict[str, list[RunResult]] = defaultdict(list)
        for _iter, _pair, name, r in self.runs:
            by_solver[name].append(r)
        return dict(by_solver)

    def correctness_report(self) -> tuple[bool, str]:
        """
        Check that every solver that found a solution found a shortest one.
        Returns (all_correct, message).
        """
        # Group by (iteration, pair_index) -> list of (solver_name, moves)
        by_pair: dict[tuple[int, int], list[tuple[str, int]]] = defaultdict(list)
        for it, pair, name, r in self.runs:
            if r.moves is not None:
                by_pair[(it, pair)].append((name, r.moves))

        all_ok = True
        mismatches: list[str] = []
        for (it, pair), results in by_pair.items():
            if not results:
                continue
            best = min(m for _, m in results)
            for name, m in results:
                if m != best:
                    all_ok = False
                    mismatches.append(f"  iter={it} pair={pair} {name}: {m} moves (best={best})")
        msg = "All algorithms found optimal (shortest) solutions."
        if not all_ok:
            msg = f"Correctness issues ({len(mismatches)}):\n" + "\n".join(mismatches[:20])
            if len(mismatches) > 20:
                msg += f"\n  ... and {len(mismatches) - 20} more"
        return all_ok, msg

    def leaderboard_states_visited(self) -> list[tuple[str, float, int]]:
        """Rank solvers by mean states visited (only solved runs). Returns [(name, mean, total_runs)]."""
        by_solver = self._group_by_solver()
        rows: list[tuple[str, float, int]] = []
        for name, results in by_solver.items():
            solved = [r for r in results if r.visited_states is not None]
            if not solved:
                rows.append((name, float("nan"), 0))
                continue
            mean = sum(r.visited_states for r in solved) / len(solved)
            rows.append((name, mean, len(solved)))
        rows.sort(key=lambda x: (x[1] if x[1] == x[1] else float("inf"), x[0]))
        return rows

    def leaderboard_wall_time(self) -> list[tuple[str, float, float, int]]:
        """Rank solvers by total wall time and mean time. Returns [(name, total_s, mean_s, runs)]."""
        by_solver = self._group_by_solver()
        rows: list[tuple[str, float, float, int]] = []
        for name, results in by_solver.items():
            solved = [r for r in results if r.wall_time_seconds is not None]
            if not solved:
                rows.append((name, 0.0, float("nan"), 0))
                continue
            total = sum(r.wall_time_seconds for r in solved)
            mean = total / len(solved)
            rows.append((name, total, mean, len(solved)))
        rows.sort(key=lambda x: (x[1], x[0]))  # by total time
        return rows

    def by_algorithm(self) -> dict[str, list[str]]:
        """Group solver names by algorithm: BFS, A*, BiDirectional."""
        algo_map: dict[str, list[str]] = {"BFS": [], "A*": [], "BiDirectional": []}
        for name, _ in SOLVERS:
            if name == "BFS":
                algo_map["BFS"].append(name)
            elif name.startswith("A*"):
                algo_map["A*"].append(name)
            elif name.startswith("BiDirectional"):
                algo_map["BiDirectional"].append(name)
        return algo_map

    def by_ai(self) -> dict[str, list[str]]:
        """Group solver names by AI: Cursor vs Gippity."""
        cursor: list[str] = []
        gippity: list[str] = []
        for name, _ in SOLVERS:
            if "Cursor" in name:
                cursor.append(name)
            elif "Gippity" in name:
                gippity.append(name)
        return {"Cursor": cursor, "Gippity": gippity}

    def algorithm_comparison(self) -> str:
        """Compare algorithms (BFS vs A* vs BiDirectional) by mean states and mean time."""
        by_solver = self._group_by_solver()
        algo = self.by_algorithm()
        lines = ["--- By algorithm ---"]
        for algo_name, solver_names in algo.items():
            if not solver_names:
                continue
            all_visits = []
            all_times = []
            for name in solver_names:
                for r in by_solver.get(name, []):
                    if r.visited_states is not None:
                        all_visits.append(r.visited_states)
                    if r.wall_time_seconds is not None:
                        all_times.append(r.wall_time_seconds)
            if all_visits and all_times:
                avg_v = sum(all_visits) / len(all_visits)
                avg_t = sum(all_times) / len(all_times)
                lines.append(f"  {algo_name}: mean states visited = {avg_v:.1f}, mean time = {avg_t * 1000:.2f} ms")
        return "\n".join(lines)

    def ai_comparison(self) -> str:
        """Compare Cursor vs Gippity by mean states and mean time."""
        by_solver = self._group_by_solver()
        ai = self.by_ai()
        lines = ["--- By AI (Cursor vs Gippity) ---"]
        for ai_name, solver_names in ai.items():
            if not solver_names:
                continue
            all_visits = []
            all_times = []
            for name in solver_names:
                for r in by_solver.get(name, []):
                    if r.visited_states is not None:
                        all_visits.append(r.visited_states)
                    if r.wall_time_seconds is not None:
                        all_times.append(r.wall_time_seconds)
            if all_visits and all_times:
                avg_v = sum(all_visits) / len(all_visits)
                avg_t = sum(all_times) / len(all_times)
                lines.append(f"  {ai_name}: mean states = {avg_v:.1f}, mean time = {avg_t * 1000:.2f} ms")
        return "\n".join(lines)

    def extra_stats(self) -> str:
        """Extra fun stats."""
        by_solver = self._group_by_solver()
        lines = ["--- Extra stats ---"]

        # Solve rate per solver
        self.iterations * (self.chain_length - 1)
        for name, results in sorted(by_solver.items()):
            solved = sum(1 for r in results if r.moves is not None)
            rate = 100 * solved / len(results) if results else 0
            lines.append(f"  {name}: solved {solved}/{len(results)} ({rate:.1f}%)")

        # Move distribution (from first solver that has data)
        all_moves: list[int] = []
        by_pair: dict[tuple[int, int], int] = {}
        for it, pair, _, r in self.runs:
            if r.moves is not None:
                by_pair[(it, pair)] = min(by_pair.get((it, pair), r.moves), r.moves)
        all_moves = list(by_pair.values())
        if all_moves:
            lines.append(
                f"  Move count: min={min(all_moves)}, max={max(all_moves)}, mean={sum(all_moves) / len(all_moves):.1f}"
            )

        # Which solver was fastest (by time) most often per pair
        wins: dict[str, int] = defaultdict(int)
        by_pair_solver: dict[tuple[int, int], list[tuple[str, float]]] = defaultdict(list)
        for it, pair, name, r in self.runs:
            if r.wall_time_seconds is not None:
                by_pair_solver[(it, pair)].append((name, r.wall_time_seconds))
        for _key, times in by_pair_solver.items():
            best = min(t for _, t in times)
            for name, t in times:
                if t == best:
                    wins[name] += 1
        solver_names = [name for name, _ in SOLVERS]
        if solver_names:
            lines.append("  Wins (fastest wall time per pair):")
            for name in sorted(solver_names, key=lambda n: -wins[n]):
                lines.append(f"    {name}: {wins[name]}")

        return "\n".join(lines)


def run_one(
    solver: DrEurekaSolver,
    solver_name: str,
    start: GameState,
    goal: GameState,
) -> RunResult:
    """Run a single solver on one pair; return RunResult with moves, visited, and time."""
    t0 = time.perf_counter()
    result: SearchResult | None = solver.solve(start, goal)
    t1 = time.perf_counter()
    if result is None:
        return RunResult(
            solver_name=solver_name,
            moves=None,
            visited_states=None,
            wall_time_seconds=t1 - t0,
        )
    return RunResult(
        solver_name=solver_name,
        moves=result.move_count,
        visited_states=result.visited_states,
        wall_time_seconds=t1 - t0,
    )


def run_testbed(
    chain_length: int,
    iterations: int,
    *,
    seed: int | None = None,
    flip_chance: float = 0.0,
) -> TestbedResult:
    """
    Run the testbed: `iterations` chains of `chain_length` states each.
    For each chain, solve each consecutive pair with every solver.
    """
    rng = random.Random(seed)
    tb = TestbedResult(
        chain_length=chain_length,
        iterations=iterations,
        seed=seed,
    )

    for it in trange(iterations, desc="Chains", unit="chain"):
        states = generate_random_states(
            chain_length,
            rng=rng,
            flip_chance=flip_chance,
        )
        for pair_idx in range(len(states) - 1):
            start, goal = states[pair_idx], states[pair_idx + 1]
            for name, solver in SOLVERS:
                r = run_one(solver, name, start, goal)
                tb.runs.append((it, pair_idx, name, r))

    return tb


def print_report(tb: TestbedResult) -> None:
    """Print correctness, leaderboards, and extra stats."""
    print("\n" + "=" * 60)
    print("CORRECTNESS")
    print("=" * 60)
    ok, msg = tb.correctness_report()
    print(msg)

    print("\n" + "=" * 60)
    print("LEADERBOARD: Fewest states visited (mean over solved runs)")
    print("=" * 60)
    for name, mean_v, count in tb.leaderboard_states_visited():
        mean_str = f"{mean_v:.1f}" if mean_v == mean_v else "N/A"
        print(f"  {name}: {mean_str} states (n={count})")

    print("\n" + "=" * 60)
    print("LEADERBOARD: Fastest wall time (total then mean ms)")
    print("=" * 60)
    for name, total_s, mean_s, count in tb.leaderboard_wall_time():
        mean_str = f"{mean_s * 1000:.2f}" if mean_s == mean_s else "N/A"
        print(f"  {name}: total {total_s:.2f} s, mean {mean_str} ms (n={count})")

    print("\n" + tb.algorithm_comparison())
    print("\n" + tb.ai_comparison())
    print("\n" + tb.extra_stats())
    print()


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Run performance testbed for Dr Eureka solvers (n-chain × m iterations)."
    )
    parser.add_argument(
        "-n",
        "--chain-length",
        type=int,
        default=6,
        metavar="N",
        help="Number of states per chain (default: 6).",
    )
    parser.add_argument(
        "-m",
        "--iterations",
        type=int,
        default=100,
        metavar="M",
        help="Number of chains to run (default: 100).",
    )
    parser.add_argument(
        "-s",
        "--seed",
        type=int,
        default=None,
        metavar="SEED",
        help="Random seed for reproducible testbed.",
    )
    parser.add_argument(
        "--flip-chance",
        type=float,
        default=0.0,
        metavar="P",
        help="Probability a tube is base-up when generating states (default: 0).",
    )
    args = parser.parse_args()

    if args.chain_length < 2:
        parser.error("Chain length must be at least 2.")
    if args.iterations < 1:
        parser.error("Iterations must be at least 1.")

    print("Dr Eureka — performance testbed")
    print("=" * 60)
    print(f"Chain length: {args.chain_length}  Iterations: {args.iterations}")
    if args.seed is not None:
        print(f"Seed: {args.seed}")
    print("Running...")

    tb = run_testbed(
        args.chain_length,
        args.iterations,
        seed=args.seed,
        flip_chance=args.flip_chance,
    )
    print_report(tb)


if __name__ == "__main__":
    main()
