"""
Object extraction and manipulation primitives for ARC-AGI tasks.

Provides connected-component extraction, color-based splitting,
list reduction (largest/smallest/most_common), spatial queries,
and composition operations (overlay/paint/underpaint).
"""
import numpy as np
from collections import Counter, deque


# ---------------------------------------------------------------------------
#  Connected component extraction
# ---------------------------------------------------------------------------

def _flood_fill(grid, start, visited, connectivity=4, univalued=True):
    """BFS flood fill from start. Returns set of (color, (r, c)) cells."""
    h, w = grid.shape
    r0, c0 = start
    seed_color = int(grid[r0, c0])
    comp = set()
    queue = deque([(r0, c0)])
    visited[r0, c0] = True

    if connectivity == 8:
        deltas = [(-1,-1),(-1,0),(-1,1),(0,-1),(0,1),(1,-1),(1,0),(1,1)]
    else:
        deltas = [(-1,0),(1,0),(0,-1),(0,1)]

    while queue:
        r, c = queue.popleft()
        val = int(grid[r, c])
        if univalued and val != seed_color:
            continue
        comp.add((val, (r, c)))
        for dr, dc in deltas:
            nr, nc = r + dr, c + dc
            if 0 <= nr < h and 0 <= nc < w and not visited[nr, nc]:
                nval = int(grid[nr, nc])
                if univalued:
                    if nval == seed_color:
                        visited[nr, nc] = True
                        queue.append((nr, nc))
                else:
                    visited[nr, nc] = True
                    queue.append((nr, nc))
    return comp


def extract_objects(grid, univalued=True, connectivity=4, without_bg=True):
    """Extract connected components from grid.

    Args:
        grid: 2D numpy array (int)
        univalued: if True, each component is single-color
        connectivity: 4 or 8
        without_bg: if True, skip the most common color (background)

    Returns:
        list of objects, each object is a set of (color, (row, col))
        sorted by size descending
    """
    grid = np.array(grid, dtype=int)
    h, w = grid.shape
    bg = most_common_color(grid) if without_bg else -1
    visited = np.zeros((h, w), dtype=bool)
    objects = []

    for r in range(h):
        for c in range(w):
            if visited[r, c]:
                continue
            val = int(grid[r, c])
            if val == bg:
                visited[r, c] = True
                continue
            comp = _flood_fill(grid, (r, c), visited, connectivity, univalued)
            if comp:
                objects.append(comp)

    objects.sort(key=len, reverse=True)
    return objects


def split_by_color(grid, without_bg=True):
    """Split grid into per-color masks. Returns list of (color, grid) pairs
    where each grid has only that color's pixels (rest = 0)."""
    grid = np.array(grid, dtype=int)
    bg = most_common_color(grid) if without_bg else -1
    colors = sorted(set(grid.flatten()) - {bg})
    result = []
    for c in colors:
        mask_grid = np.zeros_like(grid)
        mask_grid[grid == c] = c
        result.append((c, mask_grid))
    return result


# ---------------------------------------------------------------------------
#  Object to grid conversion
# ---------------------------------------------------------------------------

def object_to_grid(obj, shape, bg=0):
    """Render an object (set of (color, (r,c))) onto a grid of given shape."""
    grid = np.full(shape, bg, dtype=int)
    for color, (r, c) in obj:
        if 0 <= r < shape[0] and 0 <= c < shape[1]:
            grid[r, c] = color
    return grid


def object_to_cropped_grid(obj, bg=0):
    """Render object cropped to its bounding box."""
    if not obj:
        return np.array([[bg]], dtype=int)
    rows = [r for _, (r, c) in obj]
    cols = [c for _, (r, c) in obj]
    rmin, rmax = min(rows), max(rows)
    cmin, cmax = min(cols), max(cols)
    h, w = rmax - rmin + 1, cmax - cmin + 1
    grid = np.full((h, w), bg, dtype=int)
    for color, (r, c) in obj:
        grid[r - rmin, c - cmin] = color
    return grid


def normalize_object(obj):
    """Shift object so its top-left corner is at (0, 0)."""
    if not obj:
        return obj
    rows = [r for _, (r, c) in obj]
    cols = [c for _, (r, c) in obj]
    rmin, cmin = min(rows), min(cols)
    return {(color, (r - rmin, c - cmin)) for color, (r, c) in obj}


def shift_object(obj, dr, dc):
    """Shift all cells by (dr, dc)."""
    return {(color, (r + dr, c + dc)) for color, (r, c) in obj}


# ---------------------------------------------------------------------------
#  Object queries
# ---------------------------------------------------------------------------

def object_color(obj):
    """Color of a univalued object."""
    colors = {c for c, _ in obj}
    if len(colors) == 1:
        return colors.pop()
    return max(colors, key=lambda c: sum(1 for cc, _ in obj if cc == c))


