"""
TNFR Riemann Zeta Function Implementation
=========================================

This module provides the canonical implementation of the Riemann Zeta function
and related spectral functions for the TNFR engine.

It wraps high-precision arithmetic (mpmath) to ensure structural fidelity
during critical line exploration.

Functions:
- zeta(s): The Riemann Zeta function.
- chi_factor(s): The symmetry factor χ(s).
- structural_potential(s): Φ_s = log|ζ(s)|.
- structural_pressure(s): ΔNFR = |log|χ(s)||.
"""

from typing import Any

# Import TNFR Cache Infrastructure
from .unified_cache import CacheLevel, cache_tnfr_computation
from .unified_numerical import np

_CACHE_AVAILABLE = True

# Import Unified Backend
try:
    HAS_BACKEND = True
except ImportError:
    HAS_BACKEND = False

# We use mpmath for high precision (FFT-based arithmetic for large numbers)
try:
    from mpmath import fabs as mp_fabs
    from mpmath import gamma as mp_gamma
    from mpmath import log as mp_log
    from mpmath import mp
    from mpmath import pi as mp_pi
    from mpmath import power as mp_power
    from mpmath import sin as mp_sin
    from mpmath import zeta as mp_zeta
    from mpmath import zetazero as mp_zetazero

    HAS_MPMATH = True
except ImportError:
    HAS_MPMATH = False

# Default precision
if HAS_MPMATH:
    mp.mp.dps = 25


def _ensure_complex(s: Any) -> Any:
    if HAS_MPMATH:
        return mp.mpc(s)
    return complex(s)


@cache_tnfr_computation(level=CacheLevel.DERIVED_METRICS, dependencies=set())
def zeta_zero(n: int) -> complex:
    """
    Return the n-th non-trivial zero of the Riemann Zeta function.
    """
    if HAS_MPMATH:
        return mp_zetazero(n)
    raise ImportError("mpmath required for zeta zeros.")


@cache_tnfr_computation(level=CacheLevel.DERIVED_METRICS, dependencies=set())
def zeta_function(s: complex | float | Any) -> Any:
    """
    Compute the Riemann Zeta function ζ(s).

    Integration:
    - Uses mpmath for high-precision FFT-based arithmetic.
    - Uses TNFR Cache for structural field memoization.
    - Compatible with Unified Backend for future vectorization.
    """
    if HAS_MPMATH:
        return mp_zeta(s)

    # Fallback to scipy if available, else raise
    try:
        from scipy.special import zeta as scipy_zeta

        return scipy_zeta(s)
    except ImportError:
        raise ImportError("mpmath or scipy required for zeta function.")


@cache_tnfr_computation(level=CacheLevel.DERIVED_METRICS, dependencies=set())
def chi_factor(s: complex | float | Any) -> Any:
    """
    Compute the Riemann xi factor χ(s).
    χ(s) = 2^s * π^(s-1) * sin(πs/2) * Γ(1-s)
    """
    if HAS_MPMATH:
        s_mp = mp.mpc(s)
        return (
            mp_power(2, s_mp)
            * mp_power(mp_pi, s_mp - 1)
            * mp_sin(mp_pi * s_mp / 2)
            * mp_gamma(1 - s_mp)
        )

    # Standard float implementation
    s_c = complex(s)
    return (2**s_c) * (np.pi ** (s_c - 1)) * np.sin(np.pi * s_c / 2) * np.math.gamma(1 - s_c)  # type: ignore


@cache_tnfr_computation(level=CacheLevel.DERIVED_METRICS, dependencies=set())
def structural_potential(s: complex | float | Any) -> float:
    """
    Compute the Structural Potential Φ_s = log|ζ(s)|.
    """
    z = zeta_function(s)
    mag = abs(z)
    if HAS_MPMATH:
        return float(mp_log(mag + 1e-20))
    return float(np.log(mag + 1e-20))


@cache_tnfr_computation(level=CacheLevel.DERIVED_METRICS, dependencies=set())
def structural_pressure(s: complex | float | Any) -> float:
    """
    Compute the Structural Pressure ΔNFR = |log|χ(s)||.
    This is the derived form from symmetry breaking.
    """
    chi = chi_factor(s)
    mag = abs(chi)
    if HAS_MPMATH:
        return float(mp_fabs(mp_log(mag)))
    return float(abs(np.log(mag)))
    return float(abs(np.log(mag)))