import numpy
from scipy import linalg as dense_linalg

from ...backend_availability import CUPY_AVAILABLE, cp
from ..config import MaskedPoissonConfig
from ..domain import MaskedDomain, compute_boundary_mask, sample_boundary_coords


class CapacitanceFFTSolver:
    """Masked Poisson via rectangular FFT bulk solve + capacitance boundary correction."""

    def __init__(self, config: MaskedPoissonConfig, bulk_solver):
        self.config = config
        self.bulk_solver = bulk_solver
        self._capfft_cache = None
        self._capfft_domain = None

    def _use_cupy_backend(self):
        return (
            getattr(self.bulk_solver, "poisson_backend", "cpu") == "cupy"
            and CUPY_AVAILABLE
        )

    def _bulk_solve(self, rhs):
        """Solve one full-grid Poisson problem, using CuPy when configured."""
        if not self._use_cupy_backend():
            return self.bulk_solver.solve(rhs)

        if not getattr(self.bulk_solver, "ready", False):
            self.bulk_solver.prepare(False)

        rhs_gpu = cp.asarray(rhs, dtype=cp.float32)
        field_gpu = self.bulk_solver._solve_gpu_internal(rhs_gpu)
        cp.cuda.get_current_stream().synchronize()
        return cp.asnumpy(field_gpu)

    def prepare(self, domain: MaskedDomain):
        """Precompute capacitance matrix for boundary correction.

        domain: MaskedDomain with built Laplacian
        returns: capacitance cache dict (also stored internally)
        """
        self._capfft_domain = domain
        self._capfft_cache = self._build_capfft_cache(domain)
        return self._capfft_cache

    def _build_capfft_cache(self, domain: MaskedDomain):
        active = domain.mask
        boundary_coords = numpy.argwhere(compute_boundary_mask(active))
        boundary_coords = sample_boundary_coords(boundary_coords, self.config)
        if boundary_coords.size == 0:
            pinned_rc = numpy.array(domain.active_coords[domain.pinned], dtype=numpy.int32)
            boundary_coords = pinned_rc.reshape(1, 2)

        boundary_active_indices = domain.index_map[boundary_coords[:, 0], boundary_coords[:, 1]]
        boundary_count = int(len(boundary_active_indices))

        cap_matrix = numpy.zeros((boundary_count, boundary_count), dtype=numpy.float64)
        zero_rhs = numpy.zeros(domain.num_active, dtype=numpy.float64)
        for col, (row, col_idx) in enumerate(boundary_coords):
            charge = numpy.zeros(active.shape, dtype=numpy.float32)
            charge[row, col_idx] = 1.0
            response = numpy.asarray(self._bulk_solve(charge), dtype=numpy.float64)
            residual = domain.residual(response, zero_rhs)
            cap_matrix[:, col] = residual[boundary_active_indices]

        reg_scale = self.config.reg
        if reg_scale <= 0:
            reg_scale = 1e-8
        reg = reg_scale * max(1.0, float(numpy.linalg.norm(cap_matrix, ord=numpy.inf)))
        cap_matrix_reg = cap_matrix + reg * numpy.eye(boundary_count, dtype=numpy.float64)
        lu, piv = dense_linalg.lu_factor(cap_matrix_reg)

        return {
            "boundary_coords": boundary_coords,
            "boundary_active_indices": boundary_active_indices,
            "boundary_count": boundary_count,
            "cap_matrix": cap_matrix,
            "cap_matrix_reg": cap_matrix_reg,
            "cap_lu": lu,
            "cap_piv": piv,
            "reg": reg,
        }

    def get_cache(self, domain: MaskedDomain):
        """Return cached capacitance data, building if needed."""
        if self._capfft_cache is None or self._capfft_domain is not domain:
            self.prepare(domain)
        return self._capfft_cache

    def solve(self, f, domain: MaskedDomain):
        """Solve masked Poisson with capacitance-corrected rectangular FFT.

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

        rhs_full = numpy.zeros_like(f, dtype=numpy.float32)
        rhs_full[domain.mask] = rhs_active.astype(numpy.float32)

        field = numpy.asarray(self._bulk_solve(rhs_full), dtype=numpy.float64)
        cap_cache = self.get_cache(domain)
        if not cap_cache or cap_cache["boundary_count"] == 0:
            field_active = field[domain.mask]
            if field_active.size:
                field_active = field_active.copy()
                field_active -= field_active[domain.pinned]
            return domain.scatter(field_active, f)

        boundary_rows = cap_cache["boundary_active_indices"]
        rhs_scale = max(1.0, float(numpy.linalg.norm(rhs_active)))
        for _ in range(self.config.max_iters):
            residual = domain.residual(field, rhs_active)
            boundary_residual = residual[boundary_rows]
            if float(numpy.linalg.norm(boundary_residual)) <= self.config.tol * rhs_scale:
                break

            alpha = dense_linalg.lu_solve((cap_cache["cap_lu"], cap_cache["cap_piv"]), boundary_residual)
            charges = numpy.zeros_like(f, dtype=numpy.float32)
            coords = cap_cache["boundary_coords"]
            charges[coords[:, 0], coords[:, 1]] = alpha.astype(numpy.float32)
            correction = numpy.asarray(self._bulk_solve(charges), dtype=numpy.float64)
            field -= correction

        field_active = field[domain.mask]
        if field_active.size:
            field_active = field_active.copy()
            field_active -= field_active[domain.pinned]
        return domain.scatter(field_active, f)