"""DSL primitives for ARC-AGI grid transformations.

This is the İrade organ's vocabulary: each primitive is a *selection*
(AX3) — one concrete way to transform a grid.  The synthesis engine
picks which selections to compose.

Primitives are organised by the framework lens they most resemble:
  Geometric  → Mereoloji   (part-whole spatial transforms)
  Colour     → Kavram Sözlüğü (conceptual remapping)
  Object     → Ontoloji    (entity-level operations)
  Pattern    → FOL / Holografik  (rule-based / structural)
"""

from __future__ import annotations

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

from . import grid as G
from .types import Color, Grid, Object, Transform


# =====================================================================
#   § 1  GEOMETRIC PRIMITIVES
# =====================================================================

_GEOMETRIC: List[Transform] = [
    Transform("identity", lambda g: g, complexity=0),
    Transform("rot90", G.rotate_90),
    Transform("rot180", G.rotate_180),
    Transform("rot270", G.rotate_270),
    Transform("flip_h", G.flip_h),
    Transform("flip_v", G.flip_v),
    Transform("transpose", G.transpose),
    Transform("crop", G.crop_to_content),
]


# =====================================================================
#   § 2  GRAVITY PRIMITIVES
# =====================================================================

_GRAVITY: List[Transform] = [
    Transform("gravity_down", G.gravity_down),
    Transform("gravity_up", G.gravity_up),
    Transform("gravity_left", G.gravity_left),
    Transform("gravity_right", G.gravity_right),
]


# =====================================================================
#   § 3  COLOUR PRIMITIVES
# =====================================================================

def _swap_bg_and_fg(g: Grid) -> Grid:
    """Swap background colour with the most-common foreground colour."""
    bg = G.detect_background(g)
    counts = G.color_counts(g)
    non_bg = {c: n for c, n in counts.items() if c != bg}
    if not non_bg:
        return g
    fg = max(non_bg, key=non_bg.get)
    return G.recolor(g, {bg: fg, fg: bg})


def _keep_only_color(color: Color) -> Callable[[Grid], Grid]:
    """Return a function that zeros-out everything except *color*."""
    def fn(g: Grid) -> Grid:
        bg = G.detect_background(g)
        return [[c if c == color else bg for c in row] for row in g]
    return fn


def _remove_color(color: Color) -> Callable[[Grid], Grid]:
    """Return a function that removes *color* (replaces with background)."""
    def fn(g: Grid) -> Grid:
        bg = G.detect_background(g)
        return [[bg if c == color else c for c in row] for row in g]
    return fn


_COLOUR: List[Transform] = [
    Transform("swap_bg_fg", _swap_bg_and_fg),
    Transform("fill_enclosed", G.fill_enclosed),
]


# =====================================================================
#   § 4  OBJECT-LEVEL PRIMITIVES
# =====================================================================

def _extract_largest(g: Grid) -> Optional[Grid]:
    """Crop to the bounding box of the largest foreground object."""
    objs = G.extract_objects(g)
    if not objs:
        return None
    biggest = max(objs, key=lambda o: o.size)
    bg = G.detect_background(g)
    h = biggest.bbox.height
    w = biggest.bbox.width
    out = G.make_grid(h, w, bg)
    for r, c in biggest.pixels:
        out[r - biggest.bbox.r0][c - biggest.bbox.c0] = g[r][c]
    return out


def _extract_smallest(g: Grid) -> Optional[Grid]:
    """Crop to the bounding box of the smallest foreground object."""
    objs = G.extract_objects(g)
    if not objs:
        return None
    smallest = min(objs, key=lambda o: o.size)
    bg = G.detect_background(g)
    h = smallest.bbox.height
    w = smallest.bbox.width
    out = G.make_grid(h, w, bg)
    for r, c in smallest.pixels:
        out[r - smallest.bbox.r0][c - smallest.bbox.c0] = g[r][c]
    return out


def _sort_objects_by_size(g: Grid) -> Grid:
    """Re-arrange objects top-to-bottom by ascending size."""
    bg = G.detect_background(g)
    objs = G.extract_objects(g)
    if not objs:
        return g
    objs_sorted = sorted(objs, key=lambda o: o.size)
    h, w = G.grid_size(g)
    result = G.make_grid(h, w, bg)
    row_cursor = 0
    for obj in objs_sorted:
        sub = G.make_grid(obj.bbox.height, obj.bbox.width, bg)
        for r, c in obj.pixels:
            sub[r - obj.bbox.r0][c - obj.bbox.c0] = g[r][c]
        for r in range(obj.bbox.height):
            for c in range(obj.bbox.width):
                if row_cursor + r < h and c < w:
                    result[row_cursor + r][c] = sub[r][c]
        row_cursor += obj.bbox.height
    return result


def _outline_objects(g: Grid) -> Grid:
    """Replace each object with just its border pixels."""
    bg = G.detect_background(g)
    h, w = G.grid_size(g)
    result = G.make_grid(h, w, bg)
    objs = G.extract_objects(g)
    for obj in objs:
        for r, c in obj.pixels:
            is_border = False
            for dr, dc in G.DIRS4:
                nr, nc = r + dr, c + dc
                if not (0 <= nr < h and 0 <= nc < w) or g[nr][nc] != obj.color:
                    is_border = True
                    break
            if is_border:
                result[r][c] = obj.color
    return result


