"""
Run Dr Eureka solvers on a chain of states (pairwise: each goal becomes next start).
Uses all implemented algorithms: BFS, A* (Cursor), A* (Gippity), BiDirectional (Cursor), BiDirectional (Gippity).
States are generated randomly; chain length is configurable.
"""

import argparse
import random

from dr_eureka.a_star import AStarCursorSolver, AStarGippitySolver
from dr_eureka.bfs import BFSSolver
from dr_eureka.bi_directional import BiDirectionalCursorSolver, BiDirectionalGippitySolver
from dr_eureka.core import Colour, GameState, Tube
from dr_eureka.solver import DrEurekaSolver, SearchResult

# Valid partitions of 6 balls into 3 tubes with capacity 4: (n0, n1, n2) with 0 <= ni <= 4.
_PARTITIONS: list[tuple[int, int, int]] = [
    (4, 1, 1),
    (4, 2, 0),
    (3, 2, 1),
    (3, 3, 0),
    (2, 2, 2),
]


def random_state(
    *,
    num_tubes: int = 3,
    balls_per_color: int = 2,
    tube_capacity: int = 4,
    flip_chance: float = 0.0,
    rng: random.Random | None = None,
) -> GameState:
    """
    Generate a random valid GameState.

    Default: 3 tubes, 2 balls of each color (G, P, R) = 6 balls, capacity 4.
    Optionally flip each tube base-up with probability flip_chance (0 = all upright).
    """
    if rng is None:
        rng = random.Random()
    colors = [Colour.GREEN, Colour.PURPLE, Colour.RED]
    balls: list[Colour] = []
    for c in colors:
        balls.extend([c] * balls_per_color)
    total = len(balls)
    assert num_tubes >= 1 and total <= num_tubes * tube_capacity

    # Pick a random partition of total into num_tubes parts, each part <= tube_capacity.
    if num_tubes == 3 and total == 6 and tube_capacity == 4:
        sizes = list(rng.choice(_PARTITIONS))
        rng.shuffle(sizes)
    else:
        # General case: random assignment, retry if any tube overflows.
        while True:
            sizes = [0] * num_tubes
            for _ in range(total):
                sizes[rng.randint(0, num_tubes - 1)] += 1
            if all(s <= tube_capacity for s in sizes):
                break

    rng.shuffle(balls)
    idx = 0
    tubes: list[Tube] = []
    for size in sizes:
        tube_balls = tuple(balls[idx : idx + size])
        idx += size
        base_up = rng.random() < flip_chance if tube_balls else False
        if base_up:
            tube_balls = tube_balls[::-1]
        tubes.append(Tube(balls=tube_balls, base_up=base_up, capacity=tube_capacity))
    return GameState(tuple(tubes))


def generate_random_states(
    length: int,
    *,
    rng: random.Random | None = None,
    **kwargs: object,
) -> list[GameState]:
    """Generate a list of `length` random valid states."""
    if rng is None:
        rng = random.Random()
    return [random_state(rng=rng, **kwargs) for _ in range(length)]


SOLVERS: list[tuple[str, DrEurekaSolver]] = [
    ("BFS", BFSSolver()),
    ("A* (Cursor)", AStarCursorSolver()),
    ("A* (Gippity)", AStarGippitySolver()),
    ("BiDirectional (Cursor)", BiDirectionalCursorSolver()),
    ("BiDirectional (Gippity)", BiDirectionalGippitySolver()),
]


def run_pair(
    start: GameState,
    goal: GameState,
    pair_index: int,
) -> None:
    print(f"\n--- Pair {pair_index + 1}: state[{pair_index}] → state[{pair_index + 1}] ---")
    print("Start:", start)
    print("Goal: ", goal)

    for name, solver in SOLVERS:
        result: SearchResult | None = solver.solve(start, goal)
        if result is None:
            print(f"  {name}: no solution")
            continue
        print(f"  {name}: {result.move_count} moves, {result.visited_states} states visited")


def main() -> None:
    parser = argparse.ArgumentParser(description="Run Dr Eureka solvers on a chain of random states (pairwise).")
    parser.add_argument(
        "-n",
        "--length",
        type=int,
        default=6,
        metavar="N",
        help="Number of random states in the chain (default: 6).",
    )
    parser.add_argument(
        "-s",
        "--seed",
        type=int,
        default=None,
        metavar="SEED",
        help="Random seed for reproducible runs.",
    )
    parser.add_argument(
        "--flip-chance",
        type=float,
        default=0.0,
        metavar="P",
        help="Probability that a tube is generated base-up (default: 0).",
    )
    args = parser.parse_args()

    if args.length < 2:
        parser.error("Length must be at least 2.")
        return

    rng = random.Random(args.seed)
    states = generate_random_states(
        args.length,
        rng=rng,
        flip_chance=args.flip_chance,
    )

    print("Dr Eureka — chain of random states, all algorithms (pairwise)")
    print("=" * 60)
    print(f"Chain length: {args.length} states ({args.length - 1} pairs)")
    if args.seed is not None:
        print(f"Seed: {args.seed}")

    for i in range(len(states) - 1):
        run_pair(states[i], states[i + 1], i)

    print("\nDone.")


if __name__ == "__main__":
    main()
