from __future__ import annotations
import cmath
import math
import random
from typing import List

from .util import poly_eval, poly_eval_dp_ddp, _USE_NUMBA


def halley_refine(coeffs, x, steps=2):
    for _ in range(steps):
        p, dp, ddp = poly_eval_dp_ddp(coeffs, x)
        denom = 2 * dp * dp - p * ddp
        if denom == 0:
            break
        x = x - (2 * p * dp) / denom
    return x


def estimate_multiplicity(coeffs, x) -> int:
    """Estimate local multiplicity m using p, p', p'' at x.
    m ≈ round( Re( p * p'' / p'^2 ) ), clamped to [1..3].
    """
    p, dp, ddp = poly_eval_dp_ddp(coeffs, x)
    if dp == 0:
        return 2
    ratio = (p * ddp) / (dp * dp)
    m = int(round(ratio.real))
    if m < 1:
        m = 1
    if m > 3:
        m = 3
    return m


def halley_refine_multiplicity(coeffs, x, steps=3):
    """Halley-like refinement aware of multiplicity m when m>=2.
    Update: x <- x - m * p / ( m*p' - (m-1) * p*p''/p' ).
    Falls back to standard Halley when estimates are unstable.
    """
    for _ in range(steps):
        p, dp, ddp = poly_eval_dp_ddp(coeffs, x)
        if dp == 0:
            break
        m = estimate_multiplicity(coeffs, x)
        if m <= 1:
            # standard Halley step
            denom = 2 * dp * dp - p * ddp
            if denom == 0:
                break
            x = x - (2 * p * dp) / denom
        else:
            # multiplicity-aware update
            denom = (m * dp) - (m - 1) * (p * ddp / dp)
            if denom == 0:
                break
            x = x - (m * p) / denom
    return x


def durand_kerner(coeffs, iters=400, tol=1e-14, restarts=4):
    n = len(coeffs) - 1
    an = coeffs[-1]
    if an == 0:
        raise ValueError("Leading coefficient is zero.")
    base_R = 1 + max((abs(a) / abs(an) for a in coeffs[:-1]), default=0)

    best_roots = None
    best_res = float("inf")
    for attempt in range(restarts):
        R = base_R * (1 + 0.2 * (random.random() - 0.5))
        roots = [
            R * cmath.exp(2j * math.pi * (k + 0.25 * attempt) / n) for k in range(n)
        ]
        for it in range(iters):
            new_roots = []
            max_delta = 0.0
            for i, z in enumerate(roots):
                denom = 1.0 + 0j
                for j, w in enumerate(roots):
                    if j != i:
                        denom *= z - w
                if denom == 0:
                    z += (random.random() - 0.5) * 1e-8 + 1j * (
                        random.random() - 0.5
                    ) * 1e-8
                    denom = 1.0 + 0j
                    for j, w in enumerate(roots):
                        if j != i:
                            denom *= z - w
                val = poly_eval(coeffs, z)
                z_new = z - val / denom
                new_roots.append(z_new)
                max_delta = max(max_delta, abs(z_new - z))
            roots = new_roots
            if max_delta < tol:
                break
        # residual measure
        res = max(abs(poly_eval(coeffs, r)) for r in roots)
        if res < best_res:
            best_res = res
            best_roots = roots
    return best_roots


