"""Deliberately synthetic examples of leakage and prevalence arithmetic."""
from dataclasses import dataclass
import math


@dataclass(frozen=True)
class Visit:
    patient: int
    visit: int
    label: int


def fixture():
    return [Visit(patient, visit, patient % 2)
            for patient in range(40) for visit in range(2)]


def assert_disjoint_patients(train, test):
    overlap = {row.patient for row in train} & {row.patient for row in test}
    if overlap:
        raise ValueError(f"patient overlap: {len(overlap)}")


def memorizer_accuracy(train, test):
    if not train or not test:
        raise ValueError("both splits need rows")
    labels = {row.patient: row.label for row in train}
    # A lookup baseline, not a clinical model: unseen patients get class 0.
    return sum(labels.get(row.patient, 0) == row.label for row in test) / len(test)


def expected_counts(population, prevalence, sensitivity, specificity):
    if type(population) is not int or population <= 0:
        raise ValueError("population must be a positive integer")
    for value in (prevalence, sensitivity, specificity):
        if isinstance(value, bool) or not isinstance(value, (int, float)):
            raise ValueError("rates must be numeric")
        if not math.isfinite(value) or not 0 <= value <= 1:
            raise ValueError("rates must be finite and between 0 and 1")
    positives = population * prevalence
    tp = positives * sensitivity
    fp = (population - positives) * (1 - specificity)
    return {"true_positive": tp, "false_positive": fp,
            "false_negative": positives - tp,
            "true_negative": population - positives - fp,
            "precision": tp / (tp + fp) if tp + fp else None}


if __name__ == "__main__":
    rows = fixture()
    row_train = [r for r in rows if r.visit == 0]
    row_test = [r for r in rows if r.visit == 1]
    group_train = [r for r in rows if r.patient < 20]
    group_test = [r for r in rows if r.patient >= 20]
    assert_disjoint_patients(group_train, group_test)
    print(f"row split accuracy: {memorizer_accuracy(row_train, row_test):.0%}")
    print(f"patient split accuracy: {memorizer_accuracy(group_train, group_test):.0%}")
    for prevalence in (0.01, 0.1):
        counts = expected_counts(10000, prevalence, 0.9, 0.95)
        print(f"prevalence={prevalence:.0%} TP={counts['true_positive']:.0f} "
              f"FP={counts['false_positive']:.0f} precision={counts['precision']:.1%}")