"""Fast finite-zero explicit-formula fields.

The NUFFT methods are exploration-tier floating point.  They are intended for
large candidate searches followed by direct or interval verification at the
small set of retained points.
"""

from __future__ import annotations

import math

import finufft
import numpy as np

__all__ = ["SkewesZeroField"]


class SkewesZeroField:
    """The finite critical-line field used to locate ``pi-li`` crossovers.

    For positive zero ordinates ``gamma`` and decimal exponent ``w``, this
    evaluates

    ``baseline - 2 Re sum exp(i gamma w log(10)) / (1/2 + i gamma)``.
    """

    def __init__(self, ordinates: np.ndarray) -> None:
        values = np.asarray(ordinates, dtype=np.float64)
        if values.ndim != 1 or values.size == 0:
            raise ValueError("ordinates must be a nonempty one-dimensional array")
        if not np.all(np.isfinite(values)) or not np.all(values > 0):
            raise ValueError("ordinates must be finite and positive")
        if not np.all(np.diff(values) > 0):
            raise ValueError("ordinates must be strictly increasing")
        self.ordinates = values
        self.coefficients = -2.0 / (0.5 + 1j * values)
        self.log10 = math.log(10.0)
        self.scale = values[-1] / (0.95 * math.pi)
        self.sources = values / self.scale

    def values(
        self,
        decimal_exponents: np.ndarray,
        *,
        baseline: float = -1.0,
        eps: float = 1e-12,
        nthreads: int = 1,
    ) -> np.ndarray:
        """Evaluate a target array with a type-3 NUFFT."""
        targets = np.asarray(decimal_exponents, dtype=np.float64)
        if targets.ndim != 1:
            raise ValueError("decimal_exponents must be one-dimensional")
        transformed = finufft.nufft1d3(
            self.sources,
            self.coefficients,
            targets * (self.log10 * self.scale),
            eps=eps,
            isign=1,
            nthreads=nthreads,
        )
        return baseline + transformed.real

    def direct_value(
        self,
        decimal_exponent: float,
        *,
        baseline: float = -1.0,
    ) -> float:
        """Evaluate one point by a direct complex sum."""
        phase = self.ordinates * (decimal_exponent * self.log10)
        return baseline + float(
            np.real(
                np.sum(
                    self.coefficients * np.exp(1j * phase),
                    dtype=np.complex128,
                )
            )
        )

    def direct_derivative(self, decimal_exponent: float) -> float:
        """Evaluate the derivative with respect to decimal exponent."""
        phase = self.ordinates * (decimal_exponent * self.log10)
        return float(
            np.real(
                np.sum(
                    self.coefficients
                    * (1j * self.ordinates * self.log10)
                    * np.exp(1j * phase),
                    dtype=np.complex128,
                )
            )
        )

    def second_derivative_absolute_bound(self) -> float:
        """Return the triangle-inequality bound for ``sup |F''|``."""
        ordinates = self.ordinates
        return 2.0 * self.log10 * self.log10 * float(
            np.sum(
                ordinates
                * ordinates
                / np.sqrt(ordinates * ordinates + 0.25),
                dtype=np.float64,
            )
        )

    def grid_interpolation_reserve(self, step: float) -> float:
        """Bound a local maximum above its nearest uniform-grid sample.

        At an interior local maximum the first derivative vanishes.  Taylor's
        theorem and a nearest grid point at distance at most ``step/2`` give
        this reserve.  Endpoints must themselves be included in the grid.
        """
        if step <= 0:
            raise ValueError("step must be positive")
        return (
            0.5
            * (step / 2.0) ** 2
            * self.second_derivative_absolute_bound()
        )