def aberth_ehrlich(
    coeffs, iters=400, tol=1e-14, restarts=2, warm_starts: List[complex] | None = None
):
    """
    Aberth–Ehrlich simultaneous iteration with adaptive damping and minimal repulsion:
    z_i <- z_i - alpha_i * p(z_i) / ( p'(z_i) - p(z_i) * sum_{j!=i} 1/(z_i - z_j) )
    where alpha_i in (0,1] chosen to reduce residual and temper steps near clusters.
    """
    n = len(coeffs) - 1
    an = coeffs[-1]
    if an == 0:
        raise ValueError("Leading coefficient is zero.")
    base_R = 1 + max((abs(a) / abs(an) for a in coeffs[:-1]), default=0)

    best_roots = None
    best_res = float("inf")
    for attempt in range(restarts):
        if warm_starts and len(warm_starts) >= 1:
            roots = []
            for ws in warm_starts[: min(len(warm_starts), n)]:
                roots.append(
                    ws
                    + 0.01 * base_R * (random.random() - 0.5)
                    + 1j * 0.01 * base_R * (random.random() - 0.5)
                )
            while len(roots) < n:
                k = len(roots)
                roots.append(
                    base_R * cmath.exp(2j * math.pi * (k + 0.17 * attempt) / n)
                )
        else:
            roots = [
                base_R * cmath.exp(2j * math.pi * (k + 0.37 * attempt) / n)
                for k in range(n)
            ]

        # Precompute reversed coefficients for Horner once (used when Numba off)
        a_rev = list(reversed(coeffs))
        for it in range(iters):
            max_step = 0.0
            new_roots = roots[:]
            # Precompute p and p' at all points
            pvals = []
            dpvals = []
            for z in roots:
                if _USE_NUMBA:
                    p, dp, _ = poly_eval_dp_ddp(coeffs, z)
                else:
                    p = 0j
                    dp = 0j
                    for a in a_rev:
                        dp = dp * z + p
                        p = p * z + a
                pvals.append(p)
                dpvals.append(dp)

            # Precompute 1/(z_i - z_j) matrix for all pairs (i!=j)
            inv = [[0j] * n for _ in range(n)]
            for i in range(n):
                zi = roots[i]
                for j in range(i + 1, n):
                    d = zi - roots[j]
                    if d == 0:
                        val = 0j
                    else:
                        val = 1 / d
                    inv[i][j] = val
                    inv[j][i] = -val

            # Precompute cluster metric per root
            crowd = []
            for i, z in enumerate(roots):
                s = 0.0
                for j, w in enumerate(roots):
                    if j != i:
                        d = abs(z - w)
                        if d > 0:
                            s += 1.0 / d
                crowd.append(s)
            crowd_max = max(crowd) if crowd else 1.0

            for i, z in enumerate(roots):
                p = pvals[i]
                dp = dpvals[i]
                if dp == 0:
                    z = (
                        z
                        + (random.random() - 0.5) * 1e-6
                        + 1j * (random.random() - 0.5) * 1e-6
                    )
                    p = poly_eval(coeffs, z)
                    eps = 1e-8
                    dp = (poly_eval(coeffs, z + eps) - poly_eval(coeffs, z - eps)) / (
                        2 * eps
                    )
                # Aberth correction
                # Use precomputed reciprocal sums
                S = 0j
                row = inv[i]
                for j in range(n):
                    if j != i:
                        S += row[j]
                denom = dp - p * S
                if denom == 0:
                    continue
                delta = p / denom

                # Adaptive damping: smaller alpha if crowding is high
                local_crowd = crowd[i]
                crowd_factor = (local_crowd / crowd_max) if crowd_max > 0 else 0.0
                alpha = 1.0 / (1.0 + 2.0 * crowd_factor)
                alpha = max(0.2, min(1.0, alpha))

                # Backtracking line search on residual
                z_trial = z - alpha * delta
                res0 = abs(p)
                res1 = abs(poly_eval(coeffs, z_trial))
                back = 0
                while res1 > res0 and back < 4:
                    alpha *= 0.5
                    z_trial = z - alpha * delta
                    res1 = abs(poly_eval(coeffs, z_trial))
                    back += 1

                new_roots[i] = z_trial
                max_step = max(max_step, abs(alpha * delta))

            # Minimal repulsion for very close pairs
            eps_close = 1e-10 * (1 + base_R)
            for i in range(n):
                for j in range(i + 1, n):
                    if abs(new_roots[i] - new_roots[j]) < eps_close:
                        jitter = (1e-8 * base_R) * (
                            math.cos(2 * math.pi * random.random())
                            + 1j * math.sin(2 * math.pi * random.random())
                        )
                        new_roots[j] += jitter

            roots = new_roots
            if max_step < tol:
                break

        res = max(abs(poly_eval(coeffs, r)) for r in roots)
        if res < best_res:
            best_res = res
            best_roots = roots
    return best_roots
