from src.eval.error_analysis import (
    ErrorRecord,
    is_false_negative,
    is_false_positive,
    partition_by_context_requirement,
    partition_by_cwe_head_tail,
    partition_by_cwe_training_frequency,
    partition_by_dataflow_distance,
    partition_by_function_length_quartile,
    partition_by_project,
    sample_for_human_audit,
)


def _rec(**kw):
    base = dict(
        sample_id="s1", project="p", cwe="CWE-1", label_true=1, label_pred=1,
        func_len=10, cyclomatic_complexity=1, has_context=False, trace_gt_len=0,
        train_cwe_frequency=100,
    )
    base.update(kw)
    return ErrorRecord(**base)


def test_is_false_positive():
    assert is_false_positive(_rec(label_true=0, label_pred=1)) is True
    assert is_false_positive(_rec(label_true=1, label_pred=1)) is False


def test_is_false_negative():
    assert is_false_negative(_rec(label_true=1, label_pred=0)) is True
    assert is_false_negative(_rec(label_true=0, label_pred=0)) is False


def test_partition_by_project_computes_rates():
    records = [
        _rec(sample_id="a", project="proj1", label_true=1, label_pred=0),  # FN
        _rec(sample_id="b", project="proj1", label_true=1, label_pred=1),  # correct
        _rec(sample_id="c", project="proj2", label_true=0, label_pred=1),  # FP
    ]
    out = partition_by_project(records)
    assert out["proj1"]["n"] == 2
    assert out["proj1"]["fn_rate"] == 0.5
    assert out["proj2"]["fp_rate"] == 1.0


def test_partition_by_cwe_head_tail():
    records = [
        _rec(sample_id="a", cwe="CWE-1", label_true=1, label_pred=0),
        _rec(sample_id="b", cwe="CWE-2", label_true=1, label_pred=0),
    ]
    counts = {"CWE-1": 100, "CWE-2": 1}
    out = partition_by_cwe_head_tail(records, counts, head_k=1)
    assert out["head"]["n"] == 1
    assert out["tail"]["n"] == 1


def test_partition_by_cwe_head_tail_unknown_cwe():
    records = [_rec(sample_id="a", cwe=None)]
    out = partition_by_cwe_head_tail(records, {}, head_k=12)
    assert out["unknown"]["n"] == 1


def test_partition_by_function_length_quartile_groups_all_records():
    records = [_rec(sample_id=f"s{i}", func_len=i) for i in range(1, 21)]
    out = partition_by_function_length_quartile(records)
    total = sum(v["n"] for v in out.values())
    assert total == 20
    assert set(out.keys()) <= {"Q1", "Q2", "Q3", "Q4"}


def test_partition_by_context_requirement():
    records = [
        _rec(sample_id="a", has_context=True),
        _rec(sample_id="b", has_context=False),
    ]
    out = partition_by_context_requirement(records)
    assert out["needs_callee"]["n"] == 1
    assert out["intra_function"]["n"] == 1


def test_partition_by_dataflow_distance_buckets():
    records = [
        _rec(sample_id="a", trace_gt_len=0),
        _rec(sample_id="b", trace_gt_len=2),
        _rec(sample_id="c", trace_gt_len=4),
        _rec(sample_id="d", trace_gt_len=10),
    ]
    out = partition_by_dataflow_distance(records)
    assert out["no_trace_gt"]["n"] == 1
    assert out["short (<=2 hops)"]["n"] == 1
    assert out["medium (3-5 hops)"]["n"] == 1
    assert out["long (>5 hops)"]["n"] == 1


def test_partition_by_cwe_training_frequency():
    records = [
        _rec(sample_id="a", cwe="CWE-1", train_cwe_frequency=5),
        _rec(sample_id="b", cwe="CWE-2", train_cwe_frequency=500),
    ]
    out = partition_by_cwe_training_frequency(records, low_freq_threshold=10)
    assert out["low_frequency"]["n"] == 1
    assert out["high_frequency"]["n"] == 1


def test_sample_for_human_audit_respects_per_split_cap():
    by_split = {
        "chronological": [f"c{i}" for i in range(100)],
        "repo_disjoint": [f"r{i}" for i in range(100)],
        "ood": [f"o{i}" for i in range(10)],  # fewer than n_per_split available
    }
    sample = sample_for_human_audit(by_split, n_per_split=50, seed=1)
    by_split_count = {}
    for s in sample:
        by_split_count[s.split_name] = by_split_count.get(s.split_name, 0) + 1
    assert by_split_count["chronological"] == 50
    assert by_split_count["repo_disjoint"] == 50
    assert by_split_count["ood"] == 10  # capped at what's available, no crash


def test_sample_for_human_audit_deterministic_given_seed():
    by_split = {"chronological": [f"c{i}" for i in range(100)]}
    a = sample_for_human_audit(by_split, n_per_split=50, seed=7)
    b = sample_for_human_audit(by_split, n_per_split=50, seed=7)
    assert [s.sample_id for s in a] == [s.sample_id for s in b]