"""
Tests for mlx-fisher.

Author: Sheng-Kai Huang <akai@fawstudio.com>
"""

import mlx.core as mx
import numpy as np
import pytest

from mlx_fisher import (
    FisherMatrix,
    NaturalGradientOptimizer,
    fisher_matrix_cl,
    kl_divergence,
    quantum_relative_entropy,
)
from mlx_fisher.natural_gradient import natural_gradient_descent


# ======================================================================
# FisherMatrix basics
# ======================================================================


class TestFisherMatrixBasics:
    def test_creation(self):
        F = FisherMatrix(mx.eye(3))
        assert F.n_params == 3

    def test_non_square_raises(self):
        with pytest.raises(ValueError):
            FisherMatrix(mx.zeros((2, 3)))

    def test_inverse_identity(self):
        F = FisherMatrix(mx.eye(4))
        F_inv = F.inverse()
        np.testing.assert_allclose(np.array(F_inv), np.eye(4), atol=1e-6)

    def test_inverse_with_reg(self):
        F = FisherMatrix(mx.eye(3) * 2.0)
        F_inv = F.inverse(reg=1.0)
        # (2I + I)^{-1} = (1/3) I
        np.testing.assert_allclose(np.array(F_inv), np.eye(3) / 3.0, atol=1e-6)

    def test_marginal_errors(self):
        # For F = diag(4, 9, 16), marginal errors = 1/sqrt(diag) = [0.5, 1/3, 0.25]
        F = FisherMatrix(mx.diag(mx.array([4.0, 9.0, 16.0])))
        errs = F.marginal_errors()
        np.testing.assert_allclose(np.array(errs), [0.5, 1.0 / 3.0, 0.25], atol=1e-5)

    def test_eigenvalues(self):
        F = FisherMatrix(mx.diag(mx.array([1.0, 3.0, 5.0])))
        ev = F.eigenvalues()
        np.testing.assert_allclose(sorted(np.array(ev)), [1.0, 3.0, 5.0], atol=1e-5)

    def test_condition_number(self):
        F = FisherMatrix(mx.diag(mx.array([1.0, 10.0])))
        assert abs(F.condition_number() - 10.0) < 0.1

    def test_add(self):
        F1 = FisherMatrix(mx.eye(2) * 2.0)
        F2 = FisherMatrix(mx.eye(2) * 3.0)
        F3 = F1 + F2
        np.testing.assert_allclose(np.array(F3.matrix), np.eye(2) * 5.0, atol=1e-6)

    def test_to_numpy(self):
        F = FisherMatrix(mx.eye(2))
        arr = F.to_numpy()
        assert isinstance(arr, np.ndarray)
        assert arr.shape == (2, 2)

    def test_repr(self):
        F = FisherMatrix(mx.eye(3))
        r = repr(F)
        assert "n_params=3" in r


# ======================================================================
# FisherMatrix.from_model (Gaussian Fisher)
# ======================================================================


class TestGaussianFisher:
    def test_linear_model(self):
        """
        For model mu(theta) = A @ theta with known A and uniform sigma,
        the Fisher matrix should be (1/sigma^2) A^T A.
        """
        n_data, n_params = 100, 3
        key = mx.random.key(42)
        A = mx.random.normal((n_data, n_params), key=key)
        sigma_val = 2.0
        sigma = mx.ones((n_data,)) * sigma_val

        def model_fn(theta):
            return A @ theta

        theta0 = mx.zeros((n_params,))
        F = FisherMatrix.from_model(model_fn, theta0, sigma)

        expected = np.array(A.T @ A) / (sigma_val ** 2)
        np.testing.assert_allclose(F.to_numpy(), expected, rtol=1e-3)

    def test_symmetry(self):
        """Fisher matrix must be symmetric."""
        def model_fn(theta):
            return mx.array([theta[0] ** 2, mx.sin(theta[1]), theta[0] * theta[1]])

        theta0 = mx.array([1.0, 2.0])
        sigma = mx.ones((3,))
        F = FisherMatrix.from_model(model_fn, theta0, sigma)
        F_np = F.to_numpy()
        np.testing.assert_allclose(F_np, F_np.T, atol=1e-5)

    def test_positive_semidefinite(self):
        """Fisher matrix must be positive semi-definite."""
        def model_fn(theta):
            return mx.array([theta[0] + theta[1], theta[0] - theta[1], theta[0] * 2])

        theta0 = mx.array([1.0, 1.0])
        sigma = mx.ones((3,)) * 0.5
        F = FisherMatrix.from_model(model_fn, theta0, sigma)
        eigvals = np.linalg.eigvalsh(F.to_numpy())
        assert np.all(eigvals >= -1e-6)