def object_size(obj):
    return len(obj)


def object_bbox(obj):
    """Returns (rmin, cmin, rmax, cmax)."""
    rows = [r for _, (r, c) in obj]
    cols = [c for _, (r, c) in obj]
    return min(rows), min(cols), max(rows), max(cols)


def object_height(obj):
    rmin, _, rmax, _ = object_bbox(obj)
    return rmax - rmin + 1


def object_width(obj):
    _, cmin, _, cmax = object_bbox(obj)
    return cmax - cmin + 1


def object_center(obj):
    rows = [r for _, (r, c) in obj]
    cols = [c for _, (r, c) in obj]
    return (sum(rows) / len(rows), sum(cols) / len(cols))


# ---------------------------------------------------------------------------
#  List reducers
# ---------------------------------------------------------------------------

def largest_object(objects):
    """Return the largest object by cell count."""
    return max(objects, key=len) if objects else None


def smallest_object(objects):
    """Return the smallest object by cell count."""
    return min(objects, key=len) if objects else None


def most_common_object(objects):
    """Return the object whose normalized shape appears most frequently."""
    if not objects:
        return None
    normed = [frozenset(normalize_object(o)) for o in objects]
    counter = Counter(normed)
    most_common_shape = counter.most_common(1)[0][0]
    for o, n in zip(objects, normed):
        if n == most_common_shape:
            return o
    return objects[0]


def unique_object(objects):
    """If exactly one unique normalized shape exists, return it. Else None."""
    normed = [frozenset(normalize_object(o)) for o in objects]
    counter = Counter(normed)
    uniques = [shape for shape, count in counter.items() if count == 1]
    if len(uniques) == 1:
        for o, n in zip(objects, normed):
            if n == uniques[0]:
                return o
    return None


def filter_by_color(objects, color):
    """Keep only objects of the given color."""
    return [o for o in objects if object_color(o) == color]


def filter_by_size(objects, size):
    """Keep only objects of the given size."""
    return [o for o in objects if len(o) == size]


# ---------------------------------------------------------------------------
#  Color utilities
# ---------------------------------------------------------------------------

def most_common_color(grid):
    """Most frequent color in the grid (= background)."""
    grid = np.array(grid, dtype=int)
    counts = Counter(grid.flatten().tolist())
    return counts.most_common(1)[0][0]


def least_common_color(grid):
    """Least frequent color in the grid."""
    grid = np.array(grid, dtype=int)
    counts = Counter(grid.flatten().tolist())
    return counts.most_common()[-1][0]


def palette(grid):
    """Set of all colors in grid."""
    return set(np.array(grid, dtype=int).flatten().tolist())


def color_normalize(grid):
    """Remap colors by frequency: most common -> 0, next -> 1, etc."""
    grid = np.array(grid, dtype=int)
    counts = Counter(grid.flatten().tolist())
    ranked = [c for c, _ in counts.most_common()]
    remap = {c: i for i, c in enumerate(ranked)}
    return np.vectorize(remap.get)(grid)


# ---------------------------------------------------------------------------
#  Composition / overlay
# ---------------------------------------------------------------------------

def paint(grid, obj):
    """Paint object onto grid. Object cells OVERWRITE grid cells."""
    result = np.array(grid, dtype=int).copy()
    for color, (r, c) in obj:
        if 0 <= r < result.shape[0] and 0 <= c < result.shape[1]:
            result[r, c] = color
    return result


def underpaint(grid, obj):
    """Paint object onto grid, but ONLY where grid has background color."""
    result = np.array(grid, dtype=int).copy()
    bg = most_common_color(result)
    for color, (r, c) in obj:
        if 0 <= r < result.shape[0] and 0 <= c < result.shape[1]:
            if result[r, c] == bg:
                result[r, c] = color
    return result


def overlay_grids(base, foreground):
    """Overlay foreground onto base. Foreground non-zero pixels overwrite."""
    base = np.array(base, dtype=int).copy()
    fg = np.array(foreground, dtype=int)
    h = min(base.shape[0], fg.shape[0])
    w = min(base.shape[1], fg.shape[1])
    mask = fg[:h, :w] != 0
    base[:h, :w][mask] = fg[:h, :w][mask]
    return base


def cover(grid, obj):
    """Erase object from grid (replace with background color)."""
    result = np.array(grid, dtype=int).copy()
    bg = most_common_color(result)
    for _, (r, c) in obj:
        if 0 <= r < result.shape[0] and 0 <= c < result.shape[1]:
            result[r, c] = bg
    return result


def canvas(bg_color, shape):
    """Create a blank grid filled with bg_color."""
    return np.full(shape, bg_color, dtype=int)