"""Program synthesis engine — rule abstraction from I/O pairs.

This is the İrade organ of the solver: it *specifies* which DSL
program should be selected from the space of possibilities (AX3).

Strategy (three stages corresponding to three-phase actualization):
  Stage 1 — Perceive (İlim): extract clues from each I/O pair.
  Stage 2 — Generate candidates (İrade): narrow search via heuristics.
  Stage 3 — Validate (Kudret): test every candidate on *all* pairs.
"""

from __future__ import annotations

from typing import Dict, List, Optional, Set, Tuple

from . import dsl
from . import grid as G
from .types import Color, Grid, Pair, PairClues, Task, Transform


# =====================================================================
#   § 1  PAIR ANALYSIS (İlim phase)
# =====================================================================

def analyze_pair(pair: Pair) -> PairClues:
    """Extract structural clues from a single input→output pair."""
    ih, iw = G.grid_size(pair.input)
    oh, ow = G.grid_size(pair.output)

    in_colors = G.grid_colors(pair.input)
    out_colors = G.grid_colors(pair.output)

    bg_in = G.detect_background(pair.input)
    bg_out = G.detect_background(pair.output)

    # Attempt to infer a bijective colour mapping
    cmap = _infer_color_mapping(pair.input, pair.output) if (ih == oh and iw == ow) else None

    objs_in = G.extract_objects(pair.input, bg_in)
    objs_out = G.extract_objects(pair.output, bg_out)

    return PairClues(
        same_size=(ih == oh and iw == ow),
        height_ratio=oh / ih if ih else 1.0,
        width_ratio=ow / iw if iw else 1.0,
        input_colors=in_colors,
        output_colors=out_colors,
        added_colors=out_colors - in_colors,
        removed_colors=in_colors - out_colors,
        color_mapping=cmap,
        input_object_count=len(objs_in),
        output_object_count=len(objs_out),
    )


def _infer_color_mapping(
    inp: Grid, out: Grid
) -> Optional[Dict[Color, Color]]:
    """If input→output can be explained by a per-pixel colour map, return it."""
    ih, iw = G.grid_size(inp)
    oh, ow = G.grid_size(out)
    if ih != oh or iw != ow:
        return None
    mapping: Dict[Color, Color] = {}
    for r in range(ih):
        for c in range(iw):
            ic, oc = inp[r][c], out[r][c]
            if ic in mapping:
                if mapping[ic] != oc:
                    return None          # not a valid mapping
            else:
                mapping[ic] = oc
    return mapping


# =====================================================================
#   § 2  COMMON ANALYSIS ACROSS PAIRS
# =====================================================================

def common_clues(pairs: List[Pair]) -> PairClues:
    """Merge clues across all pairs — keep only consistent signals."""
    analyses = [analyze_pair(p) for p in pairs]
    if not analyses:
        return PairClues()

    a0 = analyses[0]
    same_size = all(a.same_size for a in analyses)
    h_ratios = {round(a.height_ratio, 2) for a in analyses}
    w_ratios = {round(a.width_ratio, 2) for a in analyses}

    # colour mapping valid only if ALL pairs agree
    cmaps = [a.color_mapping for a in analyses]
    shared_cmap = cmaps[0] if all(m is not None and m == cmaps[0] for m in cmaps) else None

    return PairClues(
        same_size=same_size,
        height_ratio=h_ratios.pop() if len(h_ratios) == 1 else 0.0,
        width_ratio=w_ratios.pop() if len(w_ratios) == 1 else 0.0,
        input_colors=set.union(*(a.input_colors for a in analyses)),
        output_colors=set.union(*(a.output_colors for a in analyses)),
        added_colors=set.intersection(*(a.added_colors for a in analyses)) if analyses else set(),
        removed_colors=set.intersection(*(a.removed_colors for a in analyses)) if analyses else set(),
        color_mapping=shared_cmap,
        input_object_count=a0.input_object_count,   # first pair as proxy
        output_object_count=a0.output_object_count,
    )


# =====================================================================
#   § 3  CANDIDATE GENERATION (İrade phase)
# =====================================================================