# ======================================================================
# FisherMatrix.from_loglikelihood
# ======================================================================


class TestFromLoglikelihood:
    def test_gaussian_loglikelihood(self):
        """
        For data ~ N(mu, sigma^2) with known sigma, the Fisher information
        for mu is 1/sigma^2 (per-sample expectation).
        from_loglikelihood returns the empirical average of the outer product
        of score vectors, so F ~ 1/sigma^2.
        """
        sigma_true = 1.5
        np.random.seed(123)
        data = mx.array(np.random.normal(0, sigma_true, 500).astype(np.float32))

        def log_likelihood(theta):
            # theta[0] = mu
            mu = theta[0]
            return -0.5 * ((data - mu) / sigma_true) ** 2

        theta0 = mx.array([0.0])
        F = FisherMatrix.from_loglikelihood(log_likelihood, theta0, n_samples=500)
        # Expected: F ~ 1 / sigma^2 = 1 / 2.25 ~ 0.444
        expected_F = 1.0 / (sigma_true ** 2)
        assert abs(F.to_numpy()[0, 0] - expected_F) / expected_F < 0.15


# ======================================================================
# FisherMatrix.from_samples
# ======================================================================


class TestFromSamples:
    def test_normal_samples(self):
        """Empirical Fisher from normal samples should approximate n/sigma^2."""
        sigma_true = 2.0
        np.random.seed(456)
        samples = mx.array(np.random.normal(0, sigma_true, 200).astype(np.float32))

        def log_prob(theta, x):
            mu = theta[0]
            return -0.5 * ((x - mu) / sigma_true) ** 2

        theta0 = mx.array([0.0])
        F = FisherMatrix.from_samples(log_prob, theta0, samples)
        expected_F = 1.0 / (sigma_true ** 2)
        # Empirical Fisher should be close to theoretical (within 30%)
        assert abs(F.to_numpy()[0, 0] - expected_F) / expected_F < 0.30


# ======================================================================
# fisher_matrix_cl
# ======================================================================


class TestFisherMatrixCl:
    def test_simple_power_law(self):
        """
        Toy C_l = A * l^(-n): Fisher matrix for (A, n) should be
        well-conditioned and positive definite.
        """
        l_max = 100

        def cl_fn(theta):
            A = theta[0]
            n = theta[1]
            ell = mx.arange(l_max + 1, dtype=mx.float32) + 1.0  # avoid l=0
            cl = A * mx.power(ell, -n)
            return cl

        theta0 = mx.array([1.0, 2.0])
        F = fisher_matrix_cl(cl_fn, theta0, f_sky=0.7, l_min=2, l_max=l_max)

        assert F.n_params == 2
        F_np = F.to_numpy()
        # Symmetric
        np.testing.assert_allclose(F_np, F_np.T, atol=1e-4)
        # Positive definite
        eigvals = np.linalg.eigvalsh(F_np)
        assert np.all(eigvals > 0)

    def test_with_noise(self):
        """Adding noise should reduce Fisher information."""
        l_max = 50

        def cl_fn(theta):
            A = theta[0]
            ell = mx.arange(l_max + 1, dtype=mx.float32) + 1.0
            return A / ell

        theta0 = mx.array([1.0])

        F_no_noise = fisher_matrix_cl(cl_fn, theta0, l_min=2, l_max=l_max)
        noise = mx.ones((l_max + 1,)) * 0.1
        F_with_noise = fisher_matrix_cl(cl_fn, theta0, noise_cl=noise, l_min=2, l_max=l_max)

        # More noise => less information
        assert F_with_noise.to_numpy()[0, 0] < F_no_noise.to_numpy()[0, 0]


