import numpy as np
import pytest
import torch

from deception_detection.metrics import (
    get_auroc,
    get_auroc_manual,
    get_fpr_tpr,
    get_tpr_at_fpr_from_control_dataset,
    get_tpr_at_fpr_from_paired_dataset,
)


def test_auroc_random_data():
    n_samples = 500
    np.random.seed(42)
    honest_scores = torch.tensor(np.random.normal(loc=0.3, scale=0.2, size=n_samples).clip(0, 1))
    deceptive_scores = torch.tensor(np.random.normal(loc=0.7, scale=0.2, size=n_samples).clip(0, 1))

    auroc_custom = get_auroc_manual(honest_scores, deceptive_scores, N=200)
    auroc_sklearn = get_auroc(honest_scores, deceptive_scores)

    assert pytest.approx(auroc_custom, abs=1e-2) == auroc_sklearn
    assert 0.5 < auroc_sklearn < 1.0  # Should be better than random but not perfect


def test_auroc_perfect_separation():
    n_samples = 100
    honest_scores = torch.zeros(n_samples)
    deceptive_scores = torch.ones(n_samples)

    auroc_custom = get_auroc_manual(honest_scores, deceptive_scores, N=200)
    auroc_sklearn = get_auroc(honest_scores, deceptive_scores)

    assert pytest.approx(auroc_custom, abs=1e-4) == 1.0
    assert pytest.approx(auroc_sklearn, abs=1e-4) == 1.0


def test_auroc_complex_case():
    honest_scores = torch.cat([torch.zeros(90), torch.ones(10)])
    deceptive_scores = torch.cat([torch.zeros(10), torch.ones(90)])

    auroc_custom = get_auroc_manual(honest_scores, deceptive_scores, N=200)
    auroc_sklearn = get_auroc(honest_scores, deceptive_scores)

    assert pytest.approx(auroc_custom, abs=1e-2) == 0.9
    assert pytest.approx(auroc_sklearn, abs=1e-4) == 0.9


def test_auroc_no_separation():
    n_samples = 500
    np.random.seed(42)
    scores = torch.tensor(np.random.uniform(0, 1, size=2 * n_samples))
    honest_scores = scores[:n_samples]
    deceptive_scores = scores[n_samples:]

    auroc_custom = get_auroc_manual(honest_scores, deceptive_scores, N=200)
    auroc_sklearn = get_auroc(honest_scores, deceptive_scores)

    assert pytest.approx(auroc_sklearn, abs=1e-1) == 0.5
    assert pytest.approx(auroc_custom, abs=1e-1) == 0.5


def test_get_fpr_tpr():
    honest_scores = torch.tensor([0.1, 0.2, 0.3, 0.4, 0.5])
    deceptive_scores = torch.tensor([0.6, 0.7, 0.8, 0.9, 1.0])

    fpr, tpr, thresholds = get_fpr_tpr(honest_scores, deceptive_scores)

    assert len(fpr) == len(tpr) == len(thresholds)
    assert all(0 <= f <= 1 for f in fpr)
    assert all(0 <= t <= 1 for t in tpr)
    assert fpr[0] == tpr[0] == 0
    assert fpr[-1] == tpr[-1] == 1


def test_get_tpr_at_fpr_from_paired_dataset():
    np.random.seed(42)
    n_samples = 1000
    honest_scores = torch.tensor(np.random.normal(loc=0.3, scale=0.2, size=n_samples).clip(0, 1))
    deceptive_scores = torch.tensor(np.random.normal(loc=0.7, scale=0.2, size=n_samples).clip(0, 1))

    tpr_at_fpr_01 = get_tpr_at_fpr_from_paired_dataset(honest_scores, deceptive_scores, 0.1)
    tpr_at_fpr_05 = get_tpr_at_fpr_from_paired_dataset(honest_scores, deceptive_scores, 0.5)

    assert 0 < tpr_at_fpr_01 < tpr_at_fpr_05 < 1


def test_get_tpr_at_fpr_from_control_dataset():
    np.random.seed(42)
    n_samples = 1000
    deceptive_scores = torch.tensor(np.random.normal(loc=0.7, scale=0.2, size=n_samples).clip(0, 1))

    tpr_at_threshold_05 = get_tpr_at_fpr_from_control_dataset(deceptive_scores, 0.5)
    tpr_at_threshold_08 = get_tpr_at_fpr_from_control_dataset(deceptive_scores, 0.8)

    assert 0 < tpr_at_threshold_08 < tpr_at_threshold_05 < 1


def test_edge_cases():
    # Test with empty tensors
    with pytest.raises(ValueError):
        get_auroc(torch.tensor([]), torch.tensor([]))

    # Test with single element tensors
    assert get_auroc(torch.tensor([0]), torch.tensor([1])) == 1.0
    assert get_auroc(torch.tensor([1]), torch.tensor([0])) == 0.0

    # Test with all identical scores
    n_samples = 100
    identical_scores = torch.ones(n_samples)
    assert pytest.approx(get_auroc(identical_scores, identical_scores), abs=1e-4) == 0.5


def test_auroc_consistency():
    n_samples = 1000
    np.random.seed(42)
    honest_scores = torch.tensor(np.random.beta(2, 5, size=n_samples))
    deceptive_scores = torch.tensor(np.random.beta(5, 2, size=n_samples))

    auroc1 = get_auroc(honest_scores, deceptive_scores)
    auroc2 = get_auroc(honest_scores, deceptive_scores)

    assert auroc1 == auroc2


if __name__ == "__main__":
    pytest.main([__file__])
