"""Tests for fce.fidelity with real assertions."""

import numpy as np
import pytest
from scipy import linalg

import sys, os
sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..'))

from fce.fidelity import (
    uhlmann_fidelity, infidelity, bures_distance,
    pure_state_fidelity, von_neumann_entropy, purity,
    linear_entropy, validate_density_matrix,
)
from tests.conftest import random_density_matrix, random_pure_state


class TestUhlmannFidelity:
    """Tests for the Uhlmann-Jozsa fidelity."""

    def test_identical_states_fidelity_one(self, pure_zero):
        """F(rho, rho) = 1 for any state."""
        assert abs(uhlmann_fidelity(pure_zero, pure_zero) - 1.0) < 1e-10

    def test_identical_mixed_state(self, maximally_mixed_2):
        """F(I/2, I/2) = 1."""
        f = uhlmann_fidelity(maximally_mixed_2, maximally_mixed_2)
        assert abs(f - 1.0) < 1e-10

    def test_orthogonal_pure_states_fidelity_zero(self, pure_zero, pure_one):
        """F(|0><0|, |1><1|) = 0."""
        f = uhlmann_fidelity(pure_zero, pure_one)
        assert f < 1e-10

    def test_pure_state_against_mixed(self, pure_zero, maximally_mixed_2):
        """F(|0><0|, I/2) = 0.5 -- THIS GUARDS AGAINST THE OLD BUG."""
        f = uhlmann_fidelity(pure_zero, maximally_mixed_2)
        assert abs(f - 0.5) < 0.01, \
            f"F(|0><0|, I/2) should be 0.5, got {f}. Old code returned 1.0."

    def test_symmetry(self, rng):
        """F(rho, sigma) = F(sigma, rho)."""
        rho = random_density_matrix(4, rng)
        sigma = random_density_matrix(4, rng)
        f1 = uhlmann_fidelity(rho, sigma)
        f2 = uhlmann_fidelity(sigma, rho)
        assert abs(f1 - f2) < 1e-10

    def test_bounded_zero_one(self, rng):
        """Fidelity is always in [0, 1]."""
        for _ in range(10):
            rho = random_density_matrix(4, rng)
            sigma = random_density_matrix(4, rng)
            f = uhlmann_fidelity(rho, sigma)
            assert 0.0 - 1e-10 <= f <= 1.0 + 1e-10

    def test_pure_state_overlap(self):
        """For pure states, F = |<psi|phi>|^2."""
        psi = np.array([1.0, 0.0], dtype=complex)
        phi = np.array([np.cos(0.3), np.sin(0.3)], dtype=complex)
        rho = np.outer(psi, psi.conj())
        sigma = np.outer(phi, phi.conj())

        f = uhlmann_fidelity(rho, sigma)
        expected = abs(psi.conj() @ phi) ** 2
        assert abs(f - expected) < 1e-10

    def test_random_density_matrices(self, rng):
        """Self-fidelity of random states is 1."""
        for _ in range(5):
            rho = random_density_matrix(8, rng)
            assert abs(uhlmann_fidelity(rho, rho) - 1.0) < 1e-8

    def test_higher_dimension(self, rng):
        """Works for 16-dimensional Hilbert space (4 qubits)."""
        rho = random_density_matrix(16, rng)
        sigma = random_density_matrix(16, rng)
        f = uhlmann_fidelity(rho, sigma)
        assert 0.0 <= f <= 1.0 + 1e-10


class TestInfidelity:
    def test_perfect_match(self, pure_zero):
        assert infidelity(pure_zero, pure_zero) < 1e-10

    def test_orthogonal(self, pure_zero, pure_one):
        assert abs(infidelity(pure_zero, pure_one) - 1.0) < 1e-10


class TestBuresDistance:
    def test_zero_for_identical(self, pure_zero):
        assert bures_distance(pure_zero, pure_zero) < 1e-10

    def test_positive_for_different(self, pure_zero, pure_one):
        assert bures_distance(pure_zero, pure_one) > 0


class TestPureStateFidelity:
    def test_against_own_dm(self):
        psi = np.array([1.0, 0.0], dtype=complex)
        rho = np.outer(psi, psi.conj())
        assert abs(pure_state_fidelity(psi, rho) - 1.0) < 1e-10

    def test_against_orthogonal(self):
        psi = np.array([1.0, 0.0], dtype=complex)
        rho = np.array([[0, 0], [0, 1]], dtype=complex)
        assert pure_state_fidelity(psi, rho) < 1e-10


class TestVonNeumannEntropy:
    def test_pure_state_zero_entropy(self, pure_zero):
        assert von_neumann_entropy(pure_zero) < 1e-10

    def test_maximally_mixed_max_entropy(self, maximally_mixed_2):
        s = von_neumann_entropy(maximally_mixed_2)
        expected = np.log(2)  # log(d) for d=2
        assert abs(s - expected) < 1e-10

    def test_entropy_nonnegative(self, rng):
        for _ in range(10):
            rho = random_density_matrix(4, rng)
            assert von_neumann_entropy(rho) >= -1e-10


class TestPurity:
    def test_pure_state(self, pure_zero):
        assert abs(purity(pure_zero) - 1.0) < 1e-10

    def test_maximally_mixed(self, maximally_mixed_4):
        # Purity of I/4 = 1/4 = 0.25
        assert abs(purity(maximally_mixed_4) - 0.25) < 1e-10


class TestValidateDensityMatrix:
    def test_valid_pure_state(self, pure_zero):
        checks = validate_density_matrix(pure_zero)
        assert checks['valid']

    def test_valid_mixed_state(self, maximally_mixed_2):
        checks = validate_density_matrix(maximally_mixed_2)
        assert checks['valid']

    def test_invalid_trace(self):
        rho = np.eye(2, dtype=complex)  # Trace = 2, not 1
        checks = validate_density_matrix(rho)
        assert not checks['trace_one']
        assert not checks['valid']

    def test_invalid_not_hermitian(self):
        rho = np.array([[0.5, 1j], [0, 0.5]], dtype=complex)
        checks = validate_density_matrix(rho)
        assert not checks['hermitian']