import numpy

from ..domain import MaskedDomain


class MaskedSparseSolver:
    """Reference masked Poisson solver via factorized sparse Laplacian."""

    def solve(self, f, domain: MaskedDomain):
        """Solve masked Poisson on active domain pixels.

        f: (H, W) scalar source field
        domain: MaskedDomain with prebuilt Laplacian
        returns: (H, W) scalar potential; zero outside mask
        """
        rhs = domain.prepare_rhs(f)
        if rhs.size == 0:
            return numpy.zeros_like(f, dtype=numpy.float32)

        # MaskedDomain builds a positive diagonal / negative off-diagonal stencil
        # for (-Laplace). Negate the direct solve result so the masked sparse path
        # matches the FFT solver convention used elsewhere: Laplace(u) = rhs.
        sol_active = -domain.factor(rhs)  # (N,) reduced-system potential
        return domain.scatter(sol_active, f)

    def solve_pair(self, f0, f1, domain: MaskedDomain):
        """Solve two masked Poisson RHS with one reused sparse factorization call."""
        rhs_pair = domain.prepare_rhs_pair(f0, f1)
        if rhs_pair.shape[0] == 0:
            empty = numpy.zeros_like(f0, dtype=numpy.float32)
            return empty, empty.copy()

        try:
            sol_pair = -domain.factor(rhs_pair)  # (N, 2) reduced-system potentials
        except Exception:
            sol_pair = -numpy.column_stack(
                [domain.factor(rhs_pair[:, 0]), domain.factor(rhs_pair[:, 1])]
            )
        return domain.scatter_pair(sol_pair, f0)