import numpy as np
from .solver_core import tile_transform, fill_enclosed, Transform


# ---------------------------------------------------------------------------
#  Existing transforms (cleaned up to import Transform from solver_core)
# ---------------------------------------------------------------------------

def tile_to_target_shifted(shift=(1, 1), tile_factor=3):
    """Tile the input tile_factor×tile_factor times, then roll by shift."""
    def fn(phi):
        h_in, w_in = phi.shape
        out_shape = (h_in * tile_factor, w_in * tile_factor)
        tiled = tile_transform(phi, out_shape)
        tiled = np.roll(tiled, shift=shift, axis=(0, 1))
        return tiled
    return Transform(fn, f"ShiftedTile_s{shift}_f{tile_factor}")


def FillEnclosedHarmonic(boundary_mask=None):
    def fn(phi):
        bm = (phi != 0) if boundary_mask is None else boundary_mask
        return fill_enclosed(phi, bm)
    return Transform(fn, "FillEnclosedHarmonic")


def Rotate(k=1):
    def fn(phi):
        return np.rot90(phi, k)
    return Transform(fn, f"Rotate_{90 * k}")


def Reflect(axis='h'):
    def fn(phi):
        if axis == 'h':
            return np.flipud(phi)
        return np.fliplr(phi)
    return Transform(fn, f"Reflect_{axis}")


def ColorMap(mapping):
    def fn(phi):
        out = phi.copy()
        for k, v in mapping.items():
            out[phi == k] = v
        return out
    return Transform(fn, f"ColorMap_{mapping}")


# ---------------------------------------------------------------------------
#  NEW transforms — Kronecker / self‑similar family
# ---------------------------------------------------------------------------

def KroneckerSelfSimilar():
    """output = kron((input != 0).astype(int), input)

    This is the exact transform for ARC task 007bbfb7 and the general
    family of "use the input's own nonzero pattern as a meta‑layout
    for placing copies of itself".  Works for any 2‑color grid.
    """
    def fn(phi):
        mask = (phi != 0).astype(phi.dtype)
        return np.kron(mask, phi)
    return Transform(fn, "KroneckerSelfSimilar")


def KroneckerSelfSimilarInv():
    """output = kron(input, (input != 0).astype(int))

    Mirror variant — identical result for symmetric inputs, differs
    for asymmetric ones.
    """
    def fn(phi):
        mask = (phi != 0).astype(phi.dtype)
        return np.kron(phi, mask)
    return Transform(fn, "KroneckerSelfSimilarInv")


# ---------------------------------------------------------------------------
#  NEW transforms — mirror / kaleidoscope tiling
# ---------------------------------------------------------------------------

def MirrorTileH():
    """Horizontal mirror tile: [abc] → [abc|cba]"""
    def fn(phi):
        return np.hstack([phi, np.fliplr(phi)])
    return Transform(fn, "MirrorTileH")


def MirrorTileV():
    """Vertical mirror tile: stack input then its vertical reflection."""
    def fn(phi):
        return np.vstack([phi, np.flipud(phi)])
    return Transform(fn, "MirrorTileV")


def MirrorTile4Way():
    """Full kaleidoscope: 2×2 mirror tile (D4 reflections)."""
    def fn(phi):
        top = np.hstack([phi, np.fliplr(phi)])
        return np.vstack([top, np.flipud(top)])
    return Transform(fn, "MirrorTile4Way")


# ---------------------------------------------------------------------------
#  NEW transforms — upscale / zoom
# ---------------------------------------------------------------------------

def Upscale(k=2):
    """Pixel‑repeat upscale: each pixel becomes a k×k block."""
    def fn(phi):
        return np.kron(phi, np.ones((k, k), dtype=phi.dtype))
    return Transform(fn, f"Upscale_{k}x")


def Downscale(k=2):
    """Downsample by taking every k‑th pixel (inverse of Upscale)."""
    def fn(phi):
        return phi[::k, ::k].copy()
    return Transform(fn, f"Downscale_{k}x")


