"""What the model learns from: the rows, their splits and how they are made."""

from dataclasses import dataclass


@dataclass(frozen=True)
class Sample:
    """One labelled example; `split` is train, validation or test."""

    x: float
    y: float
    split: str


def samples() -> list[Sample]:
    """Tiny synthetic rows for y = 2x. Replace with your own source."""
    points = [
        (-1.0, "train"),
        (-0.5, "train"),
        (0.5, "train"),
        (1.0, "train"),
        (0.25, "validation"),
        (0.75, "test"),
    ]
    return [Sample(x, 2.0 * x, split) for x, split in points]


def part(rows: list[Sample], split: str) -> list[Sample]:
    """The rows of one split."""
    return [row for row in rows if row.split == split]
