"""Tests for bex/reduce.py — Algorithm 4 (TODS 2010)."""

import pytest
from bex.reduce import (
    build_soa_with_support, soa_distance, adjunct_support,
    reduce_contexts, minimize_contexts, reduce_and_infer,
)
from bex.twotinf import build_soa
from bex.rwr0 import rwr0


class TestBuildSoaWithSupport:
    """Test support-annotated SOA building."""

    def test_single_sequence(self):
        seqs = [["a", "b", "c"]]
        soa, supp = build_soa_with_support(seqs)
        assert len(supp) > 0
        # Should have edges: (src, a), (a, b), (b, c), (c, sink)
        assert len(supp) == 4

    def test_multiple_sequences_share_edges(self):
        seqs = [["a", "b"], ["a", "b"]]
        soa, supp = build_soa_with_support(seqs)
        # Support keys are labeled tuples: (('SRC',), ('a',)), etc.
        src_key = (('SRC',), ('a',))
        assert src_key in supp
        assert supp[src_key] == 2

    def test_different_sequences(self):
        seqs = [["a", "b"], ["a", "c"]]
        soa, supp = build_soa_with_support(seqs)
        # Both share (src, a) but diverge after
        src_key = (('SRC',), ('a',))
        assert src_key in supp
        assert supp[src_key] == 2

    def test_empty_sequence(self):
        seqs = [[]]
        soa, supp = build_soa_with_support(seqs)
        empty_key = (('SRC',), ('SINK',))
        assert empty_key in supp


class TestSoaDistance:
    """Test SOA edit distance (Definition 14)."""

    def test_identical_soas(self):
        seqs_a = [["a", "b"]]
        seqs_b = [["a", "b"]]
        _, supp_a = build_soa_with_support(seqs_a)
        _, supp_b = build_soa_with_support(seqs_b)
        assert soa_distance(supp_a, supp_b) == 0.0

    def test_completely_disjoint(self):
        seqs_a = [["a", "b"]]
        seqs_b = [["x", "y"]]
        _, supp_a = build_soa_with_support(seqs_a)
        _, supp_b = build_soa_with_support(seqs_b)
        dist = soa_distance(supp_a, supp_b)
        # They share src->first and last->sink edges
        # But the middle edges are completely different
        assert dist > 0.5

    def test_partial_overlap(self):
        seqs_a = [["a", "b", "c"]]
        seqs_b = [["a", "b", "d"]]
        _, supp_a = build_soa_with_support(seqs_a)
        _, supp_b = build_soa_with_support(seqs_b)
        dist = soa_distance(supp_a, supp_b)
        # Each SOA has 2 unique edges out of 4 total
        # dist = 2/4 + 2/4 = 1.0 (formula sums both sides)
        assert dist == 1.0

    def test_asymmetric_support(self):
        """One SOA with high support shared edges, other with unique edges."""
        seqs_a = [["a", "b"]] * 10  # (a,b) has support 10
        seqs_b = [["a", "c"]]       # (a,c) has support 1
        _, supp_a = build_soa_with_support(seqs_a)
        _, supp_b = build_soa_with_support(seqs_b)
        dist = soa_distance(supp_a, supp_b)
        # dist_a = 0 (all A edges in B? no - (a,b) not in B)
        # Actually: E_A = {(SRC,a):10, (a,b):10, (b,SINK):10}
        #           F_B = {(SRC,a):1, (a,c):1, (c,SINK):1}
        # E-F = {(a,b):10, (b,SINK):10} → only_a = 20, total_a = 30 → 2/3
        # F-E = {(a,c):1, (c,SINK):1} → only_b = 2, total_b = 3 → 2/3
        # dist = 2/3 + 2/3 ≈ 1.33... wait, that can't be right
        # Actually dist can exceed 1.0 because it's two separate fractions summed
        assert dist > 0.5

    def test_symmetry(self):
        seqs_a = [["a", "b", "c"]]
        seqs_b = [["x", "y", "z"]]
        _, supp_a = build_soa_with_support(seqs_a)
        _, supp_b = build_soa_with_support(seqs_b)
        assert soa_distance(supp_a, supp_b) == soa_distance(supp_b, supp_a)

    def test_support_weighted(self):
        # Edge with high support should contribute more to distance
        seqs_a = [["a", "b"]] * 10 + [["a", "c"]]
        seqs_b = [["a", "b"]] * 10 + [["a", "d"]]
        _, supp_a = build_soa_with_support(seqs_a)
        _, supp_b = build_soa_with_support(seqs_b)
        dist = soa_distance(supp_a, supp_b)
        # Distance should be small because (a,b) has high support in both
        assert dist < 0.2


class TestAdjunctSupport:
    """Test SOA adjunction (line 6, Algorithm 4)."""

    def test_basic_adjunction(self):
        supp_a = {(0, 1): 5, (1, 2): 3}
        supp_b = {(0, 1): 2, (1, 3): 4}
        combined = adjunct_support(supp_a, supp_b)
        assert combined[(0, 1)] == 7  # 5 + 2
        assert combined[(1, 2)] == 3
        assert combined[(1, 3)] == 4

    def test_disjoint_edges(self):
        supp_a = {(0, 1): 1}
        supp_b = {(2, 3): 1}
        combined = adjunct_support(supp_a, supp_b)
        assert len(combined) == 2

    def test_empty(self):
        supp_a = {}
        supp_b = {}
        combined = adjunct_support(supp_a, supp_b)
        assert combined == {}


