"""Exact polynomial conversions and univariate certificates."""

from __future__ import annotations

from dataclasses import dataclass
from fractions import Fraction
from functools import lru_cache
from math import factorial
from typing import Mapping, Sequence

from flint import fmpq, fmpq_poly

from .certificates import CertificationError

__all__ = [
    "ChebyshevLowerBoundCertificate",
    "certify_chebyshev_lower_bound_fraction",
    "chebyshev_to_power_fraction",
    "power_to_binomial_fraction",
]


@dataclass(frozen=True)
class ChebyshevLowerBoundCertificate:
    degree: int
    lower_bound: Fraction
    value_at_minus_one: Fraction
    value_at_one: Fraction
    enclosed_roots: int


def _add(
    left: Sequence[Fraction],
    right: Sequence[Fraction],
) -> tuple[Fraction, ...]:
    size = max(len(left), len(right))
    return tuple(
        (left[index] if index < len(left) else Fraction(0))
        + (right[index] if index < len(right) else Fraction(0))
        for index in range(size)
    )


def _scale(
    polynomial: Sequence[Fraction],
    scalar: Fraction,
) -> tuple[Fraction, ...]:
    return tuple(scalar * coefficient for coefficient in polynomial)


@lru_cache(maxsize=None)
def _chebyshev_basis(
    maximum_degree: int,
) -> tuple[tuple[Fraction, ...], ...]:
    if maximum_degree < 0:
        raise CertificationError("Chebyshev degree must be nonnegative")
    if maximum_degree == 0:
        return ((Fraction(1),),)
    basis = [
        (Fraction(1),),
        (Fraction(0), Fraction(1)),
    ]
    for _degree in range(2, maximum_degree + 1):
        twice_shifted = (
            Fraction(0),
            *(
                2 * coefficient
                for coefficient in basis[-1]
            ),
        )
        basis.append(
            _add(twice_shifted, _scale(basis[-2], Fraction(-1)))
        )
    return tuple(basis)


def chebyshev_to_power_fraction(
    coefficients: Sequence[Fraction | int],
) -> tuple[Fraction, ...]:
    """Convert ``sum a_k T_k(x)`` to ascending exact power coefficients."""
    exact = tuple(Fraction(value) for value in coefficients)
    if not exact:
        raise CertificationError("Chebyshev polynomial needs coefficients")
    basis = _chebyshev_basis(len(exact) - 1)
    power = (Fraction(0),)
    for coefficient, mode in zip(exact, basis, strict=True):
        power = _add(power, _scale(mode, coefficient))
    while len(power) > 1 and power[-1] == 0:
        power = power[:-1]
    return power


@lru_cache(maxsize=None)
def _stirling_second(power: int, order: int) -> int:
    if power < 0 or order < 0:
        raise CertificationError("Stirling indices must be nonnegative")
    if power == order == 0:
        return 1
    if power == 0 or order == 0 or order > power:
        return 0
    return (
        _stirling_second(power - 1, order - 1)
        + order * _stirling_second(power - 1, order)
    )


def power_to_binomial_fraction(
    terms: Mapping[tuple[int, ...], Fraction | int],
) -> dict[tuple[int, ...], Fraction]:
    """Convert multivariate power coefficients to a binomial basis.

    The input represents ``sum_e c_e prod_i x_i**e_i``. The output
    represents the same polynomial as
    ``sum_k a_k prod_i binomial(x_i, k_i)``.
    """
    if not terms:
        return {}
    dimensions = {len(exponents) for exponents in terms}
    if len(dimensions) != 1:
        raise CertificationError("power terms have inconsistent dimensions")
    if any(
        exponent < 0
        for exponents in terms
        for exponent in exponents
    ):
        raise CertificationError("power exponents must be nonnegative")

    converted = {
        exponents: Fraction(coefficient)
        for exponents, coefficient in terms.items()
        if coefficient
    }
    variable_count = next(iter(dimensions))
    for position in range(variable_count):
        next_terms: dict[tuple[int, ...], Fraction] = {}
        for exponents, coefficient in converted.items():
            power = exponents[position]
            for order in range(power + 1):
                multiplier = _stirling_second(power, order) * factorial(order)
                if not multiplier:
                    continue
                target = list(exponents)
                target[position] = order
                key = tuple(target)
                next_terms[key] = (
                    next_terms.get(key, Fraction(0))
                    + coefficient * multiplier
                )
        converted = {
            exponents: coefficient
            for exponents, coefficient in next_terms.items()
            if coefficient
        }
    return converted


def certify_chebyshev_lower_bound_fraction(
    coefficients: Sequence[Fraction | int],
    lower_bound: Fraction | int,
    *,
    label: str = "Chebyshev polynomial",
) -> ChebyshevLowerBoundCertificate:
    """Certify ``sum a_k T_k(x) > lower_bound`` on ``[-1, 1]``."""
    exact = tuple(Fraction(value) for value in coefficients)
    if not exact:
        raise CertificationError(f"{label} needs coefficients")
    bound = Fraction(lower_bound)
    value_at_one = sum(exact)
    value_at_minus_one = sum(
        coefficient if index % 2 == 0 else -coefficient
        for index, coefficient in enumerate(exact)
    )
    if value_at_minus_one <= bound or value_at_one <= bound:
        raise CertificationError(
            f"{label} endpoint does not exceed {bound}: "
            f"{value_at_minus_one}, {value_at_one}"
        )

    shifted = list(exact)
    shifted[0] -= bound
    power = chebyshev_to_power_fraction(shifted)
    if len(power) == 1:
        if power[0] <= 0:
            raise CertificationError(
                f"{label} constant gap is nonpositive: {power[0]}"
            )
        return ChebyshevLowerBoundCertificate(
            degree=0,
            lower_bound=bound,
            value_at_minus_one=value_at_minus_one,
            value_at_one=value_at_one,
            enclosed_roots=0,
        )

    polynomial = fmpq_poly(
        [
            fmpq(coefficient.numerator, coefficient.denominator)
            for coefficient in power
        ]
    )
    roots = polynomial.complex_roots()
    enclosed_roots = sum(multiplicity for _root, multiplicity in roots)
    if enclosed_roots != polynomial.degree():
        raise CertificationError(
            f"{label} root enclosure incomplete: "
            f"{enclosed_roots} of degree {polynomial.degree()}"
        )
    for root, _multiplicity in roots:
        if (
            root.imag.contains(0)
            and not (root.real < -1 or root.real > 1)
        ):
            raise CertificationError(
                f"{label} has an unresolved root in [-1, 1]: {root}"
            )

    return ChebyshevLowerBoundCertificate(
        degree=polynomial.degree(),
        lower_bound=bound,
        value_at_minus_one=value_at_minus_one,
        value_at_one=value_at_one,
        enclosed_roots=enclosed_roots,
    )
