import numpy as np
from collections import Counter, deque


def initialize_potential(grid):
    return np.array(grid, dtype=float)


def discrete_gradient(phi):
    gx = np.zeros_like(phi)
    gy = np.zeros_like(phi)
    gx[:-1, :] = phi[1:, :] - phi[:-1, :]
    gy[:, :-1] = phi[:, 1:] - phi[:, :-1]
    mag = np.sqrt(gx**2 + gy**2)
    return gx, gy, mag


def dirichlet_energy(phi):
    _, _, mag = discrete_gradient(phi)
    return np.sum(mag**2)


def sigma_l1(phi_current, phi_target):
    return np.sum(np.abs(phi_target - phi_current))


def tile_transform(phi_in, out_shape):
    h_in, w_in = phi_in.shape
    h_out, w_out = out_shape
    out = np.zeros(out_shape, dtype=float)
    for i in range(h_out):
        for j in range(w_out):
            out[i, j] = phi_in[i % h_in, j % w_in]
    return out


def fill_enclosed(phi, boundary_mask=None):
    """
    Fill zero regions that are fully enclosed by non-zero boundary.
    Returns a new array with enclosed holes filled with modal neighbor color.
    """
    arr = np.array(phi, dtype=int)
    h, w = arr.shape
    if boundary_mask is None:
        boundary = (arr != 0)
    else:
        boundary = np.array(boundary_mask, dtype=bool)

    visited = np.zeros((h, w), dtype=bool)
    filled = arr.copy()

    def flood(start_r, start_c):
        q = deque()
        q.append((start_r, start_c))
        comp = []
        touches_border = False
        visited[start_r, start_c] = True
        while q:
            r, c = q.popleft()
            comp.append((r, c))
            if r == 0 or c == 0 or r == h-1 or c == w-1:
                touches_border = True
            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 visited[nr, nc]:
                    if arr[nr, nc] == 0:
                        visited[nr, nc] = True
                        q.append((nr, nc))
        return comp, touches_border

    for i in range(h):
        for j in range(w):
            if arr[i, j] == 0 and not visited[i, j]:
                comp, touches_border = flood(i, j)
                if not touches_border:
                    neighbor_vals = []
                    for (r, c) in comp:
                        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:
                                v = arr[nr, nc]
                                if v != 0:
                                    neighbor_vals.append(int(v))
                    if neighbor_vals:
                        mode_color = Counter(neighbor_vals).most_common(1)[0][0]
                    else:
                        mode_color = 1
                    for (r, c) in comp:
                        filled[r, c] = mode_color
    return filled


class Transform:
    def __init__(self, func, name, params=None):
        self.func = func
        self.name = name
        self.params = params or {}

    def apply(self, phi):
        return self.func(phi, **self.params)

    def compose(self, other):
        def composed(phi):
            return other.apply(self.apply(phi))
        return Transform(lambda p: composed(p), f"{self.name}∘{other.name}")

    def __repr__(self):
        return f"<Transform {self.name}>"


def beam_minimize(phi_in, phi_target, atomic_library, beam_width=4, max_depth=4, lock_coeff=0.1):
    identity = Transform(lambda p: p, "Id")
    beam = [(identity, phi_in, 0.0)]
    best = None

    for depth in range(max_depth):
        candidates = []
        for T_cur, _, _ in beam:
            for T_atomic in atomic_library:
                T_new = T_cur.compose(T_atomic)
                phi_new = T_new.apply(phi_in)
                residue = sigma_l1(phi_new, phi_target)
                energy = dirichlet_energy(phi_new)
                score = residue + lock_coeff * energy
                candidates.append((T_new, phi_new, score))

        if not candidates:
            break

        candidates.sort(key=lambda x: x[2])
        beam = candidates[:beam_width]
        best = beam[0]

        if sigma_l1(best[1], phi_target) == 0:
            break

    return best[0], best[1]