"""Search / execution engine — candidate generation & evaluation.

This is the Kudret organ: it *effects* the transition from candidate
programs to concrete predictions on unseen test inputs (AX4).

Architecture:
  ArcSolver.solve_task  →  synthesize  →  apply to test inputs
  ArcSolver.batch_solve →  iterate over many tasks, collect stats
"""

from __future__ import annotations

import time
from typing import Dict, List, Optional

from . import grid as G
from .synthesis import synthesize, validates_all
from .types import Grid, Solution, Task, Transform


class ArcSolver:
    """ARC-AGI solver orchestrating perception → synthesis → execution."""

    def __init__(
        self,
        *,
        max_compose_depth: int = 2,
        timeout_candidates: int = 5000,
        verbose: bool = False,
    ):
        self.max_compose_depth = max_compose_depth
        self.timeout_candidates = timeout_candidates
        self.verbose = verbose

    # ── Per-task solving ──────────────────────────────────────
    def solve_task(self, task: Task) -> Solution:
        """Solve a single ARC task.  Returns predictions for every test input."""
        t0 = time.perf_counter()

        transform = synthesize(
            task,
            max_compose_depth=self.max_compose_depth,
            timeout_candidates=self.timeout_candidates,
        )

        elapsed = time.perf_counter() - t0
        if self.verbose:
            status = transform.name if transform else "UNSOLVED"
            print(f"  {task.task_id}: {status} ({elapsed:.2f}s)")

        if transform is None:
            return Solution(
                task_id=task.task_id,
                predictions=[],
                transform=None,
                confidence=0.0,
            )

        predictions = []
        for inp in task.test_inputs:
            result = transform(inp)
            if result is not None:
                predictions.append(result)
            else:
                # Return empty grid as fallback
                predictions.append([])

        return Solution(
            task_id=task.task_id,
            predictions=predictions,
            transform=transform,
            confidence=1.0,
        )

    # ── Batch solving with statistics ─────────────────────────
    def batch_solve(self, tasks: List[Task]) -> Dict:
        """Solve many tasks and return aggregate statistics."""
        solutions: List[Solution] = []
        correct = 0
        attempted = 0
        total = len(tasks)

        t_start = time.perf_counter()
        for task in tasks:
            sol = self.solve_task(task)
            solutions.append(sol)

            if sol.predictions:
                attempted += 1
                # Check against known test outputs
                if task.test_outputs and len(sol.predictions) == len(task.test_outputs):
                    all_match = all(
                        G.grids_equal(pred, expected)
                        for pred, expected in zip(sol.predictions, task.test_outputs)
                    )
                    if all_match:
                        correct += 1

        t_total = time.perf_counter() - t_start
        accuracy = correct / total if total else 0.0

        stats = {
            "total": total,
            "attempted": attempted,
            "correct": correct,
            "accuracy": accuracy,
            "time_seconds": round(t_total, 2),
        }
        if self.verbose:
            print(f"\n── Results: {correct}/{total} correct "
                  f"({accuracy:.1%}) in {t_total:.1f}s ──")

        return {"stats": stats, "solutions": solutions}
