"""
Tests for Quantum Backend Module.
"""

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

import torch
import pytest
from src.quantum_backend import QuantumBackend, BackendType, ClassicalSurrogateUnitary


def test_classical_surrogate_unitary():
    surrogate = ClassicalSurrogateUnitary(n_qubits=4, n_layers=2, n_outputs=4)
    x = torch.randn(2, 6, 4)
    out = surrogate(x)
    assert out.shape == (2, 6, 4)
    # Pauli-Z expectation values bounded in [-1, 1]
    assert out.min() >= -1.01
    assert out.max() <= 1.01
    print("✓ test_classical_surrogate_unitary passed")


def test_quantum_backend_execution():
    backend = QuantumBackend(
        backend_type=BackendType.CLASSICAL_SURROGATE,
        n_qubits=4,
        n_layers=2,
        d_model=64,
    )
    x = torch.randn(2, 8, 64)
    out, meta = backend(x)
    assert out.shape == (2, 8, 64)
    assert meta["classification"] in ["MEASURED", "SIMULATED"]
    print("✓ test_quantum_backend_execution passed")


def test_quantum_kernel_fidelity():
    backend = QuantumBackend(
        backend_type=BackendType.CLASSICAL_SURROGATE,
        n_qubits=4,
        n_layers=2,
        d_model=32,
    )
    B, H, T_q, T_k, D = 1, 2, 4, 4, 8
    q = torch.randn(B, H, T_q, D)
    k = torch.randn(B, H, T_k, D)

    K = backend.compute_kernel_matrix(q, k)
    assert K.shape == (B, H, T_q, T_k)
    # Self-fidelity K(q, q) should be close to 1.0
    K_self = backend.compute_kernel_matrix(q, q)
    diag = torch.diagonal(K_self[0, 0])
    assert torch.all(diag > 0.8), f"Self-fidelity diagonal should be near 1.0, got {diag}"
    print("✓ test_quantum_kernel_fidelity passed")


if __name__ == "__main__":
    test_classical_surrogate_unitary()
    test_quantum_backend_execution()
    test_quantum_kernel_fidelity()
    print("All Quantum Backend tests passed!")