def generate_candidates(task: Task) -> List[Transform]:
    """Generate an ordered list of candidate transforms, guided by clues."""
    clues = common_clues(task.train)
    candidates: List[Transform] = []

    # ── Always try all fixed primitives first ──────────────────
    candidates.extend(dsl.all_primitives())

    # ── Colour mapping ─────────────────────────────────────────
    if clues.color_mapping is not None:
        # Only add if it's a non-trivial mapping
        trivial = all(k == v for k, v in clues.color_mapping.items())
        if not trivial:
            candidates.insert(0, dsl.make_recolor(clues.color_mapping))

    # ── Per-colour keep / remove ───────────────────────────────
    for color in clues.input_colors - {G.detect_background(task.train[0].input)}:
        candidates.append(dsl.make_keep_color(color))
        candidates.append(dsl.make_remove_color(color))

    # ── Removed colours → colour-specific extraction ──────────
    for color in clues.removed_colors:
        candidates.append(dsl.make_remove_color(color))

    # ── Size-based heuristics ─────────────────────────────────
    if clues.height_ratio == 2.0 and clues.width_ratio == 2.0:
        candidates.insert(0, dsl.make_scale(2))
    elif clues.height_ratio == 3.0 and clues.width_ratio == 3.0:
        candidates.insert(0, dsl.make_scale(3))

    # ── Fill-enclosed with specific colours ───────────────────
    for color in clues.added_colors:
        candidates.append(dsl.make_fill_enclosed(color))

    # ── Per-object versions of simple transforms ──────────────
    for base in [dsl.REGISTRY.get("rot90"), dsl.REGISTRY.get("flip_h")]:
        if base is not None:
            candidates.append(dsl.apply_per_object(base))

    # ── Crop + colour transforms ──────────────────────────────
    if not clues.same_size:
        crop = dsl.REGISTRY["crop"]
        for color in clues.input_colors - {G.detect_background(task.train[0].input)}:
            candidates.append(dsl.compose(dsl.make_keep_color(color), crop))

    return candidates


# =====================================================================
#   § 4  VALIDATION (Kudret phase)
# =====================================================================

def validates_all(transform: Transform, pairs: List[Pair]) -> bool:
    """Return ``True`` if *transform* maps every pair's input to output."""
    for p in pairs:
        result = transform(p.input)
        if result is None or not G.grids_equal(result, p.output):
            return False
    return True


def validate_score(transform: Transform, pairs: List[Pair]) -> float:
    """Return fraction of pairs where transform produces exact output."""
    if not pairs:
        return 0.0
    correct = sum(
        1 for p in pairs
        if (r := transform(p.input)) is not None and G.grids_equal(r, p.output)
    )
    return correct / len(pairs)


# =====================================================================
#   § 5  MAIN SYNTHESIS
# =====================================================================

def synthesize(
    task: Task,
    *,
    max_compose_depth: int = 2,
    timeout_candidates: int = 5000,
) -> Optional[Transform]:
    """Find a DSL program that maps all training inputs to their outputs.

    Returns the simplest valid transform, or ``None``.
    """
    # Phase 1 — single-step candidates (guided by clues)
    candidates = generate_candidates(task)
    for t in candidates:
        if validates_all(t, task.train):
            return t

    # Phase 2 — two-step compositions
    if max_compose_depth >= 2:
        # Only compose certain promising pairs to keep search bounded
        base_set = candidates[:30]   # cap to manage combinatorics
        composed_count = 0
        for t1 in base_set:
            for t2 in base_set:
                if t1.name == t2.name == "identity":
                    continue
                ct = dsl.compose(t1, t2)
                if validates_all(ct, task.train):
                    return ct
                composed_count += 1
                if composed_count >= timeout_candidates:
                    break
            if composed_count >= timeout_candidates:
                break

    # Phase 3 — three-step compositions (very selective)
    if max_compose_depth >= 3:
        # Only try top candidates
        short = candidates[:10]
        count = 0
        for t1 in short:
            for t2 in short:
                for t3 in short:
                    ct = dsl.compose(dsl.compose(t1, t2), t3)
                    if validates_all(ct, task.train):
                        return ct
                    count += 1
                    if count >= timeout_candidates:
                        return None

    return None