class TestReduceContexts:
    """Test Reduce algorithm (Algorithm 4)."""

    def test_no_merge_high_threshold(self):
        """With very high threshold, nothing should merge."""
        contexts = {
            "a": [["x", "y"]] * 5,
            "b": [["p", "q"]] * 5,
        }
        merged, info = reduce_contexts(contexts, threshold=0.99)
        assert info["merges"] == 0
        assert len(merged) == 2

    def test_no_merge_identical(self):
        """Identical contexts (dist=0) are NOT merged by Reduce.
        
        From the paper: M := {(s, t) | 0 < distD(s, t) < ε}
        dist=0 means already equivalent → unified by Minimize, not Reduce.
        """
        contexts = {
            "a": [["x", "y"]] * 5,
            "b": [["x", "y"]] * 5,
        }
        merged, info = reduce_contexts(contexts, threshold=0.5)
        assert info["merges"] == 0  # No merge for identical contexts
        assert len(merged) == 2

    def test_merge_similar(self):
        """Similar contexts (0 < dist < threshold) should merge.
        
        Use asymmetric support: many shared sequences, few unique ones.
        This makes the unique edges have low support relative to shared edges.
        """
        seqs_a = [["a", "b"]] * 10 + [["a", "c"]]  # 10 shared, 1 unique
        seqs_b = [["a", "b"]] * 10 + [["a", "d"]]  # 10 shared, 1 unique
        contexts = {
            "a": seqs_a,
            "b": seqs_b,
        }
        merged, info = reduce_contexts(contexts, threshold=0.5)
        # dist ≈ 0.12 (2/33 + 2/33), well below threshold 0.5
        assert info["merges"] >= 1

    def test_small_contexts_preserved(self):
        """Contexts with < 2 sequences should be preserved."""
        contexts = {
            "a": [["x", "y"]],  # Only 1 sequence
            "b": [["p", "q"]] * 5,
        }
        merged, info = reduce_contexts(contexts, threshold=0.5)
        assert "a" in merged  # Preserved
        assert "b" in merged

    def test_merge_info_contains_log(self):
        """Merge log should record each merge."""
        contexts = {
            "a": [["a", "b", "c"]] * 5,
            "b": [["a", "b", "d"]] * 5,
        }
        merged, info = reduce_contexts(contexts, threshold=0.5)
        if info["merges"] > 0:
            assert len(info["merge_log"]) == 1
            assert "merged_into" in info["merge_log"][0]
            assert "removed" in info["merge_log"][0]

    def test_iterative_merging(self):
        """Multiple iterations should work correctly."""
        # Create contexts that form a chain: a~b, b~c, but a!~c
        contexts = {
            "a": [["a", "b"]] * 5,
            "b": [["a", "b", "c"]] * 5,
            "c": [["a", "b", "c", "d"]] * 5,
        }
        merged, info = reduce_contexts(contexts, threshold=0.6)
        # Should merge at least some pairs
        assert info["iterations"] >= 1


class TestMinimizeContexts:
    """Test Minimize (Line 15, Algorithm 4)."""

    def test_merge_same_sore(self):
        """Contexts with identical SOREs should be merged."""
        # Both contexts produce the same SORE
        contexts = {
            "a": [["mockk"]] * 5,
            "b": [["mockk"]] * 5,
        }
        minimized = minimize_contexts(contexts)
        # Should merge into one context
        assert len(minimized) == 1
        total_seqs = sum(len(v) for v in minimized.values())
        assert total_seqs == 10

    def test_keep_different_sores(self):
        """Contexts with different SOREs should stay separate."""
        contexts = {
            "a": [["a", "b"]] * 5,
            "b": [["x", "y"]] * 5,
        }
        minimized = minimize_contexts(contexts)
        assert len(minimized) == 2

    def test_empty_contexts(self):
        """Empty contexts should be handled gracefully."""
        contexts = {}
        minimized = minimize_contexts(contexts)
        assert len(minimized) == 0


class TestReduceAndInfer:
    """Test full pipeline: Reduce + Minimize + Infer."""

    def test_basic_pipeline(self):
        contexts = {
            "a": [["a", "b", "c"]] * 5,
            "b": [["a", "b", "d"]] * 5,
        }
        result = reduce_and_infer(contexts, threshold=0.5)
        assert "merged" in result
        assert "infer_results" in result
        assert "merge_info" in result
        assert "minimize_info" in result

    def test_coverage_count(self):
        """Coverage count should equal total methods in successful contexts."""
        contexts = {
            "a": [["a", "b"]] * 5,
            "b": [["x", "y"]] * 5,
        }
        result = reduce_and_infer(contexts, threshold=0.99)  # No merging
        assert result["coverage_count"] >= 0