#!/usr/bin/env python3
"""Independent NumPy checks for the spectral-space AI-REML equations.

This file deliberately does not import the production C++ implementation.  It
serves as a small, readable oracle for gradient and transformation tests.
"""

from __future__ import annotations

import math

import numpy as np


def restricted_loglik_and_derivatives(
    theta: np.ndarray,
    eigenvalues: np.ndarray,
    x_star: np.ndarray,
    y_star: np.ndarray,
) -> tuple[float, np.ndarray, np.ndarray]:
    """Return paper-scale REML log likelihood, score, and AI matrix.

    Parameters are the standard-deviation coordinates ``(sigma_e, sigma_g)``.
    The function forms P explicitly because it is only used on tiny test data.
    """

    sigma_e, sigma_g = (float(theta[0]), float(theta[1]))
    diagonal = sigma_e**2 + sigma_g**2 * eigenvalues
    if np.any(diagonal <= 0.0):
        raise ValueError("H is not positive definite")
    w = 1.0 / diagonal
    h_inv = np.diag(w)
    c = x_star.T @ h_inv @ x_star
    c_inv = np.linalg.inv(c)
    p = h_inv - h_inv @ x_star @ c_inv @ x_star.T @ h_inv
    py = p @ y_star

    loglik = -0.5 * (
        np.linalg.slogdet(c)[1]
        + np.log(diagonal).sum()
        + float(y_star @ py)
    )

    score_e = -sigma_e * (np.trace(p) - float(py @ py))
    score_g = -sigma_g * (
        np.trace(p @ np.diag(eigenvalues))
        - float(py @ (eigenvalues * py))
    )

    p_py = p @ py
    lambda_py = eigenvalues * py
    p_lambda_py = p @ lambda_py
    ai_ee = 2.0 * sigma_e**2 * float(py @ p_py)
    ai_eg = 2.0 * sigma_e * sigma_g * float(py @ p_lambda_py)
    ai_gg = 2.0 * sigma_g**2 * float(lambda_py @ p_lambda_py)
    ai = np.asarray([[ai_ee, ai_eg], [ai_eg, ai_gg]], dtype=np.float64)
    return loglik, np.asarray([score_e, score_g]), ai


def finite_difference_gradient(
    theta: np.ndarray,
    eigenvalues: np.ndarray,
    x_star: np.ndarray,
    y_star: np.ndarray,
    step: float = 1e-6,
) -> np.ndarray:
    result = np.empty(2, dtype=np.float64)
    for index in range(2):
        delta = np.zeros(2, dtype=np.float64)
        delta[index] = step
        high = restricted_loglik_and_derivatives(
            theta + delta, eigenvalues, x_star, y_star
        )[0]
        low = restricted_loglik_and_derivatives(
            theta - delta, eigenvalues, x_star, y_star
        )[0]
        result[index] = (high - low) / (2.0 * step)
    return result


def test_score_matches_finite_difference() -> None:
    rng = np.random.default_rng(70123)
    n = 17
    p = 4
    eigenvalues = np.linspace(0.05, 2.1, n)
    x_star = np.column_stack((np.ones(n), rng.normal(size=(n, p - 1))))
    y_star = rng.normal(size=n)
    theta = np.asarray([0.83, 0.57])
    _, score, ai = restricted_loglik_and_derivatives(
        theta, eigenvalues, x_star, y_star
    )
    numerical = finite_difference_gradient(theta, eigenvalues, x_star, y_star)
    np.testing.assert_allclose(score, numerical, rtol=2e-6, atol=2e-6)
    np.testing.assert_allclose(ai, ai.T, rtol=0.0, atol=1e-12)
    assert np.linalg.eigvalsh(ai).min() > 0.0


def test_orthogonal_transformation_preserves_reml() -> None:
    rng = np.random.default_rng(70124)
    n = 15
    q, _ = np.linalg.qr(rng.normal(size=(n, n)))
    eigenvalues = np.linspace(0.1, 1.8, n)
    grm = q @ np.diag(eigenvalues) @ q.T
    x = np.column_stack((np.ones(n), rng.normal(size=(n, 3))))
    y = rng.normal(size=n)
    theta = np.asarray([0.76, 0.62])
    sigma_e, sigma_g = theta

    h = sigma_g**2 * grm + sigma_e**2 * np.eye(n)
    h_inv = np.linalg.inv(h)
    c = x.T @ h_inv @ x
    p = h_inv - h_inv @ x @ np.linalg.inv(c) @ x.T @ h_inv
    original = -0.5 * (
        np.linalg.slogdet(c)[1]
        + np.linalg.slogdet(h)[1]
        + float(y @ p @ y)
    )

    transformed = restricted_loglik_and_derivatives(
        theta, eigenvalues, q.T @ x, q.T @ y
    )[0]
    assert math.isclose(original, transformed, rel_tol=2e-12, abs_tol=2e-12)


if __name__ == "__main__":
    test_score_matches_finite_difference()
    test_orthogonal_transformation_preserves_reml()
    print("PASS: spectral AI-REML reference checks")