"""Validation metrics for comparing discovered vs ground truth DAGs.

Metrics: SHD, edge F1, orientation accuracy, causal sufficiency, FDR.
"""

from __future__ import annotations

from dataclasses import dataclass


@dataclass(frozen=True)
class ValidationMetrics:
    """Results of comparing predicted edges to ground truth."""

    shd: int
    precision: float
    recall: float
    f1: float
    orientation_accuracy: float
    fdr: float
    true_positives: int
    false_positives: int
    false_negatives: int
    reversed_edges: int
    total_true_edges: int
    total_predicted_edges: int


def compute_metrics(
    true_edges: list[tuple[str, str]],
    predicted_edges: list[tuple[str, str]],
) -> ValidationMetrics:
    """Compare predicted DAG edges against ground truth.

    Args:
        true_edges: Ground truth (source, target) pairs.
        predicted_edges: Discovered (source, target) pairs.

    Returns:
        ValidationMetrics with SHD, F1, orientation accuracy, FDR.
    """
    true_set = set(true_edges)
    pred_set = set(predicted_edges)

    # Undirected versions for skeleton comparison.
    true_undirected = {_undirected(e) for e in true_edges}
    pred_undirected = {_undirected(e) for e in predicted_edges}

    # Skeleton matches (edge exists in both, ignoring direction).
    skeleton_tp = true_undirected & pred_undirected
    skeleton_fp = pred_undirected - true_undirected
    skeleton_fn = true_undirected - pred_undirected

    # Directed matches.
    tp = len(true_set & pred_set)

    # Reversed edges: skeleton match but wrong direction.
    reversed_count = 0
    for edge in predicted_edges:
        u_edge = _undirected(edge)
        if u_edge in true_undirected and edge not in true_set:
            rev = (edge[1], edge[0])
            if rev in true_set:
                reversed_count += 1

    # SHD = missing + extra + reversed.
    shd = len(skeleton_fn) + len(skeleton_fp) + reversed_count

    # Precision, recall, F1 on directed edges.
    total_pred = len(pred_set)
    total_true = len(true_set)

    precision = tp / total_pred if total_pred > 0 else 0.0
    recall = tp / total_true if total_true > 0 else 0.0
    f1 = 2 * precision * recall / (precision + recall) if (precision + recall) > 0 else 0.0

    # Orientation accuracy: of skeleton matches, fraction correctly oriented.
    correctly_oriented = 0
    for u_edge in skeleton_tp:
        a, b = u_edge
        if (a, b) in true_set and (a, b) in pred_set:
            correctly_oriented += 1
        elif (b, a) in true_set and (b, a) in pred_set:
            correctly_oriented += 1

    orientation_acc = correctly_oriented / len(skeleton_tp) if skeleton_tp else 1.0

    # FDR = false positives / total predicted.
    fp = total_pred - tp
    fdr = fp / total_pred if total_pred > 0 else 0.0
    fn = total_true - tp

    return ValidationMetrics(
        shd=shd,
        precision=precision,
        recall=recall,
        f1=f1,
        orientation_accuracy=orientation_acc,
        fdr=fdr,
        true_positives=tp,
        false_positives=fp,
        false_negatives=fn,
        reversed_edges=reversed_count,
        total_true_edges=total_true,
        total_predicted_edges=total_pred,
    )


def _undirected(edge: tuple[str, str]) -> tuple[str, str]:
    """Canonical undirected form: alphabetically ordered."""
    a, b = edge
    return (a, b) if a <= b else (b, a)