# ---------------------------------------------------------------------------
#  NEW transforms — stacking
# ---------------------------------------------------------------------------

def StackH(n=2):
    """Tile horizontally n times: [A] → [A|A|...|A]."""
    def fn(phi):
        return np.tile(phi, (1, n))
    return Transform(fn, f"StackH_{n}")


def StackV(n=2):
    """Tile vertically n times."""
    def fn(phi):
        return np.tile(phi, (n, 1))
    return Transform(fn, f"StackV_{n}")


# ---------------------------------------------------------------------------
#  NEW transforms — color manipulation
# ---------------------------------------------------------------------------

def RetainColor(color):
    """Keep only pixels of the given color; zero the rest."""
    def fn(phi):
        out = np.zeros_like(phi)
        out[phi == color] = color
        return out
    return Transform(fn, f"RetainColor_{color}")


def RemoveColor(color):
    """Zero out all pixels of the given color."""
    def fn(phi):
        out = phi.copy()
        out[phi == color] = 0
        return out
    return Transform(fn, f"RemoveColor_{color}")


def InvertColors():
    """Swap black (0) with the most common non‑zero color."""
    def fn(phi):
        nonzero = phi[phi != 0]
        if nonzero.size == 0:
            return phi.copy()
        from collections import Counter
        top_color = Counter(nonzero.flatten().astype(int).tolist()).most_common(1)[0][0]
        out = phi.copy()
        mask_zero = (phi == 0)
        mask_top = (phi == top_color)
        out[mask_zero] = top_color
        out[mask_top] = 0
        return out
    return Transform(fn, "InvertColors")


# ---------------------------------------------------------------------------
#  NEW transforms — gravity
# ---------------------------------------------------------------------------

def GravityDown():
    """Non‑zero pixels fall to the bottom within each column."""
    def fn(phi):
        out = np.zeros_like(phi)
        h, w = phi.shape
        for c in range(w):
            col = phi[:, c]
            nonzero = col[col != 0]
            if nonzero.size > 0:
                out[h - nonzero.size:, c] = nonzero
        return out
    return Transform(fn, "GravityDown")


def GravityUp():
    """Non‑zero pixels rise to the top within each column."""
    def fn(phi):
        out = np.zeros_like(phi)
        h, w = phi.shape
        for c in range(w):
            col = phi[:, c]
            nonzero = col[col != 0]
            if nonzero.size > 0:
                out[:nonzero.size, c] = nonzero
        return out
    return Transform(fn, "GravityUp")


# ---------------------------------------------------------------------------
#  NEW transforms — overlay / composition
# ---------------------------------------------------------------------------

def OverlayTransparent(background):
    """Overlay: background pixels are replaced by foreground where foreground != 0."""
    bg = np.array(background, dtype=float)
    def fn(phi):
        out = bg.copy()
        mask = (phi != 0)
        if phi.shape != out.shape:
            p = tile_transform(phi, out.shape)
            m = (p != 0)
            out[m] = p[m]
        else:
            out[mask] = phi[mask]
        return out
    return Transform(fn, "OverlayTransparent")


# ---------------------------------------------------------------------------
#  NEW transforms — border / crop helpers
# ---------------------------------------------------------------------------

def CropToContent():
    """Crop grid to bounding box of non‑zero content."""
    def fn(phi):
        rows = np.any(phi != 0, axis=1)
        cols = np.any(phi != 0, axis=0)
        if not rows.any():
            return phi.copy()
        rmin, rmax = np.where(rows)[0][[0, -1]]
        cmin, cmax = np.where(cols)[0][[0, -1]]
        return phi[rmin:rmax + 1, cmin:cmax + 1].copy()
    return Transform(fn, "CropToContent")


def Transpose():
    """Matrix transpose."""
    def fn(phi):
        return phi.T.copy()
    return Transform(fn, "Transpose")