"""Example 9 — Evaluate two training runs on a held-out test split; compare aggregate metrics.

This is the end-to-end ML shape and where the **dataset** earns its place: you convert a
workset into a stable, reusable dataset (the bridge in 08), split it train/val/test, and
score two candidate runs on the **same** held-out test set, then compare. Splits are cut
over a dataset — the reproducible source you train against.

    labeled store ─workset─▶ convert ─▶ dataset ─split_set─▶ train / val / TEST
                                                                  │  (per run)
                                          ┌───────────────────────┴───────────────────────┐
                                          ▼                                                ▼
                              eval store  (Add ← test split)                   eval store  (Add ← test split)
                              produce pred.cls  (inference)                    produce pred.cls
                              produce correct.cls (metric vs label)            produce correct.cls
                                          │                                                │
                                          ▼                                                ▼
                              mean(correct.cls)  = accuracy  ◀── compare ──▶  mean(correct.cls) = accuracy

Two ideas do the work:

* **A split set is worker-cut into child datasets** ``{split_set}/{train,val,test}``. Each
  is an ordinary, ready dataset; we seed each run's eval store from the **same frozen test
  split** with an ``Add`` (zero-copy, see 02_add_to_store.py), so both runs are scored on an
  identical held-out set.
* **Metrics are numeric scalar keys** (``.cls``), so the aggregate is a single server-side
  call: ``get_dataset_scalar_stats`` returns ``mean(correct.cls)`` — the run's accuracy —
  without streaming a byte to the client. Comparing runs is comparing two means.

Predictions and the metric are stand-ins (no real model runs); a run that "agrees with the
label" more often scores higher, which is all the comparison needs.

    DATAVO_SERVER=https://dev.datavo.io python examples/09_train_eval_compare.py
"""

from __future__ import annotations

import tempfile
from pathlib import Path

from _common import (
    build_dataset,
    cleanup,
    cleanup_datasets,
    cleanup_split_sets,
    make_client,
    seed_source_store,
    step,
    unique,
    wait_for_split_set,
)

from datavo_sdk import (
    DatasetScalarStatsRequest,
    SourceSpec,
    DatasetSplitSetCreateRequest,
    HasKeys,
    KeyGroupSpec,
    AttachedBy,
    Not,
    SampleStream,
    ShardPlan,
    SplitOutputSpec,
    SplitRatioSelector,
    TarWriter,
)

# A labeled store: raw audio + a ground-truth class label (``.cls`` = a numeric scalar key,
# so it is filterable/aggregatable server-side, unlike an opaque .npy blob).
N = 60
SAMPLES = [
    {"__key__": f"{i:04d}", "audio.wav": f"audio-{i}".encode(), "label.cls": i % 2}
    for i in range(N)
]

# Two simulated training runs. The number is the fraction of the test set the run predicts
# correctly: a "better" run agrees with the ground-truth label more often, so its aggregate
# accuracy comes out higher. (Deterministic stand-in — no real model runs.)
RUNS = [
    {"id": "run_a", "accuracy": 0.92},
    {"id": "run_b", "accuracy": 0.67},
]


def read_scalars(client, sample_ids: list[str], keys: list[str]) -> dict[str, dict[str, int]]:
    """Read scalar ``keys`` per sample id straight from the store — a server-side attribute
    read (``get_sample``), no streaming needed. A scalar key lands as a sample attribute
    with a ``numeric_value``, which is exactly what an eval step needs to score.
    """
    return {
        sid: {
            key: int(value)
            for a in client.get_sample(sid).attributes
            if (key := a.name) in keys and (value := a.numeric_value) is not None
        }
        for sid in sample_ids
    }


