"""Error analysis (spec §12): partitions FP/FN by CWE (head/tail), project,
function length quartile, cyclomatic complexity, context requirement
(intra-function vs needs-callee), data-flow distance (trace_gt length),
training-frequency of CWE, and trace-validity of the emitted explanation.

Pure partitioning/aggregation functions over prediction records; plotting and
the human-audit sampling tool are separate (matplotlib-dependent) concerns
layered on top, not needed for the partition logic itself to be correct and
testable.
"""
from __future__ import annotations

import random
from collections import defaultdict
from dataclasses import dataclass, field

from src.eval.metrics import head_tail_cwe_split


@dataclass
class ErrorRecord:
    sample_id: str
    project: str
    cwe: str | None
    label_true: int
    label_pred: int
    func_len: int
    cyclomatic_complexity: int
    has_context: bool
    trace_gt_len: int  # 0 if no trace ground truth
    train_cwe_frequency: int = 0


def is_false_positive(r: ErrorRecord) -> bool:
    return r.label_true == 0 and r.label_pred == 1


def is_false_negative(r: ErrorRecord) -> bool:
    return r.label_true == 1 and r.label_pred == 0


def partition_by_cwe_head_tail(
    records: list[ErrorRecord], cwe_counts: dict[str, int], head_k: int = 12
) -> dict[str, dict]:
    head, tail = head_tail_cwe_split(cwe_counts, head_k=head_k)

    def bucket(r: ErrorRecord) -> str:
        if r.cwe is None:
            return "unknown"
        return "head" if r.cwe in head else "tail"

    return _error_rates_by_key(records, bucket)


def partition_by_project(records: list[ErrorRecord]) -> dict[str, dict]:
    return _error_rates_by_key(records, lambda r: r.project)


def _quartile_labels(values: list[int]) -> list[str]:
    if not values:
        return []
    sorted_vals = sorted(values)
    n = len(sorted_vals)
    q1 = sorted_vals[n // 4]
    q2 = sorted_vals[n // 2]
    q3 = sorted_vals[(3 * n) // 4]

    def label(v: int) -> str:
        if v <= q1:
            return "Q1"
        if v <= q2:
            return "Q2"
        if v <= q3:
            return "Q3"
        return "Q4"

    return [label(v) for v in values]


def partition_by_function_length_quartile(records: list[ErrorRecord]) -> dict[str, dict]:
    lengths = [r.func_len for r in records]
    labels = _quartile_labels(lengths)
    label_by_id = dict(zip((r.sample_id for r in records), labels))
    return _error_rates_by_key(records, lambda r: label_by_id[r.sample_id])


def partition_by_cyclomatic_complexity_quartile(records: list[ErrorRecord]) -> dict[str, dict]:
    values = [r.cyclomatic_complexity for r in records]
    labels = _quartile_labels(values)
    label_by_id = dict(zip((r.sample_id for r in records), labels))
    return _error_rates_by_key(records, lambda r: label_by_id[r.sample_id])


def partition_by_context_requirement(records: list[ErrorRecord]) -> dict[str, dict]:
    return _error_rates_by_key(records, lambda r: "needs_callee" if r.has_context else "intra_function")


def partition_by_dataflow_distance(records: list[ErrorRecord]) -> dict[str, dict]:
    def bucket(r: ErrorRecord) -> str:
        if r.trace_gt_len == 0:
            return "no_trace_gt"
        if r.trace_gt_len <= 2:
            return "short (<=2 hops)"
        if r.trace_gt_len <= 5:
            return "medium (3-5 hops)"
        return "long (>5 hops)"

    return _error_rates_by_key(records, bucket)


def partition_by_cwe_training_frequency(
    records: list[ErrorRecord], low_freq_threshold: int = 10
) -> dict[str, dict]:
    def bucket(r: ErrorRecord) -> str:
        if r.cwe is None:
            return "unknown"
        return "low_frequency" if r.train_cwe_frequency < low_freq_threshold else "high_frequency"

    return _error_rates_by_key(records, bucket)


def _error_rates_by_key(records: list[ErrorRecord], key_fn) -> dict[str, dict]:
    groups: dict[str, list[ErrorRecord]] = defaultdict(list)
    for r in records:
        groups[key_fn(r)].append(r)

    out: dict[str, dict] = {}
    for key, group in groups.items():
        n = len(group)
        n_fp = sum(1 for r in group if is_false_positive(r))
        n_fn = sum(1 for r in group if is_false_negative(r))
        n_correct = sum(1 for r in group if r.label_true == r.label_pred)
        out[key] = {
            "n": n,
            "fp_rate": n_fp / n if n else 0.0,
            "fn_rate": n_fn / n if n else 0.0,
            "accuracy": n_correct / n if n else 0.0,
        }
    return out


@dataclass
class HumanAuditSample:
    sample_id: str
    split_name: str


def sample_for_human_audit(
    by_split: dict[str, list[str]], n_per_split: int = 50, seed: int = 42
) -> list[HumanAuditSample]:
    """spec §12 human audit protocol: "sample 150 model outputs (50 each
    from chronological/repo-disjoint/OOD)". `by_split` maps split name to the
    sample_ids available in it (typically outputs from the final TraceGRPO
    model's predictions on that split's test set)."""
    rng = random.Random(seed)
    out: list[HumanAuditSample] = []
    for split_name, sample_ids in by_split.items():
        chosen = sample_ids[:]
        rng.shuffle(chosen)
        for sid in chosen[:n_per_split]:
            out.append(HumanAuditSample(sample_id=sid, split_name=split_name))
    return out