import numpy as np
from .solver_core import tile_transform, fill_enclosed

class Transform:
    def __init__(self, func, name):
        self.func = func
        self.name = name
    def apply(self, phi):
        return self.func(phi)
    def __repr__(self):
        return f"<Transform {self.name}>"

def tile_to_target_shifted(shift=(1,1), tile_factor=3):
    """
    Tile the input to a canvas tile_factor x tile_factor times, apply a small roll,
    and return the tiled canvas. This guarantees the transform output differs from input.
    """
    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"tile_to_target_shift{shift}")

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}")