"""Tests for the SP(k) Fortran-port implementation logic.

These tests validate that the equations and calibration tables hard-coded into the
CAMB-side Fortran implementation match the authoritative Python implementation in
`pyspk`.

Notes:
    These tests do not require CAMB to be compiled. They validate the numerical
    equivalence of the suppression formula and the copied calibration coefficients.
"""

from __future__ import annotations

import math

import numpy as np
import pytest

import pyspk
from pyspk.fit_vals import best_fit_vals

# --- Calibration coefficients duplicated from dev/CAMB/fortran/halofit.f90 ---
# Kept explicit (rather than sourced from pyspk) so the test will catch copy/paste
# mistakes when porting values into CAMB.
_SPK_COEFFS: dict[int, dict[str, tuple[float, float, float]]] = {
    200: {
        "alpha": (15.24311120000861, -1.2436699435560352, 0.14837558774401766),
        "beta": (14.969187892657688, -1.0993025612653198, 0.12905587245129102),
        "gamma": (0.8000441576980428, -0.01715621131893159, 0.06131887249968379),
        "lambda_a": (0.02178116280689233, -0.0077564325654746955, 0.0007915576054589781),
        "lambda_b": (3.0878286643613437, 0.4529677646796634, 0.001552571083240605),
        "mu_a": (0.6930259177449359, -0.16913553700233935, 0.04263185199898842),
        "mu_b": (3.161914061444856, 0.8616834297321924, 0.011346427353554053),
        "mu_c": (5.532188503256583, -3.0864672185252537, 0.5083422518560442),
        "nu_a": (413.00988701513904, 311.63957063032285, 37.89105940901369),
        "nu_b": (-11.243859405779181, -0.34421412616421965, 0.3343548325485801),
        "nu_c": (3.476463891168505, -0.018333059687988575, -0.08276237963970698),
    },
    500: {
        "alpha": (14.783423122120318, -0.999062404857228, 0.12062854541689262),
        "beta": (14.620528368613265, -0.9136466201011957, 0.10835389086945699),
        "gamma": (0.9671320682693298, -0.03185388045484575, 0.02650236152450093),
        "lambda_a": (0.019349810078190303, -0.007410668383424459, 0.0008334762393555539),
        "lambda_b": (2.9566773924238143, 0.6205340408676114, -0.001928273640110775),
        "mu_a": (0.715853343781141, -0.19276613600825665, 0.04948240117059147),
        "mu_b": (3.385355123440431, 0.9658906605139421, -0.06825861100375574),
        "mu_c": (4.457257708010122, -2.191853871334233, 0.45457701107254733),
        "nu_a": (478.86477329610375, 429.88795783439946, 249.25655627821902),
        "nu_b": (-11.227459319819815, -0.5581080204509223, 0.4489962047114509),
        "nu_c": (3.499449440557995, -0.08488559389068073, -0.0923847866118189),
    },
}


def _poly2(x: float, c0: float, c1: float, c2: float) -> float:
    return c2 * x * x + c1 * x + c0


def _spk_params_port(SO: int, z: float) -> dict[str, float]:
    """Compute best-fit parameters with the same clamp behavior as the CAMB port."""
    z_eff = float(min(max(z, 0.125), 3.0))
    x = 1.0 + z_eff

    coeffs = _SPK_COEFFS[int(SO)]
    return {name: float(_poly2(x, *coeffs[name])) for name in coeffs}


def _spk_suppression_powerlaw_port(
    *,
    SO: int,
    k: np.ndarray,
    z: float,
    fb_a: float,
    fb_pow: float,
    fb_pivot: float,
) -> np.ndarray:
    """Vectorized numpy port of the Fortran SP(k) suppression (power-law relation only)."""
    # CAMB port clamps z and k.
    z_eff = float(min(max(z, 0.125), 3.0))
    k_eff = np.clip(np.asarray(k, dtype=float), 1e-12, 12.0)

    p = _spk_params_port(SO, z_eff)

    best_mass = p["alpha"] - (p["alpha"] - p["beta"]) * np.power(k_eff, p["gamma"])
    m_opt = np.power(10.0, best_mass)

    if fb_pivot > 0:
        fb = float(fb_a) * np.power(m_opt / float(fb_pivot), float(fb_pow))
    else:
        fb = np.full_like(m_opt, float(fb_a))

    x = np.log10(k_eff)
    x0 = 1.0 + p["lambda_a"] * np.exp(p["lambda_b"] * x)
    x1 = p["mu_a"] + (1.0 - p["mu_a"]) / (1.0 + np.exp(p["mu_b"] * x + p["mu_c"]))
    x2 = p["nu_a"] * np.exp(-0.5 * np.square((x - p["nu_b"]) / p["nu_c"]))

    sup = x0 - (x0 - x1) * np.exp(-x2 * fb)
    return np.maximum(sup, 1e-6)