# ======================================================================
# KL divergence
# ======================================================================


class TestKLDivergence:
    def test_identical_distributions(self):
        """D_KL(p || p) = 0."""
        p = mx.array([0.25, 0.25, 0.25, 0.25])
        d = kl_divergence(p, p)
        assert abs(d.item()) < 1e-6

    def test_known_value(self):
        """D_KL(p || q) for known p, q."""
        p = mx.array([0.5, 0.5])
        q = mx.array([0.25, 0.75])
        # D_KL = 0.5*ln(0.5/0.25) + 0.5*ln(0.5/0.75)
        #       = 0.5*ln(2) + 0.5*ln(2/3)
        expected = 0.5 * np.log(2) + 0.5 * np.log(2.0 / 3.0)
        assert abs(kl_divergence(p, q).item() - expected) < 1e-5

    def test_non_negative(self):
        """KL divergence is always non-negative (Gibbs' inequality)."""
        np.random.seed(789)
        for _ in range(10):
            p_raw = np.abs(np.random.randn(20)) + 1e-6
            q_raw = np.abs(np.random.randn(20)) + 1e-6
            p = mx.array((p_raw / p_raw.sum()).astype(np.float32))
            q = mx.array((q_raw / q_raw.sum()).astype(np.float32))
            assert kl_divergence(p, q).item() >= -1e-6

    def test_asymmetry(self):
        """D_KL(p||q) != D_KL(q||p) in general."""
        p = mx.array([0.9, 0.1])
        q = mx.array([0.1, 0.9])
        d_pq = kl_divergence(p, q).item()
        d_qp = kl_divergence(q, p).item()
        # They should be equal by symmetry of the specific distributions
        # Actually for this case they ARE equal, so pick asymmetric ones
        p2 = mx.array([0.7, 0.3])
        q2 = mx.array([0.2, 0.8])
        d_pq2 = kl_divergence(p2, q2).item()
        d_qp2 = kl_divergence(q2, p2).item()
        # In general not equal (unless special symmetry)
        # Just check they're both non-negative
        assert d_pq2 >= -1e-6
        assert d_qp2 >= -1e-6


# ======================================================================
# Quantum relative entropy
# ======================================================================


class TestQuantumRelativeEntropy:
    def _make_density_matrix(self, dim, seed=0):
        np.random.seed(seed)
        A = np.random.randn(dim, dim).astype(np.float32)
        rho = A @ A.T
        rho = rho / np.trace(rho)
        return mx.array(rho)

    def test_identical_states(self):
        """D(rho || rho) = 0."""
        rho = self._make_density_matrix(4, seed=42)
        D = quantum_relative_entropy(rho, rho)
        assert abs(D.item()) < 1e-4

    def test_non_negative(self):
        """Klein's inequality: D(rho || sigma) >= 0."""
        for seed in range(5):
            rho = self._make_density_matrix(8, seed=seed)
            sigma = self._make_density_matrix(8, seed=seed + 100)
            D = quantum_relative_entropy(rho, sigma)
            assert D.item() >= -1e-4, f"D = {D.item()} for seed {seed}"

    def test_classical_limit(self):
        """
        For diagonal density matrices, quantum relative entropy
        reduces to classical KL divergence.
        """
        p = np.array([0.3, 0.5, 0.2], dtype=np.float32)
        q = np.array([0.1, 0.6, 0.3], dtype=np.float32)

        rho = mx.diag(mx.array(p))
        sigma = mx.diag(mx.array(q))

        D_quantum = quantum_relative_entropy(rho, sigma).item()
        D_classical = kl_divergence(mx.array(p), mx.array(q)).item()

        np.testing.assert_allclose(D_quantum, D_classical, rtol=1e-3)

    def test_pure_vs_mixed(self):
        """D(pure || maximally mixed) = ln(d) for a d-dimensional system."""
        d = 4
        # Pure state |0><0|
        rho = mx.zeros((d, d), dtype=mx.float32)
        rho[0, 0] = 1.0
        # Maximally mixed
        sigma = mx.eye(d, dtype=mx.float32) / d

        D = quantum_relative_entropy(rho, sigma).item()
        expected = np.log(d)
        np.testing.assert_allclose(D, expected, rtol=1e-2)


