"""Diagnostics for boundary crowding and convergence."""

from __future__ import annotations

from dataclasses import dataclass
from typing import Literal

import numpy as np
from numpy.typing import NDArray

from superellipse import Superellipse
from panel_bie import DoubleLayerLaplace
from harmonic_measure.boundary_data import corner_boundary_data


@dataclass
class RefinementRecord:
    """Record of a single refinement step."""
    round: int = 0
    panels_per_quadrant: int = 0
    m_per_panel: int = 0
    beta: float = 0.0
    n_nodes: int = 0
    arc_len: float = 0.0
    delta_abs: float | None = None
    delta_rel: float | None = None


def corner_arc_len_once(
    n: int,
    a: float,
    c: float,
    *,
    beta: float,
    panels_per_quadrant: int,
    m_per_panel: int,
) -> tuple[float, float, int]:
    """Single discretization solve for corner harmonic mass.

    Parameters
    ----------
    n : int
        Lamé exponent (p = 2n)
    a : float
        Aspect ratio
    c : float
        Corner radius parameter (ρ = c/n)
    beta : float
        Corner refinement parameter
    panels_per_quadrant : int
        Number of panels per quadrant
    m_per_panel : int
        Gauss-Legendre nodes per panel

    Returns
    -------
    arc_len : float
        Corner arc length ≈ 2π u(0)
    u0 : float
        Harmonic measure u(0)
    n_nodes : int
        Total number of boundary nodes
    """
    # Create Lamé curve with exponent 2n
    curve = Superellipse(a=1.0, b=a, p=2*n)
    disc = curve.panel_discretization(
        panels_per_quadrant=panels_per_quadrant,
        nodes_per_panel=m_per_panel,
        beta=beta,
    )

    # Boundary data for corner neighborhoods
    rho = c / n
    delta = rho / 5.0
    g = corner_boundary_data(disc.points, a, rho, delta)

    # Solve
    solver = DoubleLayerLaplace(disc)
    mu = solver.solve(g)

    # Evaluate at origin
    u0 = solver.evaluate(mu, [0.0, 0.0])
    arc_len = 2.0 * np.pi * u0

    return float(arc_len), float(u0), disc.n_nodes


def corner_arc_len_converged(
    n: int,
    a: float,
    c: float,
    *,
    tol: float = 1e-8,
    quality: Literal["fast", "mid", "high"] = "mid",
    max_rounds: int = 10,
    verbose: bool = False,
) -> tuple[float, float, dict, list[RefinementRecord]]:
    """Adaptively refine until corner harmonic mass converges.

    Uses "no whammies" approach: increase discretization parameters
    until the quantity of interest stabilizes.

    Parameters
    ----------
    n : int
        Lamé exponent (p = 2n)
    a : float
        Aspect ratio
    c : float
        Corner radius parameter
    tol : float
        Convergence tolerance (absolute or relative)
    quality : str
        'fast', 'mid', or 'high' - affects initial parameters
    max_rounds : int
        Maximum refinement rounds
    verbose : bool
        Print progress

    Returns
    -------
    arc_len : float
        Converged corner arc length
    u0 : float
        Converged harmonic measure u(0)
    final_params : dict
        Final discretization parameters
    history : list[RefinementRecord]
        Refinement history
    """
    # Initial parameters based on quality
    if quality == "fast":
        ppq, m, beta = 2, 8, 4.0
    elif quality == "mid":
        ppq, m, beta = 4, 16, 8.0
    else:  # high
        ppq, m, beta = 8, 24, 12.0

    history = []
    prev_arc = None

    for round_num in range(1, max_rounds + 1):
        arc_len, u0, n_nodes = corner_arc_len_once(
            n, a, c,
            beta=beta,
            panels_per_quadrant=ppq,
            m_per_panel=m,
        )

        delta_abs = None
        delta_rel = None
        if prev_arc is not None:
            delta_abs = abs(arc_len - prev_arc)
            delta_rel = delta_abs / (abs(arc_len) + 1e-15)

        record = RefinementRecord(
            round=round_num,
            panels_per_quadrant=ppq,
            m_per_panel=m,
            beta=beta,
            n_nodes=n_nodes,
            arc_len=arc_len,
            delta_abs=delta_abs,
            delta_rel=delta_rel,
        )
        history.append(record)

        if verbose:
            print(f"Round {round_num}: N={n_nodes}, arc={arc_len:.10f}, Δ={delta_abs}")

        # Check convergence
        if delta_abs is not None:
            if delta_abs < tol or delta_rel < tol:
                break

        # Refine for next round
        prev_arc = arc_len
        ppq = int(ppq * 1.5)
        m = min(m + 4, 32)
        beta = min(beta + 2, 20)

    final_params = {
        "panels_per_quadrant": ppq,
        "m_per_panel": m,
        "beta": beta,
    }

    return arc_len, u0, final_params, history


def adaptive_convergence(
    n: int,
    a: float,
    c: float,
    tol: float = 1e-10,
    verbose: bool = False,
) -> tuple[float, list[RefinementRecord]]:
    """Convenience wrapper for corner_arc_len_converged.

    Returns just the arc length and history.
    """
    arc, _, _, hist = corner_arc_len_converged(
        n, a, c, tol=tol, quality="mid", verbose=verbose
    )
    return arc, hist


def corner_scaling_diagnostic(
    ns: list[int],
    a: float = 1.0,
    c: float = 2.0,
    tol: float = 1e-8,
    verbose: bool = False,
) -> tuple[NDArray, NDArray, float]:
    """Compute corner harmonic mass across a range of Lamé exponents.

    Fits log(arc_len) ~ slope * log(n) to estimate scaling exponent.

    Parameters
    ----------
    ns : list of int
        Lamé exponents to test
    a : float
        Aspect ratio
    c : float
        Corner radius parameter
    tol : float
        Convergence tolerance
    verbose : bool
        Print progress

    Returns
    -------
    ns : ndarray
        Exponent values
    arc_lens : ndarray
        Corner arc lengths
    slope : float
        Fitted slope (should be ≈ -2 for right-angle corners)
    """
    arc_lens = []
    for n in ns:
        if verbose:
            print(f"Computing n={n}...")
        arc, _ = adaptive_convergence(n, a, c, tol=tol, verbose=False)
        arc_lens.append(arc)

    ns_arr = np.array(ns, dtype=float)
    arcs_arr = np.array(arc_lens)

    # Fit log-log
    slope, intercept = np.polyfit(np.log(ns_arr), np.log(arcs_arr), 1)

    if verbose:
        print(f"Fitted slope: {slope:.4f}")

    return ns_arr, arcs_arr, slope
