import numpy
from scipy import sparse
from scipy.sparse import linalg as sparse_linalg

from .config import MaskedPoissonConfig


def normalize_mask(mask, shape):
    """Normalize optional mask to bool array matching grid shape."""
    if mask is None:
        return None
    mask_arr = numpy.asarray(mask, dtype=bool)
    if mask_arr.shape != shape:
        raise ValueError("Mask dimensions should match the Poisson grid")
    return mask_arr


def compute_boundary_mask(active):
    """Return active pixels with at least one inactive 4-neighbor."""
    up = numpy.zeros_like(active)
    down = numpy.zeros_like(active)
    left = numpy.zeros_like(active)
    right = numpy.zeros_like(active)
    up[1:, :] = active[:-1, :]
    down[:-1, :] = active[1:, :]
    left[:, 1:] = active[:, :-1]
    right[:, :-1] = active[:, 1:]
    fully_surrounded = up & down & left & right
    return active & ~fully_surrounded


def build_masked_factor(lap):
    """Factorize the masked Laplacian for fast repeated triangular solves.

    The reduced system is a symmetric M-matrix (graph Laplacian of the active
    pixels with an extra Dirichlet pin), so SuperLU can skip partial pivoting
    (``SymmetricMode`` + ``diag_pivot_thresh=0``), which is ~2x faster than the
    default pivoting factorization. Factoring in float32 halves the memory
    traffic of the triangular solves; relative error against the float64
    reference stays around 1e-7 for these systems.

    ``rhs_dtype`` is returned because ``SuperLU.solve`` refuses to cast a
    float64 right-hand side down to a float32 factor.

    lap: (N, N) symmetric sparse Laplacian (any scipy dtype)
    returns: (solve_callable, rhs_dtype)
    """
    lap_csc = lap.tocsc()
    lap32 = lap_csc.astype(numpy.float32, copy=False)
    for diag_pivot_thresh in (0.0, 1.0):
        try:
            lu = sparse_linalg.splu(
                lap32,
                permc_spec="MMD_AT_PLUS_A",
                diag_pivot_thresh=diag_pivot_thresh,
                options=dict(SymmetricMode=True),
            )
            return lu.solve, numpy.float32
        except Exception:  # noqa: S112 - try the next pivoting strategy
            # Probe: try the next pivoting strategy; the float64 fallback below is
            # returned if every candidate fails, so this failure is not lost.
            continue
    # Last resort: the original pivoting LU on the float64 matrix.
    return sparse_linalg.factorized(lap_csc), numpy.float64


def sample_boundary_coords(boundary_coords, config: MaskedPoissonConfig):
    """Subsample boundary pixel coordinates for capacitance correction."""
    if boundary_coords.size == 0:
        return boundary_coords

    sampled = boundary_coords
    if config.boundary_stride > 1:
        sampled = sampled[:: config.boundary_stride]

    if config.max_boundary > 0 and len(sampled) > config.max_boundary:
        picks = numpy.linspace(0, len(sampled) - 1, num=config.max_boundary, dtype=numpy.int32)
        sampled = sampled[picks]

    return numpy.ascontiguousarray(sampled, dtype=numpy.int32)