# ======================================================================
# Natural gradient optimizer
# ======================================================================


class TestNaturalGradient:
    def test_quadratic_convergence(self):
        """
        For a quadratic loss L = 0.5 * theta^T A theta with Fisher = A,
        natural gradient should converge in one step (up to lr).
        """
        A = mx.array([[4.0, 1.0], [1.0, 3.0]])

        def loss_fn(theta):
            return 0.5 * (theta @ A @ theta)

        def fisher_fn(theta):
            return A

        theta0 = mx.array([2.0, -1.0])

        # One step with lr=1 and zero damping should land at origin
        # nat_grad = A^{-1} A theta = theta
        # theta_new = theta - 1.0 * theta = 0
        theta_final, losses = natural_gradient_descent(
            loss_fn, fisher_fn, theta0, lr=1.0, damping=0.0, n_steps=1
        )
        np.testing.assert_allclose(np.array(theta_final), [0.0, 0.0], atol=1e-2)

    def test_optimizer_class(self):
        """Test the NaturalGradientOptimizer class interface."""
        A = mx.array([[2.0, 0.0], [0.0, 2.0]])
        opt = NaturalGradientOptimizer(lr=0.5, damping=0.0)

        theta = mx.array([4.0, -2.0])
        grad = A @ theta  # gradient of 0.5 * theta^T A theta

        def fisher_est(t):
            return A

        theta_new = opt.step(theta, grad, fisher_estimator=fisher_est)
        # nat_grad = A^{-1} A theta = theta
        # theta_new = theta - 0.5 * theta = 0.5 * theta
        np.testing.assert_allclose(np.array(theta_new), [2.0, -1.0], atol=1e-4)

    def test_reset(self):
        opt = NaturalGradientOptimizer()
        A = mx.eye(2)
        opt.step(mx.array([1.0, 1.0]), mx.array([0.1, 0.1]), fisher_estimator=lambda t: A)
        assert opt.fisher_matrix is not None
        opt.reset()
        assert opt.fisher_matrix is None

    def test_missing_fisher_raises(self):
        opt = NaturalGradientOptimizer()
        with pytest.raises(ValueError):
            opt.step(mx.array([1.0]), mx.array([0.1]), fisher_estimator=None)


# ======================================================================
# Integration test
# ======================================================================


class TestIntegration:
    def test_fisher_to_natural_gradient(self):
        """
        End-to-end: compute Fisher from model, then use it
        for natural gradient on a simple regression.
        """
        # True parameters
        theta_true = mx.array([2.0, -1.0, 0.5])
        n_data = 50
        key = mx.random.key(99)
        X = mx.random.normal((n_data,), key=key)
        sigma_noise = 0.3

        # Model: y = a*x^2 + b*x + c
        def model_fn(theta):
            return theta[0] * X * X + theta[1] * X + theta[2]

        y_true = model_fn(theta_true)

        # Compute Fisher at a nearby point
        theta_start = mx.array([1.5, -0.5, 0.0])
        sigma = mx.ones((n_data,)) * sigma_noise
        F = FisherMatrix.from_model(model_fn, theta_start, sigma)

        # Fisher should be positive definite
        eigvals = np.linalg.eigvalsh(F.to_numpy())
        assert np.all(eigvals > 0)

        # Marginal errors should be finite and positive
        errs = F.marginal_errors()
        assert np.all(np.array(errs) > 0)
        assert np.all(np.isfinite(np.array(errs)))