def attach_scalar(client, store: str, producer_id: str, key: str, *, inputs: list[str], values_for) -> None:
    """Attach one scalar ``key`` under ``producer_id`` for the samples that still need it
    (a condensed 03_attach_keys.py): select the not-yet-attached samples as a workset,
    stream them for their ids, compute, attach.

    Idempotency is per ``(producer_id, key)``, so ``pred.cls`` and ``correct.cls`` are two
    independent passes for the same run producer — and doneness is an **empty workset**.
    ``values_for(ids)`` computes the values for exactly the ids that still need them.
    """
    workset = client.create_workset(
        sources=[SourceSpec(
            sample_store=store,
            keys=[inputs[0]],
            filters=[
                HasKeys(inputs),
                Not(AttachedBy(producer_id, [key])),
            ],
        )],
        shard_plan=ShardPlan(shard_size=8),
    )
    if not workset.sample_count:
        return  # nothing missing — idempotent no-op
    ids = workset_sample_ids(client, workset, read_key=inputs[0])
    values = values_for(ids)
    with tempfile.TemporaryDirectory() as tmp:
        shard = Path(tmp) / f"{producer_id}-{key}.tar"
        with TarWriter(str(shard)) as writer:
            for sid in ids:
                writer.write({"__key__": sid, key: values[sid]})
        client.attach_keys(
            store, str(shard),
            keys=[key], producer_id=producer_id,
        )


def workset_sample_ids(client, workset, *, read_key: str) -> list[str]:
    """The workset's sample ids, by iterating it — the id set *is* the work, so there is no
    separate list to walk. ``decode=False``: a producer forwards payloads, it doesn't read them.
    """
    return [
        sample["__key__"]
        for sample in SampleStream.from_workset(
            workset.workset_id, keys=[read_key], client=client, decode=False
        )
    ]


def evaluate_run(client, test_split: str, run: dict, store: str, dataset: str) -> None:
    """Score one run on the shared test split: seed an eval store from it, attach this
    run's predictions + a per-sample correctness metric, and freeze a dataset pinned to
    the run's producer so its metric can be aggregated."""
    run_id, accuracy = run["id"], run["accuracy"]

    step(f"[{run_id}] seed eval store {store!r} from the test split (zero-copy Add)")
    client.create_store(name=store, display_name=f"eval store · {run_id}")
    client.add_to_store(
        store, operation_id=f"seed-{run_id}", from_dataset=test_split,
        key_subset=["audio.wav", "label.cls"], display_name="seed from test split",
    )
    client.wait_for_add(store, f"seed-{run_id}", timeout_seconds=300)

    # One producer identity per run, declaring BOTH outputs up front: declared_output_keys
    # are fixed at registration, so a run producer can't grow a new output later. pred.cls
    # is the inference output; correct.cls is the eval metric computed against the label.
    # `metadata` is free-form, caller-owned provenance (model version, run ref, …) — the
    # server stores it verbatim and never interprets it.
    client.register_producer(
        store, run_id, declared_output_keys=["pred.cls", "correct.cls"],
        metadata={"producer_kind": "eval_run", "accuracy": accuracy},
    )

    def predictions(ids: list[str]) -> dict[str, int]:
        """Deterministic stand-in for inference: the first ``accuracy``-fraction of samples
        predict the label correctly, the rest are flipped — so the run's measured accuracy
        lands on ``accuracy``."""
        labels = read_scalars(client, ids, ["label.cls"])
        n_correct = round(accuracy * len(ids))
        return {
            sid: (labels[sid]["label.cls"] if i < n_correct else 1 - labels[sid]["label.cls"])
            for i, sid in enumerate(sorted(ids))
        }

    def correctness(ids: list[str]) -> dict[str, int]:
        """The metric reads its inputs back from the store, like any downstream producer:
        1 where this run's pred.cls agrees with the ground-truth label, else 0."""
        scored = read_scalars(client, ids, ["pred.cls", "label.cls"])
        return {sid: int(v["pred.cls"] == v["label.cls"]) for sid, v in scored.items()}

    step(f"[{run_id}] attach pred.cls (stand-in inference, ~{accuracy:.0%} accurate)")
    attach_scalar(client, store, run_id, "pred.cls", inputs=["audio.wav"], values_for=predictions)

    step(f"[{run_id}] attach correct.cls (metric: 1 when pred == label, else 0)")
    attach_scalar(
        client, store, run_id, "correct.cls",
        inputs=["audio.wav", "pred.cls"], values_for=correctness,
    )

    step(f"[{run_id}] freeze an eval dataset pinning this run's produced keys")
    build_dataset(
        client,
        dataset,
        sources=[
            SourceSpec(
                sample_store=store,
                keys=["pred.cls", "correct.cls", "label.cls"],
                # base label is unpinned; the derived keys pin THIS run's producer:
                producer_pins={"pred.cls": run_id, "correct.cls": run_id},
            )
        ],
        key_groups=[
            KeyGroupSpec(name="pred", keys=["pred.cls"]),
            KeyGroupSpec(name="metric", keys=["correct.cls"]),
            KeyGroupSpec(name="label", keys=["label.cls"]),
        ],
        shard_plan=ShardPlan(shard_size=8),
    )


