"""
Unit tests for key computational functions in mimosa.

These tests validate the correctness of individual functions from:
- mimosa/functions.py
- mimosa/comparison.py
- mimosa/models.py
"""

import numpy as np
import pytest

from mimosa.comparison import DataComparator, TomtomComparator, UniversalMotifComparator
from mimosa.functions import (
    batch_all_scores,
    cut_prc,
    cut_roc,
    format_params,
    pcm_to_pfm,
    pfm_to_pwm,
    precision_recall_curve,
    roc_curve,
    score_seq,
    scores_to_frequencies,
    standardized_pauc,
)
from mimosa.models import BammMotif, PwmMotif, RaggedScores, SitegaMotif
from mimosa.ragged import RaggedData


def test_pfm_to_pwm_basic():
    """Test basic PFM to PWM conversion"""
    # Create a simple PFM with uniform values
    pfm = np.array([[0.25, 0.25], [0.25, 0.25], [0.25, 0.25], [0.25, 0.25]])

    pwm = pfm_to_pwm(pfm)

    # Verify shape is preserved
    assert pwm.shape == pfm.shape

    # Verify that PWM values are log ratios
    expected = np.log((pfm + 0.0001) / 0.25)
    np.testing.assert_allclose(pwm, expected, rtol=1e-6)


def test_pfm_to_pwm_with_zeros():
    """Test PFM to PWM conversion with zeros"""
    pfm = np.array([[0.0, 0.5], [0.0, 0.0], [0.5, 0.0], [0.5, 0.5]])

    pwm = pfm_to_pwm(pfm)

    # Verify shape is preserved
    assert pwm.shape == pfm.shape

    # Verify that PWM values are calculated correctly even with zeros
    expected = np.log((pfm + 0.0001) / 0.25)
    np.testing.assert_allclose(pwm, expected, rtol=1e-6)


def test_pcm_to_pfm_basic():
    """Test basic PCM to PFM conversion"""
    pcm = np.array([[2, 3], [1, 1], [1, 0], [0, 0]], dtype=float)

    pfm = pcm_to_pfm(pcm)

    # Verify shape is preserved
    assert pfm.shape == pcm.shape

    # Calculate expected manually
    number_of_sites = pcm.sum(axis=0)
    expected = (pcm + 0.25) / (number_of_sites + 1)
    np.testing.assert_allclose(pfm, expected, rtol=1e-6)


def test_score_seq_basic():
    """Test basic sequence scoring"""
    # Simple scoring model
    model = np.array(
        [
            [1.0, 2.0, 3.0],
            [1.0, 2.0, 3.0],
            [1.0, 2.0, 3.0],
        ]
    )
    # DNA sequence as numerical representation [A, C, G, T] -> [0, 1, 2, 3]
    num_site = np.array([0, 1, 2], dtype=np.int8)
    kmer = 1

    score = score_seq(num_site, kmer, model)

    # With kmer=1, the function should compute the sum of model values at positions
    expected_score = 1.0 + 2.0 + 3.0
    assert score == expected_score


def test_precision_recall_curve_basic():
    """Test basic precision-recall curve calculation"""
    classification = np.array([1, 0, 1, 1, 0])
    scores = np.array([0.9, 0.8, 0.7, 0.6, 0.5])

    precision, recall, thresholds = precision_recall_curve(classification, scores)

    # Verify shapes are consistent
    assert precision.shape == recall.shape == thresholds.shape
    # Verify bounds
    assert np.all(precision >= 0) and np.all(precision <= 1.1)  # Allow for small numerical errors
    assert np.all(recall >= 0) and np.all(recall <= 1.1)


def test_roc_curve_basic():
    """Test basic ROC curve calculation"""
    classification = np.array([1, 0, 1, 1, 0])
    scores = np.array([0.9, 0.8, 0.7, 0.6, 0.5])

    tpr, fpr, thresholds = roc_curve(classification, scores)

    # Verify shapes are consistent
    assert tpr.shape == fpr.shape == thresholds.shape
    # Verify bounds
    assert np.all(tpr >= 0) and np.all(tpr <= 1.1)
    assert np.all(fpr >= 0) and np.all(fpr <= 1.1)


def test_cut_roc_basic():
    """Test basic ROC curve cutting"""
    tpr = np.array([0.0, 0.5, 1.0])
    fpr = np.array([0.0, 0.2, 0.5])
    thr = np.array([np.inf, 0.5, 0.0])
    score_cutoff = 0.6

    tpr_cut, fpr_cut, thr_cut = cut_roc(tpr, fpr, thr, score_cutoff)

    # Verify that the function returns arrays
    assert isinstance(tpr_cut, np.ndarray)
    assert isinstance(fpr_cut, np.ndarray)
    assert isinstance(thr_cut, np.ndarray)


def test_cut_prc_basic():
    """Test basic PRC curve cutting"""
    rec = np.array([0.0, 0.5, 1.0])
    prec = np.array([1.0, 0.8, 0.6])
    thr = np.array([np.inf, 0.5, 0.0])
    score_cutoff = 0.6

    rec_cut, prec_cut, thr_cut = cut_prc(rec, prec, thr, score_cutoff)

    # Verify that the function returns arrays
    assert isinstance(rec_cut, np.ndarray)
    assert isinstance(prec_cut, np.ndarray)
    assert isinstance(thr_cut, np.ndarray)


