"""Audit a hand-built group/time split; no sklearn or model training.

Run: python group-time-split.py
All records are synthetic. Days are integer instants within one toy month.
The fitting snapshot is day 6, and the offline scoring snapshot is day 15.
"""

from dataclasses import dataclass, replace
from statistics import mean
import json


@dataclass(frozen=True)
class Row:
    row_id: str
    user: str
    day: int
    label_known_day: int
    x: float
    feature_known_day: int


ROWS = (
    Row("a1", "A", 1, 3, 10, 1),
    Row("b1", "B", 2, 4, 20, 2),
    Row("a2", "A", 4, 8, 30, 4),
    Row("c1", "C", 6, 8, 100, 6),
    Row("b2", "B", 7, 9, 40, 7),
    Row("d1", "D", 8, 10, 120, 8),
    Row("c2", "C", 9, 11, 140, 10),
    Row("a3", "A", 10, 12, 50, 10),
)
FIT_DAY = 6
SCORE_DAY = 15


def audit(train, test, *, require_new_users=True):
    """Check this exercise's declared evaluation contract, not every leakage."""
    if not train or not test:
        raise ValueError("both sides must contain rows")
    train_ids = {r.row_id for r in train}
    test_ids = {r.row_id for r in test}
    overlap = sorted({r.user for r in train} & {r.user for r in test})
    features_late = sorted(r.row_id for r in (*train, *test)
                           if r.feature_known_day > r.day)
    labels_late = sorted(r.row_id for r in train if r.label_known_day > FIT_DAY)
    return {
        "row_ids_disjoint": not bool(train_ids & test_ids),
        "group_overlap": overlap,
        "group_rule_passed": not overlap or not require_new_users,
        "train_before_test": max(r.day for r in train) < min(r.day for r in test),
        "fit_snapshot_passed": all(r.day < FIT_DAY for r in train)
                              and all(r.day >= FIT_DAY for r in test),
        "training_labels_unavailable": labels_late,
        "features_unavailable_at_prediction": features_late,
        "test_labels_ready_for_scoring": all(r.label_known_day <= SCORE_DAY for r in test),
    }


def passed(result):
    return all((result["row_ids_disjoint"], result["group_rule_passed"],
                result["train_before_test"], result["fit_snapshot_passed"],
                not result["training_labels_unavailable"],
                not result["features_unavailable_at_prediction"],
                result["test_labels_ready_for_scoring"]))


def main():
    # Group isolation alone allows training rows later than held-out rows.
    group_train = tuple(r for r in ROWS if r.user in {"A", "B"})
    group_test = tuple(r for r in ROWS if r.user in {"C", "D"})
    group_only = audit(group_train, group_test)
    assert group_only["group_rule_passed"] and not group_only["train_before_test"]

    # Time order alone allows the same users on both sides and pending labels.
    time_train = tuple(r for r in ROWS if r.day < FIT_DAY)
    time_test = tuple(r for r in ROWS if r.day >= FIT_DAY)
    time_only = audit(time_train, time_test)
    assert time_only["train_before_test"] and time_only["group_overlap"] == ["A", "B"]
    assert time_only["training_labels_unavailable"] == ["a2"]
    assert time_only["features_unavailable_at_prediction"] == ["c2"]

    # Freeze a training snapshot, then evaluate its unseen users in the future.
    train = tuple(r for r in ROWS if r.day < FIT_DAY
                  and r.label_known_day <= FIT_DAY and r.feature_known_day <= r.day)
    seen = {r.user for r in train}
    test = tuple(r for r in ROWS if r.day >= FIT_DAY and r.user not in seen
                 and r.feature_known_day <= r.day and r.label_known_day <= SCORE_DAY)
    joint = audit(train, test)
    assert [r.row_id for r in train] == ["a1", "b1"]
    assert [r.row_id for r in test] == ["c1", "d1"]
    assert passed(joint)

    # Fitting a centering mean from test values changes the learned transform.
    train_mean = mean(r.x for r in train)
    leaky_mean = mean(r.x for r in (*train, *test))
    centered_test = [r.x - train_mean for r in test]
    leaky_centered_test = [r.x - leaky_mean for r in test]
    assert train_mean == 15 and leaky_mean == 62.5
    assert centered_test == [85, 105] and leaky_centered_test == [37.5, 57.5]
    changed_test = tuple(replace(r, x=r.x + 1000) for r in test)
    assert mean(r.x for r in train) == train_mean
    assert mean(r.x for r in (*train, *changed_test)) != leaky_mean

    # Demonstrate the final-fold row rule only, not sklearn's implementation.
    # Six sorted rows: train indices 0..2; gap indices 3..4; test index 5.
    regular_days = [1, 2, 3, 4, 5, 6]
    irregular_days = [1, 2, 20, 21, 22, 30]
    gap_rows = 2
    test_index = 5
    last_train_index = test_index - gap_rows - 1
    calendar_gaps = [days[test_index] - days[last_train_index]
                     for days in (regular_days, irregular_days)]
    assert calendar_gaps == [3, 10]

    # A negative control: an empty split must not silently pass the audit.
    try:
        audit((), test)
    except ValueError:
        empty_split_rejected = True
    else:
        raise AssertionError("empty split was accepted")

    report = {
        "group_only": group_only,
        "time_only": time_only,
        "joint": {"train_ids": [r.row_id for r in train],
                  "test_ids": [r.row_id for r in test], "audit": joint},
        "centering": {"train_mean": train_mean, "leaky_mean": leaky_mean,
                      "test_using_train_mean": centered_test,
                      "test_using_leaky_mean": leaky_centered_test},
        "same_gap_rows": gap_rows,
        "calendar_distance_days": calendar_gaps,
        "empty_split_rejected": empty_split_rejected,
        "result": "PASS: split, availability, centering and row-gap invariants",
        "scope": "Python standard-library arithmetic and contract checks; no sklearn or model training",
    }
    print(json.dumps(report, ensure_ascii=False, indent=2))


if __name__ == "__main__":
    if not __debug__:
        raise SystemExit("Run without -O or -OO: this lab requires active assertions.")
    main()