"""ARC-AGI solver — shared type definitions.

Maps to the İlim (Distinguishing) phase of three-phase actualization:
types are the distinctions that ground all downstream operations (AX2).
"""

from __future__ import annotations

import json
from dataclasses import dataclass, field
from pathlib import Path
from typing import (
    Callable,
    Dict,
    FrozenSet,
    List,
    Optional,
    Set,
    Tuple,
)

# ── Core type aliases ──────────────────────────────────────────────
Grid = List[List[int]]       # 2-D matrix of ints 0–9
Point = Tuple[int, int]      # (row, col)
Color = int                  # 0–9


# ── Geometry ───────────────────────────────────────────────────────
@dataclass(frozen=True)
class BBox:
    """Axis-aligned bounding box — inclusive on all sides."""

    r0: int
    c0: int
    r1: int
    c1: int

    @property
    def height(self) -> int:
        return self.r1 - self.r0 + 1

    @property
    def width(self) -> int:
        return self.c1 - self.c0 + 1

    @property
    def area(self) -> int:
        return self.height * self.width


# ── Objects ────────────────────────────────────────────────────────
@dataclass
class Object:
    """A connected component extracted from a grid."""

    color: Color               # –1 when multi-colour
    pixels: FrozenSet[Point]
    bbox: BBox

    @property
    def size(self) -> int:
        return len(self.pixels)


# ── Perception ─────────────────────────────────────────────────────
@dataclass
class GridInfo:
    """Full perception snapshot of a grid."""

    grid: Grid
    height: int
    width: int
    colors: Set[Color]
    color_counts: Dict[Color, int]
    background: Color
    objects: List[Object]
    symmetry: Dict[str, bool]


# ── Task / Pair ────────────────────────────────────────────────────
@dataclass
class Pair:
    """One training/test input→output example."""

    input: Grid
    output: Grid


@dataclass
class Task:
    """An ARC-AGI task."""

    task_id: str
    train: List[Pair]
    test_inputs: List[Grid]
    test_outputs: List[Grid]          # empty for blind evaluation

    @classmethod
    def from_json(cls, data: dict, task_id: str = "unknown") -> "Task":
        train = [Pair(p["input"], p["output"]) for p in data["train"]]
        test_inputs = [p["input"] for p in data["test"]]
        test_outputs = [p["output"] for p in data["test"] if "output" in p]
        return cls(
            task_id=task_id,
            train=train,
            test_inputs=test_inputs,
            test_outputs=test_outputs,
        )

    @classmethod
    def from_file(cls, path: str | Path) -> "Task":
        p = Path(path)
        with open(p) as f:
            data = json.load(f)
        return cls.from_json(data, task_id=p.stem)


# ── Transform / Program ───────────────────────────────────────────
@dataclass
class Transform:
    """A named grid transformation (element of the DSL)."""

    name: str
    fn: Callable[[Grid], Optional[Grid]]
    complexity: int = 1

    def __call__(self, grid: Grid) -> Optional[Grid]:
        try:
            return self.fn(grid)
        except Exception:
            return None


# ── Solution ──────────────────────────────────────────────────────
@dataclass
class Solution:
    """Solver output for one task."""

    task_id: str
    predictions: List[Grid]
    transform: Optional[Transform] = None
    confidence: float = 0.0


# ── Synthesis clues ───────────────────────────────────────────────
@dataclass
class PairClues:
    """Clues extracted from comparing one input/output pair."""

    same_size: bool = False
    height_ratio: float = 1.0
    width_ratio: float = 1.0
    input_colors: Set[Color] = field(default_factory=set)
    output_colors: Set[Color] = field(default_factory=set)
    added_colors: Set[Color] = field(default_factory=set)
    removed_colors: Set[Color] = field(default_factory=set)
    color_mapping: Optional[Dict[Color, Color]] = None
    input_object_count: int = 0
    output_object_count: int = 0