def test_standardized_pauc_basic():
    """Test basic standardized partial AUC calculation"""
    pauc_raw = 0.7
    pauc_min = 0.5
    pauc_max = 1.0

    standardized = standardized_pauc(pauc_raw, pauc_min, pauc_max)

    # Manual calculation: 0.5 * (1.0 + (0.7 - 0.5) / (1.0 - 0.5))
    expected = 0.5 * (1.0 + (0.7 - 0.5) / (1.0 - 0.5))
    assert abs(standardized - expected) < 1e-6


def test_scores_to_frequencies_basic():
    """Test basic scores to frequencies conversion"""
    # Create a simple RaggedData with some scores
    data = np.array([1.0, 2.0, 1.0, 3.0, 2.0], dtype=np.float32)
    offsets = np.array([0, 2, 4, 5], dtype=np.int64)  # Two sequences of lengths 2, 2, 1
    ragged_scores = RaggedData(data, offsets)

    freq_result = scores_to_frequencies(ragged_scores)

    # Verify the output is a RaggedData
    assert isinstance(freq_result, RaggedData)
    # Verify shape consistency
    assert freq_result.data.shape == ragged_scores.data.shape
    assert freq_result.offsets.shape == ragged_scores.offsets.shape


def test_format_params_basic():
    """Test basic parameter formatting"""
    params = {"k": 3, "metric": "pcc", "n_perm": 1000}
    formatted = format_params(params)

    # Should be in alphabetical order: k-3_metric-pcc_n_perm-1000
    expected = "k-3_metric-pcc_n_perm-1000"
    assert formatted == expected


def test_tomtom_comparator_initialization():
    """Test TomtomComparator initialization"""
    comparator = TomtomComparator(metric="pcc", n_permutations=100)

    assert comparator.metric == "pcc"
    assert comparator.n_permutations == 100
    assert comparator.name == "TomtomComparator_PCC"


def test_tomtom_comparator_invalid_metric():
    """Test TomtomComparator with invalid metric raises error"""
    with pytest.raises(ValueError):
        TomtomComparator(metric="invalid")


def test_data_comparator_initialization():
    """Test DataComparator initialization"""
    comparator = DataComparator(metric="cj", n_permutations=100)

    assert comparator.metric == "cj"
    assert comparator.n_permutations == 100
    assert comparator.name == "DataComparator"


def test_universal_motif_comparator_initialization():
    """Test UniversalMotifComparator initialization"""
    comparator = UniversalMotifComparator(metric="co", n_permutations=100)

    assert comparator.metric == "co"
    assert comparator.n_permutations == 100
    assert comparator.name == "UniversalMotifComparator"


def test_universal_motif_comparator_invalid_metric():
    """Test UniversalMotifComparator with invalid metric raises error"""
    with pytest.raises(ValueError):
        UniversalMotifComparator(metric="invalid")


def test_ragged_scores_from_numba():
    """Test RaggedScores creation from Numba RaggedData"""
    # Create mock RaggedData
    data = np.array([1.0, 2.0, 3.0, 4.0], dtype=np.float32)
    offsets = np.array([0, 2, 4], dtype=np.int64)
    ragged_data = RaggedData(data, offsets)

    # Create RaggedScores from RaggedData
    ragged_scores = RaggedScores.from_numba(ragged_data)

    assert isinstance(ragged_scores, RaggedScores)
    assert ragged_scores.values.shape == (2, 2)  # 2 sequences, max length 2
    assert ragged_scores.lengths.shape == (2,)  # 2 sequence lengths


def test_pwm_motif_creation():
    """Test PwmMotif creation"""
    matrix = np.array([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0], [7.0, 8.0], [1.0, 2.0]])  # Extended PWM
    pfm = np.array([[0.25, 0.25], [0.25, 0.25], [0.25, 0.25], [0.25, 0.25]])  # Standard PFM

    motif = PwmMotif(matrix=matrix, name="test_pwm", length=2, pfm=pfm)

    assert motif.name == "test_pwm"
    assert motif.length == 2
    assert motif.model_type == "pwm"
    assert motif.kmer == 1  # Default value


def test_bamm_motif_creation():
    """Test BammMotif creation"""
    matrix = np.random.rand(4, 4, 6)  # Example 2nd order BaMM

    motif = BammMotif(matrix=matrix, name="test_bamm", length=6, kmer=3)

    assert motif.name == "test_bamm"
    assert motif.length == 6
    assert motif.model_type == "bamm"
    assert motif.kmer == 3


def test_sitega_motif_creation():
    """Test SitegaMotif creation"""
    matrix = np.random.rand(16, 8)  # Example SiteGA matrix

    motif = SitegaMotif(matrix=matrix, name="test_sitega", length=8, kmer=2)

    assert motif.name == "test_sitega"
    assert motif.length == 8
    assert motif.model_type == "sitega"
    assert motif.kmer == 2


def test_batch_all_scores_with_simple_data():
    """Test batch_all_scores with simple RaggedData"""
    # Create proper RaggedData for testing
    data = np.array([0, 1, 2, 3, 0, 1], dtype=np.int8)
    offsets = np.array([0, 3, 6], dtype=np.int64)
    sequences = RaggedData(data, offsets)
    matrix = np.array([[1.0, 2.0], [3.0, 4.0]], dtype=np.float32)

    result = batch_all_scores(sequences, matrix, kmer=1, is_revcomp=False)
    assert hasattr(result, "data") and hasattr(result, "offsets")


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