"""Test Fisher-Rao geometry on Δ⁶³."""

import numpy as np
import pytest

from qig_core.geometry.fisher_rao import (
    bhattacharyya_coefficient,
    exp_map,
    fisher_rao_distance,
    frechet_mean,
    log_map,
    random_basin,
    slerp_sqrt,
    to_simplex,
)


class TestToSimplex:
    def test_output_sums_to_one(self):
        raw = np.random.randn(64)
        p = to_simplex(raw)
        assert abs(p.sum() - 1.0) < 1e-10

    def test_all_positive(self):
        raw = np.array([-5.0, 0.0, 3.0, 1.0])
        p = to_simplex(raw)
        assert (p > 0).all()

    def test_idempotent(self):
        p = to_simplex(np.random.rand(64))
        p2 = to_simplex(p)
        np.testing.assert_allclose(p, p2, atol=1e-12)


class TestFisherRaoDistance:
    def test_self_distance_is_zero(self, uniform_basin):
        d = fisher_rao_distance(uniform_basin, uniform_basin)
        assert abs(d) < 1e-10

    def test_symmetric(self, random_basins):
        p, q = random_basins[0], random_basins[1]
        d1 = fisher_rao_distance(p, q)
        d2 = fisher_rao_distance(q, p)
        assert abs(d1 - d2) < 1e-10

    def test_triangle_inequality(self, random_basins):
        p, q, r = random_basins[0], random_basins[1], random_basins[2]
        d_pq = fisher_rao_distance(p, q)
        d_qr = fisher_rao_distance(q, r)
        d_pr = fisher_rao_distance(p, r)
        assert d_pr <= d_pq + d_qr + 1e-10

    def test_range(self, random_basins):
        for p in random_basins:
            for q in random_basins:
                d = fisher_rao_distance(p, q)
                assert 0.0 <= d <= np.pi / 2 + 1e-10


class TestBhattacharyya:
    def test_self_is_one(self, uniform_basin):
        bc = bhattacharyya_coefficient(uniform_basin, uniform_basin)
        assert abs(bc - 1.0) < 1e-10

    def test_range(self, random_basins):
        for i in range(len(random_basins)):
            for j in range(i + 1, len(random_basins)):
                bc = bhattacharyya_coefficient(random_basins[i], random_basins[j])
                assert 0.0 <= bc <= 1.0 + 1e-10


class TestSlerp:
    def test_endpoints(self, random_basins):
        p, q = random_basins[0], random_basins[1]
        s0 = slerp_sqrt(p, q, 0.0)
        s1 = slerp_sqrt(p, q, 1.0)
        np.testing.assert_allclose(s0, p, atol=1e-8)
        np.testing.assert_allclose(s1, q, atol=1e-8)

    def test_midpoint_on_simplex(self, random_basins):
        p, q = random_basins[0], random_basins[1]
        mid = slerp_sqrt(p, q, 0.5)
        assert abs(mid.sum() - 1.0) < 1e-10
        assert (mid > 0).all()


class TestFrechetMean:
    def test_single_point(self, random_basins):
        p = random_basins[0]
        mean = frechet_mean([p])
        np.testing.assert_allclose(mean, p, atol=1e-10)

    def test_on_simplex(self, random_basins):
        mean = frechet_mean(random_basins)
        assert abs(mean.sum() - 1.0) < 1e-10
        assert (mean > 0).all()

    def test_closer_to_inputs_than_random(self, random_basins, rng):
        mean = frechet_mean(random_basins)
        rand_point = rng.dirichlet(np.ones(64))
        avg_dist_mean = np.mean([fisher_rao_distance(mean, p) for p in random_basins])
        avg_dist_rand = np.mean([fisher_rao_distance(rand_point, p) for p in random_basins])
        # Fréchet mean should generally be closer to the set than a random point
        # (not a strict guarantee for every random point, but statistically reliable)
        assert avg_dist_mean <= avg_dist_rand + 0.3


class TestLogExpMap:
    def test_roundtrip(self, random_basins):
        p, q = random_basins[0], random_basins[1]
        v = log_map(p, q)
        q_recovered = exp_map(p, v)
        np.testing.assert_allclose(q_recovered, q, atol=1e-6)


class TestRandomBasin:
    def test_on_simplex(self):
        b = random_basin(64)
        assert abs(b.sum() - 1.0) < 1e-10
        assert (b > 0).all()
        assert len(b) == 64
