"""Core HarmonicMeasure class."""

from __future__ import annotations

from dataclasses import dataclass
from typing import Callable

import numpy as np
from numpy.typing import NDArray

from panel_bie import DoubleLayerLaplace
from superellipse import Superellipse, PanelDiscretization


@dataclass
class HarmonicMeasure:
    """Harmonic measure computation for 2D domains.

    Parameters
    ----------
    geometry : Superellipse or PanelDiscretization
        Domain boundary. If Superellipse, will create panel discretization.
    panels_per_quadrant : int
        Panels per quadrant (if geometry is Superellipse)
    nodes_per_panel : int
        Gauss-Legendre nodes per panel
    beta : float
        Corner refinement parameter

    Examples
    --------
    >>> curve = Superellipse(p=16)
    >>> hm = HarmonicMeasure(curve)
    >>> omega = hm.measure_corners([0, 0], radius=0.1)
    """

    geometry: Superellipse | PanelDiscretization
    panels_per_quadrant: int = 4
    nodes_per_panel: int = 16
    beta: float = 8.0

    _disc: PanelDiscretization | None = None
    _solver: DoubleLayerLaplace | None = None

    def __post_init__(self):
        if isinstance(self.geometry, PanelDiscretization):
            self._disc = self.geometry
        else:
            self._disc = self.geometry.panel_discretization(
                panels_per_quadrant=self.panels_per_quadrant,
                nodes_per_panel=self.nodes_per_panel,
                beta=self.beta,
            )
        self._solver = DoubleLayerLaplace(self._disc)

    @property
    def discretization(self) -> PanelDiscretization:
        return self._disc

    @property
    def solver(self) -> DoubleLayerLaplace:
        return self._solver

    def measure_subset(
        self,
        basepoint: NDArray[np.floating] | list[float],
        subset_fn: Callable[[NDArray], bool] | NDArray[np.floating],
        smooth: bool = True,
        delta: float | None = None,
    ) -> float:
        """Compute harmonic measure of a boundary subset.

        Parameters
        ----------
        basepoint : array-like, shape (2,)
            Interior point from which to compute harmonic measure
        subset_fn : callable or ndarray
            If callable: returns True for points in subset
            If ndarray: directly provides boundary data values
        smooth : bool
            If True and subset_fn is callable, use smoothed indicator
        delta : float, optional
            Smoothing width (default: auto-computed)

        Returns
        -------
        omega : float
            Harmonic measure ω(basepoint, subset)
        """
        basepoint = np.asarray(basepoint)

        if callable(subset_fn):
            pts = self._disc.points
            if smooth:
                from harmonic_measure.boundary_data import smooth_indicator
                if delta is None:
                    # Estimate from mesh spacing
                    delta = np.mean(self._disc.weights) / 5
                g = np.array([
                    smooth_indicator(subset_fn(pt), delta)
                    for pt in pts
                ])
            else:
                g = np.array([float(subset_fn(pt)) for pt in pts])
        else:
            g = np.asarray(subset_fn)

        mu = self._solver.solve(g)
        return self._solver.evaluate(mu, basepoint)

    def measure_corners(
        self,
        basepoint: NDArray[np.floating] | list[float],
        radius: float,
        corners: list[tuple[float, float]] | None = None,
    ) -> float:
        """Compute harmonic measure of corner neighborhoods.

        Parameters
        ----------
        basepoint : array-like, shape (2,)
            Interior point
        radius : float
            Neighborhood radius around each corner
        corners : list of (x, y), optional
            Corner locations. Default: (±1, ±a) for Lamé curves

        Returns
        -------
        omega : float
            Total harmonic measure of all corner neighborhoods
        """
        basepoint = np.asarray(basepoint)

        if corners is None:
            # Default: rectangle corners for Lamé curves
            if isinstance(self.geometry, Superellipse):
                a, b = self.geometry.a, self.geometry.b
            else:
                # Estimate from boundary extent
                pts = self._disc.points
                a = np.max(np.abs(pts[:, 0]))
                b = np.max(np.abs(pts[:, 1]))
            corners = [(a, b), (a, -b), (-a, b), (-a, -b)]

        def in_corner(pt):
            for cx, cy in corners:
                if np.sqrt((pt[0] - cx)**2 + (pt[1] - cy)**2) <= radius:
                    return True
            return False

        return self.measure_subset(basepoint, in_corner)

    def density(
        self,
        basepoint: NDArray[np.floating] | list[float],
    ) -> NDArray[np.floating]:
        """Compute harmonic measure density at each boundary node.

        The density ρ satisfies: ω(z, A) ≈ ∫_A ρ(s) ds

        Parameters
        ----------
        basepoint : array-like, shape (2,)
            Interior point

        Returns
        -------
        rho : ndarray, shape (N,)
            Density at each boundary node
        """
        basepoint = np.asarray(basepoint)
        pts = self._disc.points
        n = pts.shape[0]

        # Compute harmonic measure for small neighborhood of each point
        # This is a discrete approximation to the density
        densities = np.zeros(n)
        w = self._disc.weights

        for i in range(n):
            # Delta function at node i
            g = np.zeros(n)
            g[i] = 1.0 / w[i]  # Normalize by weight

            mu = self._solver.solve(g)
            densities[i] = self._solver.evaluate(mu, basepoint)

        return densities

    def plot_density(
        self,
        basepoint: NDArray[np.floating] | list[float],
        ax=None,
        cmap: str = "viridis",
    ):
        """Plot harmonic measure density on the boundary.

        Requires matplotlib.
        """
        import matplotlib.pyplot as plt
        from matplotlib.collections import LineCollection

        basepoint = np.asarray(basepoint)
        pts = self._disc.points
        rho = self.density(basepoint)

        if ax is None:
            fig, ax = plt.subplots(figsize=(8, 8))

        # Create line segments colored by density
        segments = np.array([
            [pts[i], pts[(i + 1) % len(pts)]]
            for i in range(len(pts))
        ])

        lc = LineCollection(segments, cmap=cmap, linewidths=2)
        lc.set_array(rho)
        ax.add_collection(lc)

        ax.autoscale()
        ax.set_aspect("equal")
        ax.set_xlabel("x")
        ax.set_ylabel("y")
        ax.set_title(f"Harmonic measure density from ({basepoint[0]}, {basepoint[1]})")

        plt.colorbar(lc, ax=ax, label="Density")

        return ax
