"""Laplace equation solvers using boundary integral methods."""

from __future__ import annotations

from dataclasses import dataclass
from typing import Callable

import numpy as np
from numpy.typing import NDArray
from scipy.linalg import solve

from panel_bie.geometry import Geometry, validate_geometry
from panel_bie.kernels import (
    laplace_double_layer_matrix,
    laplace_single_layer_matrix,
)


@dataclass
class DoubleLayerLaplace:
    """Double-layer potential solver for interior Dirichlet problems.

    Solves: Δu = 0 in Ω, u = g on ∂Ω

    Using representation: u(x) = ∫_∂Ω μ(y) ∂G(x,y)/∂n_y ds_y

    The density μ satisfies: (-½I + K)μ = g

    Parameters
    ----------
    geometry : Geometry
        Boundary discretization with points, normals, weights, curvature

    Examples
    --------
    >>> from superellipse import Superellipse
    >>> curve = Superellipse(p=8)
    >>> disc = curve.panel_discretization()
    >>> solver = DoubleLayerLaplace(disc)
    >>> mu = solver.solve(lambda pt: pt[0])  # u = x on boundary
    >>> u_0 = solver.evaluate(mu, [0, 0])
    """

    geometry: Geometry
    _A: NDArray[np.floating] | None = None

    def __post_init__(self):
        validate_geometry(self.geometry)

    @property
    def n_nodes(self) -> int:
        return self.geometry.points.shape[0]

    def assemble(self) -> NDArray[np.floating]:
        """Assemble the discretized operator matrix (-½I + K).

        Returns
        -------
        A : ndarray, shape (N, N)
            Discretized integral operator
        """
        if self._A is not None:
            return self._A

        pts = self.geometry.points
        nrm = self.geometry.normals
        w = self.geometry.weights
        kappa = self.geometry.curvature

        # Build K matrix (double-layer operator)
        K = laplace_double_layer_matrix(
            pts, pts, nrm, w,
            curvature=kappa,
            same_points=True,
        )

        # Form (-½I + K)
        self._A = -0.5 * np.diag(np.ones(self.n_nodes)) + K
        return self._A

    def solve(
        self,
        g: Callable[[NDArray], float] | NDArray[np.floating],
    ) -> NDArray[np.floating]:
        """Solve for density μ given boundary data g.

        Parameters
        ----------
        g : callable or ndarray
            If callable: g(point) returns boundary value at point
            If ndarray: directly provides values at boundary nodes

        Returns
        -------
        mu : ndarray, shape (N,)
            Solution density
        """
        A = self.assemble()

        if callable(g):
            rhs = np.array([g(pt) for pt in self.geometry.points])
        else:
            rhs = np.asarray(g)

        return solve(A, rhs)

    def evaluate(
        self,
        mu: NDArray[np.floating],
        point: NDArray[np.floating] | list[float],
    ) -> float:
        """Evaluate solution at an interior point.

        Parameters
        ----------
        mu : ndarray, shape (N,)
            Solution density from solve()
        point : array-like, shape (2,)
            Interior point

        Returns
        -------
        u : float
            Solution value at point
        """
        point = np.asarray(point)
        pts = self.geometry.points
        nrm = self.geometry.normals
        w = self.geometry.weights

        # Single target, multiple sources
        K_row = laplace_double_layer_matrix(
            point.reshape(1, 2), pts, nrm, w,
            same_points=False,
        )

        return float(K_row @ mu)

    def evaluate_grid(
        self,
        mu: NDArray[np.floating],
        grid: NDArray[np.floating],
        mask: NDArray[np.bool_] | None = None,
    ) -> NDArray[np.floating]:
        """Evaluate solution on a grid of points.

        Parameters
        ----------
        mu : ndarray, shape (N,)
            Solution density
        grid : ndarray, shape (M, 2)
            Evaluation points
        mask : ndarray, shape (M,), optional
            Boolean mask; only evaluate where True

        Returns
        -------
        u : ndarray, shape (M,)
            Solution values (NaN where mask is False)
        """
        if mask is None:
            mask = np.ones(grid.shape[0], dtype=bool)

        u = np.full(grid.shape[0], np.nan)

        if not np.any(mask):
            return u

        pts = self.geometry.points
        nrm = self.geometry.normals
        w = self.geometry.weights

        K = laplace_double_layer_matrix(
            grid[mask], pts, nrm, w,
            same_points=False,
        )

        u[mask] = K @ mu
        return u


@dataclass
class SingleLayerLaplace:
    """Single-layer potential solver for exterior/Neumann problems.

    Uses representation: u(x) = ∫_∂Ω σ(y) G(x,y) ds_y

    Parameters
    ----------
    geometry : Geometry
        Boundary discretization
    """

    geometry: Geometry
    _A: NDArray[np.floating] | None = None

    def __post_init__(self):
        validate_geometry(self.geometry)

    @property
    def n_nodes(self) -> int:
        return self.geometry.points.shape[0]

    def assemble(self) -> NDArray[np.floating]:
        """Assemble the single-layer operator matrix."""
        if self._A is not None:
            return self._A

        pts = self.geometry.points
        w = self.geometry.weights

        self._A = laplace_single_layer_matrix(pts, pts, w)
        # Diagonal needs special treatment (log singularity)
        # For now, use a simple regularization
        for i in range(self.n_nodes):
            self._A[i, i] = 0.0  # Will be refined in future version

        return self._A

    def evaluate(
        self,
        sigma: NDArray[np.floating],
        point: NDArray[np.floating] | list[float],
    ) -> float:
        """Evaluate single-layer potential at a point."""
        point = np.asarray(point)
        pts = self.geometry.points
        w = self.geometry.weights

        S_row = laplace_single_layer_matrix(
            point.reshape(1, 2), pts, w,
        )

        return float(S_row @ sigma)