_OBJECT: List[Transform] = [
    Transform("extract_largest", _extract_largest, complexity=2),
    Transform("extract_smallest", _extract_smallest, complexity=2),
    Transform("sort_objects", _sort_objects_by_size, complexity=3),
    Transform("outline_objects", _outline_objects, complexity=2),
]


# =====================================================================
#   § 5  PATTERN PRIMITIVES
# =====================================================================

def _mirror_h_extend(g: Grid) -> Grid:
    """Extend grid by mirroring horizontally → double width."""
    flipped = G.flip_h(g)
    h, w = G.grid_size(g)
    return [g[r] + flipped[r] for r in range(h)]


def _mirror_v_extend(g: Grid) -> Grid:
    """Extend grid by mirroring vertically → double height."""
    flipped = G.flip_v(g)
    return [row[:] for row in g] + [row[:] for row in flipped]


def _tile_2x2(g: Grid) -> Grid:
    """Tile the grid 2 × 2."""
    return G.tile_grid(g, 2, 2)


def _tile_3x3(g: Grid) -> Grid:
    """Tile the grid 3 × 3."""
    return G.tile_grid(g, 3, 3)


def _scale_2(g: Grid) -> Grid:
    return G.scale_grid(g, 2)


def _scale_3(g: Grid) -> Grid:
    return G.scale_grid(g, 3)


_PATTERN: List[Transform] = [
    Transform("mirror_h_ext", _mirror_h_extend, complexity=2),
    Transform("mirror_v_ext", _mirror_v_extend, complexity=2),
    Transform("tile_2x2", _tile_2x2, complexity=2),
    Transform("tile_3x3", _tile_3x3, complexity=2),
    Transform("scale_2", _scale_2, complexity=2),
    Transform("scale_3", _scale_3, complexity=2),
]


# =====================================================================
#   § 6  COMPOSITION UTILITIES
# =====================================================================

def compose(t1: Transform, t2: Transform) -> Transform:
    """Compose two transforms: first *t1* then *t2*."""
    def fn(g: Grid) -> Optional[Grid]:
        mid = t1(g)
        if mid is None:
            return None
        return t2(mid)
    return Transform(
        name=f"{t1.name} >> {t2.name}",
        fn=fn,
        complexity=t1.complexity + t2.complexity,
    )


def apply_per_object(transform: Transform) -> Transform:
    """Apply *transform* to each object independently; reassemble on canvas."""
    def fn(g: Grid) -> Optional[Grid]:
        bg = G.detect_background(g)
        objs = G.extract_objects(g)
        if not objs:
            return None
        h, w = G.grid_size(g)
        result = G.make_grid(h, w, bg)
        for obj in objs:
            # Extract object sub-grid
            sub = G.make_grid(obj.bbox.height, obj.bbox.width, bg)
            for r, c in obj.pixels:
                sub[r - obj.bbox.r0][c - obj.bbox.c0] = g[r][c]
            # Transform
            transformed = transform(sub)
            if transformed is None:
                return None
            # Place back
            th, tw = G.grid_size(transformed)
            for r in range(th):
                for c in range(tw):
                    nr = obj.bbox.r0 + r
                    nc = obj.bbox.c0 + c
                    if 0 <= nr < h and 0 <= nc < w and transformed[r][c] != bg:
                        result[nr][nc] = transformed[r][c]
        return result
    return Transform(
        name=f"per_object({transform.name})",
        fn=fn,
        complexity=transform.complexity + 2,
    )


# =====================================================================
#   § 7  PARAMETERISED FACTORIES
# =====================================================================

def make_recolor(mapping: Dict[Color, Color]) -> Transform:
    """Create a recolouring transform from an explicit mapping."""
    tag = ",".join(f"{k}->{v}" for k, v in sorted(mapping.items()))
    return Transform(
        name=f"recolor({tag})",
        fn=lambda g: G.recolor(g, mapping),
        complexity=2,
    )


def make_keep_color(color: Color) -> Transform:
    """Create a transform that keeps only *color*."""
    return Transform(
        name=f"keep_color({color})",
        fn=_keep_only_color(color),
        complexity=2,
    )


def make_remove_color(color: Color) -> Transform:
    """Create a transform that removes *color*."""
    return Transform(
        name=f"remove_color({color})",
        fn=_remove_color(color),
        complexity=2,
    )


def make_scale(factor: int) -> Transform:
    """Create a scaling transform."""
    return Transform(
        name=f"scale({factor})",
        fn=lambda g: G.scale_grid(g, factor),
        complexity=2,
    )


def make_fill_enclosed(color: Color) -> Transform:
    """Create a fill-enclosed with specific colour."""
    return Transform(
        name=f"fill_enclosed({color})",
        fn=lambda g: G.fill_enclosed(g, color),
        complexity=2,
    )


# =====================================================================
#   § 8  REGISTRIES
# =====================================================================

# All fixed (parameter-free) primitives
PRIMITIVES: List[Transform] = (
    _GEOMETRIC + _GRAVITY + _COLOUR + _OBJECT + _PATTERN
)

# Lookup by name
REGISTRY: Dict[str, Transform] = {t.name: t for t in PRIMITIVES}


def all_primitives() -> List[Transform]:
    """Return a copy of the full primitives list."""
    return list(PRIMITIVES)