def aggregate_accuracy(client, dataset: str) -> tuple[float, int]:
    """The run's accuracy = mean(correct.cls) over the eval dataset, computed server-side."""
    stats = client.get_dataset_scalar_stats(dataset, DatasetScalarStatsRequest(stats_keys=["correct.cls"]))
    stat = next(s for s in stats.scalar_stats if s.key == "correct.cls")
    return float(stat.mean_value or 0.0), stat.present_count


def main() -> None:
    client = make_client()
    store = unique("example_eval_store")
    dataset = unique("example_eval_ds")
    split_set = unique("example_eval_split")
    eval_stores = {run["id"]: unique(f"example_eval_{run['id']}") for run in RUNS}
    eval_datasets = {run["id"]: unique(f"example_eval_{run['id']}_ds") for run in RUNS}
    try:
        step(f"seed a labeled store {store!r} with {N} samples (audio.wav + label.cls)")
        seed_source_store(client, store, samples=SAMPLES)

        step(f"freeze it into dataset {dataset!r}, then cut a 60/20/20 train/val/test split")
        # You *could* point training straight at a workset — it's iterable (see 05–07) and 10
        # does exactly that for its follow-up stages. But a run you'll reproduce, split, and
        # compare wants a stable, reusable source, so here we convert to a dataset (08) and
        # split that (split sets are cut over a dataset).
        build_dataset(
            client,
            dataset,
            sources=[
                SourceSpec(
                    sample_store=store,
                    keys=["audio.wav", "label.cls"],
                )
            ],
            key_groups=[
                KeyGroupSpec(name="audio", keys=["audio.wav"]),
                KeyGroupSpec(name="label", keys=["label.cls"]),
            ],
            shard_plan=ShardPlan(shard_size=8),
        )
        client.create_dataset_split_set(
            dataset,
            DatasetSplitSetCreateRequest(
                name=split_set,
                assignment_seed=7,
                splits=[
                    SplitOutputSpec(name="train", selector=SplitRatioSelector(kind="ratio", ratio=0.6)),
                    SplitOutputSpec(name="val", selector=SplitRatioSelector(kind="ratio", ratio=0.2)),
                    SplitOutputSpec(name="test", selector=SplitRatioSelector(kind="ratio", ratio=0.2)),
                ],
                shard_plan=ShardPlan(shard_size=8),
            ),
        )
        detail = wait_for_split_set(client, dataset, split_set)
        sizes = {d.split_name: d.sample_count for d in detail.datasets}
        print(f"    split sizes: {sizes} (children are datasets named {split_set}/<split>)")
        test_split = f"{split_set}/test"  # a normal dataset — Add-able and open-able

        step("evaluate each run on the SAME held-out test split")
        for run in RUNS:
            evaluate_run(client, test_split, run, eval_stores[run["id"]], eval_datasets[run["id"]])

        step("compare aggregate metrics between the two runs (server-side mean of correct.cls)")
        results = {run["id"]: aggregate_accuracy(client, eval_datasets[run["id"]]) for run in RUNS}
        print(f"    {'run':<8} {'accuracy':>9}   {'n':>3}")
        for run_id, (acc, n) in results.items():
            print(f"    {run_id:<8} {acc:>8.1%}   {n:>3}")
        winner = max(results, key=lambda run_id: results[run_id][0])
        print(f"    → {winner} wins on the held-out test split "
              f"({results[winner][0]:.1%} vs {min(a for a, _ in results.values()):.1%})")
    finally:
        # Teardown order: eval datasets → eval stores → split set → base dataset → base store.
        cleanup_datasets(client, *eval_datasets.values())
        cleanup(client, *eval_stores.values())
        cleanup_split_sets(client, dataset, split_set)
        cleanup_datasets(client, dataset)
        cleanup(client, store)


if __name__ == "__main__":
    main()