import numpy as np
from collections import deque
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 x 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}")


# ---------------------------------------------------------------------------
#  Kronecker / self-similar family
# ---------------------------------------------------------------------------

def KroneckerSelfSimilar():
    """output = kron((input != 0).astype(int), input)"""
    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))"""
    def fn(phi):
        mask = (phi != 0).astype(phi.dtype)
        return np.kron(phi, mask)
    return Transform(fn, "KroneckerSelfSimilarInv")


# ---------------------------------------------------------------------------
#  Mirror / kaleidoscope tiling
# ---------------------------------------------------------------------------

def MirrorTileH():
    def fn(phi):
        return np.hstack([phi, np.fliplr(phi)])
    return Transform(fn, "MirrorTileH")

def MirrorTileV():
    def fn(phi):
        return np.vstack([phi, np.flipud(phi)])
    return Transform(fn, "MirrorTileV")

def MirrorTile4Way():
    def fn(phi):
        top = np.hstack([phi, np.fliplr(phi)])
        return np.vstack([top, np.flipud(top)])
    return Transform(fn, "MirrorTile4Way")


# ---------------------------------------------------------------------------
#  Upscale / zoom
# ---------------------------------------------------------------------------

def Upscale(k=2):
    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):
    def fn(phi):
        return phi[::k, ::k].copy()
    return Transform(fn, f"Downscale_{k}x")


# ---------------------------------------------------------------------------
#  Stacking
# ---------------------------------------------------------------------------

def StackH(n=2):
    def fn(phi):
        return np.tile(phi, (1, n))
    return Transform(fn, f"StackH_{n}")

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


# ---------------------------------------------------------------------------
#  Color manipulation
# ---------------------------------------------------------------------------

def RetainColor(color):
    def fn(phi):
        out = np.zeros_like(phi)
        out[phi == color] = color
        return out
    return Transform(fn, f"RetainColor_{color}")

def RemoveColor(color):
    def fn(phi):
        out = phi.copy()
        out[phi == color] = 0
        return out
    return Transform(fn, f"RemoveColor_{color}")

def InvertColors():
    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")


# ---------------------------------------------------------------------------
#  Gravity
# ---------------------------------------------------------------------------

def GravityDown():
    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():
    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")


# ---------------------------------------------------------------------------
#  Overlay / composition
# ---------------------------------------------------------------------------

def OverlayTransparent(background):
    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")


# ---------------------------------------------------------------------------
#  Border / crop helpers
# ---------------------------------------------------------------------------

def CropToContent():
    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():
    def fn(phi):
        return phi.T.copy()
    return Transform(fn, "Transpose")


# ---------------------------------------------------------------------------
#  Object-based transforms (using object_layer)
# ---------------------------------------------------------------------------

def ExtractLargestObject():
    def fn(phi):
        from .object_layer import extract_objects, object_to_cropped_grid
        objs = extract_objects(phi, univalued=True, connectivity=4, without_bg=True)
        if not objs:
            return phi.copy()
        return object_to_cropped_grid(objs[0]).astype(float)
    return Transform(fn, "ExtractLargestObject")

def ExtractSmallestObject():
    def fn(phi):
        from .object_layer import extract_objects, object_to_cropped_grid
        objs = extract_objects(phi, univalued=True, connectivity=4, without_bg=True)
        if not objs:
            return phi.copy()
        return object_to_cropped_grid(objs[-1]).astype(float)
    return Transform(fn, "ExtractSmallestObject")

def ExtractUniqueObject():
    def fn(phi):
        from .object_layer import extract_objects, unique_object, object_to_cropped_grid
        objs = extract_objects(phi, univalued=True, connectivity=4, without_bg=True)
        u = unique_object(objs)
        if u is None:
            return phi.copy()
        return object_to_cropped_grid(u).astype(float)
    return Transform(fn, "ExtractUniqueObject")

def ExtractMostCommonObject():
    def fn(phi):
        from .object_layer import extract_objects, most_common_object, object_to_cropped_grid
        objs = extract_objects(phi, univalued=True, connectivity=4, without_bg=True)
        mc = most_common_object(objs)
        if mc is None:
            return phi.copy()
        return object_to_cropped_grid(mc).astype(float)
    return Transform(fn, "ExtractMostCommonObject")

def KeepLargestObject():
    def fn(phi):
        from .object_layer import extract_objects, object_to_grid, most_common_color
        grid = np.rint(phi).astype(int)
        bg = most_common_color(grid)
        objs = extract_objects(grid, univalued=True, connectivity=4, without_bg=True)
        if not objs:
            return phi.copy()
        return object_to_grid(objs[0], grid.shape, bg=bg).astype(float)
    return Transform(fn, "KeepLargestObject")

def KeepSmallestObject():
    def fn(phi):
        from .object_layer import extract_objects, object_to_grid, most_common_color
        grid = np.rint(phi).astype(int)
        bg = most_common_color(grid)
        objs = extract_objects(grid, univalued=True, connectivity=4, without_bg=True)
        if not objs:
            return phi.copy()
        return object_to_grid(objs[-1], grid.shape, bg=bg).astype(float)
    return Transform(fn, "KeepSmallestObject")

def SortObjectsBySize():
    def fn(phi):
        from .object_layer import extract_objects, paint, most_common_color, canvas
        grid = np.rint(phi).astype(int)
        bg = most_common_color(grid)
        objs = extract_objects(grid, univalued=True, connectivity=4, without_bg=True)
        result = canvas(bg, grid.shape)
        for obj in sorted(objs, key=len, reverse=True):
            result = paint(result, obj)
        return result.astype(float)
    return Transform(fn, "SortObjectsBySize")


# ---------------------------------------------------------------------------
#  Fill / connect / compress
# ---------------------------------------------------------------------------

def FillInterior():
    def fn(phi):
        from .object_layer import extract_objects, object_bbox, object_color
        grid = np.rint(phi).astype(int).copy()
        objs = extract_objects(grid, univalued=True, connectivity=4, without_bg=True)
        for obj in objs:
            rmin, cmin, rmax, cmax = object_bbox(obj)
            color = object_color(obj)
            h, w = rmax - rmin + 1, cmax - cmin + 1
            local_visited = np.zeros((h, w), dtype=bool)
            for _, (r, c) in obj:
                local_visited[r - rmin, c - cmin] = True
            queue = deque()
            exterior = np.zeros((h, w), dtype=bool)
            for i in range(h):
                for j in range(w):
                    if (i == 0 or i == h-1 or j == 0 or j == w-1) and not local_visited[i, j]:
                        exterior[i, j] = True
                        queue.append((i, j))
            while queue:
                r, c = queue.popleft()
                for dr, dc in [(-1,0),(1,0),(0,-1),(0,1)]:
                    nr, nc = r+dr, c+dc
                    if 0 <= nr < h and 0 <= nc < w and not exterior[nr, nc] and not local_visited[nr, nc]:
                        exterior[nr, nc] = True
                        queue.append((nr, nc))
            for i in range(h):
                for j in range(w):
                    if not local_visited[i, j] and not exterior[i, j]:
                        grid[i + rmin, j + cmin] = color
        return grid.astype(float)
    return Transform(fn, "FillInterior")

def ConnectSameColorH():
    def fn(phi):
        grid = np.rint(phi).astype(int).copy()
        h, w = grid.shape
        from .object_layer import most_common_color
        bg = most_common_color(grid)
        for r in range(h):
            colored = [(c, int(grid[r, c])) for c in range(w) if grid[r, c] != bg]
            by_color = {}
            for c, val in colored:
                by_color.setdefault(val, []).append(c)
            for val, cols in by_color.items():
                if len(cols) >= 2:
                    cols.sort()
                    for i in range(len(cols) - 1):
                        for c in range(cols[i], cols[i+1] + 1):
                            grid[r, c] = val
        return grid.astype(float)
    return Transform(fn, "ConnectSameColorH")

def ConnectSameColorV():
    def fn(phi):
        grid = np.rint(phi).astype(int).copy()
        h, w = grid.shape
        from .object_layer import most_common_color
        bg = most_common_color(grid)
        for c in range(w):
            colored = [(r, int(grid[r, c])) for r in range(h) if grid[r, c] != bg]
            by_color = {}
            for r, val in colored:
                by_color.setdefault(val, []).append(r)
            for val, rows in by_color.items():
                if len(rows) >= 2:
                    rows.sort()
                    for i in range(len(rows) - 1):
                        for r in range(rows[i], rows[i+1] + 1):
                            grid[r, c] = val
        return grid.astype(float)
    return Transform(fn, "ConnectSameColorV")

def CompressGrid():
    def fn(phi):
        grid = np.rint(phi).astype(int)
        rows = [grid[0]]
        for i in range(1, grid.shape[0]):
            if not np.array_equal(grid[i], grid[i-1]):
                rows.append(grid[i])
        result = np.array(rows, dtype=int)
        cols = [result[:, 0]]
        for j in range(1, result.shape[1]):
            if not np.array_equal(result[:, j], result[:, j-1]):
                cols.append(result[:, j])
        result = np.column_stack(cols) if cols else result
        return result.astype(float)
    return Transform(fn, "CompressGrid")

def RemoveBlackLines():
    def fn(phi):
        grid = np.rint(phi).astype(int)
        row_mask = np.any(grid != 0, axis=1)
        if row_mask.any():
            grid = grid[row_mask]
        col_mask = np.any(grid != 0, axis=0)
        if col_mask.any():
            grid = grid[:, col_mask]
        return grid.astype(float)
    return Transform(fn, "RemoveBlackLines")

def ColorByProximity():
    def fn(phi):
        grid = np.rint(phi).astype(int).copy()
        from .object_layer import most_common_color
        bg = most_common_color(grid)
        h, w = grid.shape
        dist = np.full((h, w), float('inf'))
        queue = deque()
        for r in range(h):
            for c in range(w):
                if grid[r, c] != bg:
                    dist[r, c] = 0
                    queue.append((r, c))
        while queue:
            r, c = queue.popleft()
            for dr, dc in [(-1,0),(1,0),(0,-1),(0,1)]:
                nr, nc = r+dr, c+dc
                if 0 <= nr < h and 0 <= nc < w and dist[nr, nc] > dist[r, c] + 1:
                    dist[nr, nc] = dist[r, c] + 1
                    grid[nr, nc] = grid[r, c]
                    queue.append((nr, nc))
        return grid.astype(float)
    return Transform(fn, "ColorByProximity")

def DrawBorder():
    def fn(phi):
        grid = np.rint(phi).astype(int)
        from .object_layer import most_common_color
        bg = most_common_color(grid)
        h, w = grid.shape
        result = np.full_like(grid, bg)
        for r in range(h):
            for c in range(w):
                if grid[r, c] != bg:
                    is_border = False
                    for dr, dc in [(-1,0),(1,0),(0,-1),(0,1)]:
                        nr, nc = r+dr, c+dc
                        if nr < 0 or nr >= h or nc < 0 or nc >= w or grid[nr, nc] == bg:
                            is_border = True
                            break
                    if is_border:
                        result[r, c] = grid[r, c]
        return result.astype(float)
    return Transform(fn, "DrawBorder")