def test_spk_coefficients_match_pyspk_tables() -> None:
    """Ensure the coefficients copied into the CAMB Fortran port match pyspk."""
    for so, coeffs in _SPK_COEFFS.items():
        so_key = str(int(so))
        for name, triplet in coeffs.items():
            expected = tuple(float(x) for x in best_fit_vals[so_key][name])
            assert triplet == expected


@pytest.mark.parametrize("SO", [200, 500])
def test_spk_suppression_matches_pyspk_sup_model_for_power_law_relation(SO: int) -> None:
    """Suppression computed by the port matches pyspk for finite points."""
    rng = np.random.default_rng(12345 + SO)

    # Keep k within the port clamp range; include a couple of values > 8 (pyspk default k_max)
    k = np.geomspace(1e-3, 12.0, 64)

    # Pick z safely inside calibration range to avoid differences in behavior.
    z = float(rng.uniform(0.125, 3.0))

    # Choose a simple constant baryon fraction to reduce risk of triggering pyspk fitting-limit NaNs.
    fb_a = 0.5
    fb_pow = 0.0
    fb_pivot = 1.0

    _, sup_py = pyspk.sup_model(
        SO=SO,
        z=z,
        fb_a=fb_a,
        fb_pow=fb_pow,
        fb_pivot=fb_pivot,
        k_array=k,
        errors=False,
    )

    sup_port = _spk_suppression_powerlaw_port(
        SO=SO,
        k=k,
        z=z,
        fb_a=fb_a,
        fb_pow=fb_pow,
        fb_pivot=fb_pivot,
    )

    finite = np.isfinite(sup_py)
    assert finite.any(), "pyspk returned no finite suppression values for this test input"

    # Tight tolerances: formulas should be identical for finite points.
    assert np.allclose(sup_port[finite], sup_py[finite], rtol=5e-12, atol=5e-12)


def test_spk_clamping_behavior_matches_spec() -> None:
    """The CAMB port clamps z and k; validate that behavior is consistent."""
    SO = 200
    fb_a = 0.5
    fb_pow = 0.0
    fb_pivot = 1.0

    k_hi = np.array([20.0, 50.0], dtype=float)
    s_hi = _spk_suppression_powerlaw_port(SO=SO, k=k_hi, z=0.5, fb_a=fb_a, fb_pow=fb_pow, fb_pivot=fb_pivot)
    s_12 = _spk_suppression_powerlaw_port(
        SO=SO, k=np.array([12.0, 12.0]), z=0.5, fb_a=fb_a, fb_pow=fb_pow, fb_pivot=fb_pivot
    )
    assert np.allclose(s_hi, s_12)

    k = np.array([0.1, 1.0, 10.0], dtype=float)
    s_z0 = _spk_suppression_powerlaw_port(SO=SO, k=k, z=0.0, fb_a=fb_a, fb_pow=fb_pow, fb_pivot=fb_pivot)
    s_zmin = _spk_suppression_powerlaw_port(SO=SO, k=k, z=0.125, fb_a=fb_a, fb_pow=fb_pow, fb_pivot=fb_pivot)
    assert np.allclose(s_z0, s_zmin)

    s_z10 = _spk_suppression_powerlaw_port(SO=SO, k=k, z=10.0, fb_a=fb_a, fb_pow=fb_pow, fb_pivot=fb_pivot)
    s_zmax = _spk_suppression_powerlaw_port(SO=SO, k=k, z=3.0, fb_a=fb_a, fb_pow=fb_pow, fb_pivot=fb_pivot)
    assert np.allclose(s_z10, s_zmax)


def test_spk_suppression_floor_is_positive() -> None:
    """Ensure suppression is floored to a small positive value (Fortran safety)."""
    # Use an extreme (negative) fb_a to try to create numerical issues; the port should still floor.
    k = np.geomspace(1e-3, 12.0, 16)
    sup = _spk_suppression_powerlaw_port(SO=200, k=k, z=0.5, fb_a=-100.0, fb_pow=0.0, fb_pivot=1.0)
    assert np.all(np.isfinite(sup))
    assert np.all(sup >= 1e-6)


def test_spk_port_is_monotonic_wrt_floor_only() -> None:
    """Sanity check: the port returns a vector with same shape and finite values."""
    k = np.array([1e-3, 0.1, 1.0, 12.0])
    sup = _spk_suppression_powerlaw_port(SO=500, k=k, z=1.0, fb_a=0.5, fb_pow=0.0, fb_pivot=1.0)
    assert sup.shape == k.shape
    assert np.isfinite(sup).all()
    assert math.isfinite(float(sup.mean()))