class MaskedDomain:
    """Reduced Laplacian system on active mask pixels (2D only)."""

    __slots__ = (
        "_rhs_buffer",
        "_rhs_pair_buffer",
        "_scatter_buffer",
        "_scatter_pair_buffer",
        "active_coords",
        "active_linear_indices",
        "factor",
        "factor_dtype",
        "index_map",
        "lap",
        "mask",
        "num_active",
        "pinned",
        "shape",
    )

    def __init__(
        self,
        shape,
        mask,
        index_map,
        factor,
        lap,
        num_active,
        pinned,
        active_coords,
        active_linear_indices,
        factor_dtype=numpy.float32,
    ):
        self.shape = shape
        self.mask = mask
        self.index_map = index_map
        self.factor = factor
        self.factor_dtype = numpy.dtype(factor_dtype)
        self.lap = lap
        self.num_active = num_active
        self.pinned = pinned
        self.active_coords = active_coords
        self.active_linear_indices = active_linear_indices
        self._rhs_buffer = None
        self._rhs_pair_buffer = None
        self._scatter_buffer = None
        self._scatter_pair_buffer = None

    @classmethod
    def build(cls, mask, spacings):
        """Build masked Laplacian domain from bool mask and grid spacings (dy, dx).

        mask: (H, W) bool; True = active domain pixels
        spacings: (dy, dx) grid spacing tuple
        returns: MaskedDomain with factorized sparse Laplacian
        """
        active = numpy.asarray(mask, dtype=bool)
        h, w = active.shape
        active_coords = numpy.argwhere(active)
        num_active = int(active.sum())
        if num_active == 0:
            return cls(
                shape=active.shape,
                mask=active.copy(),
                index_map=numpy.full(active.shape, -1, dtype=numpy.int32),
                factor=None,
                lap=None,
                num_active=0,
                pinned=0,
                active_coords=active_coords,
                active_linear_indices=numpy.flatnonzero(active),
            )

        index_map = numpy.full(active.shape, -1, dtype=numpy.int32)
        index_map[active] = numpy.arange(num_active, dtype=numpy.int32)
        active_linear_indices = numpy.flatnonzero(active)

        dy, dx = float(spacings[0]), float(spacings[1])
        coeff_y = 1.0 / (dy * dy)
        coeff_x = 1.0 / (dx * dx)

        rows = []
        cols = []
        data = []
        pinned = 0
        for row, col in active_coords:
            idx = int(index_map[row, col])
            if idx == pinned:
                rows.append(idx)
                cols.append(idx)
                data.append(1.0)
                continue

            diag = 0.0
            for drow, dcol, coeff in ((-1, 0, coeff_y), (1, 0, coeff_y), (0, -1, coeff_x), (0, 1, coeff_x)):
                nrow = row + drow
                ncol = col + dcol
                if nrow < 0 or nrow >= h or ncol < 0 or ncol >= w or not active[nrow, ncol]:
                    continue
                rows.append(idx)
                cols.append(int(index_map[nrow, ncol]))
                data.append(-coeff)
                diag += coeff

            if diag == 0.0:
                rows.append(idx)
                cols.append(idx)
                data.append(1.0)
            else:
                rows.append(idx)
                cols.append(idx)
                data.append(diag)

        lap = sparse.csc_matrix((data, (rows, cols)), shape=(num_active, num_active), dtype=numpy.float64)
        factor, factor_dtype = build_masked_factor(lap)
        return cls(
            shape=active.shape,
            mask=active.copy(),
            index_map=index_map,
            factor=factor,
            lap=lap,
            num_active=num_active,
            pinned=pinned,
            active_coords=active_coords,
            active_linear_indices=active_linear_indices,
            factor_dtype=factor_dtype,
        )

    def matches(self, mask):
        """True if this domain was built for the same mask array."""
        mask_arr = numpy.asarray(mask, dtype=bool)
        return self.shape == mask_arr.shape and numpy.array_equal(self.mask, mask_arr)

    def prepare_rhs(self, f):
        """Extract gauge-fixed RHS on active pixels.

        f: (H, W) scalar source field
        returns: (N,) RHS on active pixels; empty if no active pixels

        The returned dtype matches the factorization (see ``factor_dtype``):
        ``SuperLU.solve`` will not downcast a float64 RHS for a float32 factor.
        """
        if self.num_active == 0:
            return numpy.zeros(0, dtype=self.factor_dtype)
        if (
            self._rhs_buffer is None
            or self._rhs_buffer.shape != (self.num_active,)
            or self._rhs_buffer.dtype != self.factor_dtype
        ):
            self._rhs_buffer = numpy.empty(self.num_active, dtype=self.factor_dtype)
        numpy.take(
            numpy.asarray(f, dtype=self.factor_dtype).ravel(),
            self.active_linear_indices,
            out=self._rhs_buffer,
        )
        rhs = self._rhs_buffer
        rhs -= rhs.mean()
        rhs[self.pinned] = 0.0
        return rhs

    def prepare_rhs_pair(self, f0, f1):
        """Extract two gauge-fixed RHS vectors into one reusable (N, 2) buffer."""
        if self.num_active == 0:
            return numpy.zeros((0, 2), dtype=self.factor_dtype)
        if (
            self._rhs_pair_buffer is None
            or self._rhs_pair_buffer.shape != (self.num_active, 2)
            or self._rhs_pair_buffer.dtype != self.factor_dtype
        ):
            self._rhs_pair_buffer = numpy.empty((self.num_active, 2), dtype=self.factor_dtype)
        flat0 = numpy.asarray(f0, dtype=self.factor_dtype).ravel()
        flat1 = numpy.asarray(f1, dtype=self.factor_dtype).ravel()
        numpy.take(flat0, self.active_linear_indices, out=self._rhs_pair_buffer[:, 0])
        numpy.take(flat1, self.active_linear_indices, out=self._rhs_pair_buffer[:, 1])
        self._rhs_pair_buffer -= self._rhs_pair_buffer.mean(axis=0, keepdims=True)
        self._rhs_pair_buffer[self.pinned, :] = 0.0
        return self._rhs_pair_buffer

    def scatter(self, sol_active, template):
        """Map reduced solution back to full grid.

        sol_active: (N,) potential on active pixels
        template: (H, W) array used for output shape/dtype
        returns: (H, W) potential; zero outside mask
        """
        if self._scatter_buffer is None or self._scatter_buffer.shape != template.shape:
            self._scatter_buffer = numpy.zeros(template.shape, dtype=numpy.float32)
        else:
            self._scatter_buffer.fill(0)
        if sol_active.size:
            self._scatter_buffer.ravel()[self.active_linear_indices] = numpy.asarray(sol_active, dtype=numpy.float32)
        return self._scatter_buffer.copy()

    def scatter_pair(self, sol_active_pair, template):
        """Map two reduced solutions back to a pair of full-grid outputs."""
        if self._scatter_pair_buffer is None or self._scatter_pair_buffer.shape != (2,) + template.shape:
            self._scatter_pair_buffer = numpy.zeros((2,) + template.shape, dtype=numpy.float32)
        else:
            self._scatter_pair_buffer.fill(0)
        if sol_active_pair.size:
            values = numpy.asarray(sol_active_pair, dtype=numpy.float32)
            self._scatter_pair_buffer[0].ravel()[self.active_linear_indices] = values[:, 0]
            self._scatter_pair_buffer[1].ravel()[self.active_linear_indices] = values[:, 1]
        return self._scatter_pair_buffer[0].copy(), self._scatter_pair_buffer[1].copy()

    def residual(self, field, rhs_active):
        """Compute masked Laplacian residual L u - rhs on active pixels.

        field: (H, W) full-grid field
        rhs_active: (N,) RHS on active pixels (or None for homogeneous)
        returns: (N,) residual vector
        """
        vec = numpy.asarray(field, dtype=numpy.float64)[self.mask]
        residual = self.lap.dot(vec)
        if rhs_active is not None:
            residual = residual - rhs_active
        return residual

    def to_cache_dict(self, capfft=None):
        """Legacy dict view for tests and facade compatibility."""
        return {
            "shape": self.shape,
            "mask": self.mask,
            "index_map": self.index_map,
            "factor": self.factor,
            "factor_dtype": self.factor_dtype,
            "lap": self.lap,
            "num_active": self.num_active,
            "pinned": self.pinned,
            "active_coords": self.active_coords,
            "capfft": capfft,
        }