"""Example 15 — Reuse a split's membership: other keys, another store, everything else.

A dataset freezes *which samples*; ``MemberOf`` lets you reuse that membership anywhere the
samples went. It is one more row predicate on a store source::

    SourceSpec(sample_store=store, keys=["transcript.txt"],
               filters=[MemberOf(dataset, split_set=split_set, split_name="holdout")])

A sample matches when it **is** a member or **descends from** one through an Add — matching
is on ``origin_key``, the lineage every Add propagates. So the same predicate reads:

1. the holdout split with a key the dataset did not freeze, in the dataset's own store;
2. the holdout split in a store its samples were added into;
3. and, negated, everything in the store that is *not* in the holdout.

Last, the splits themselves combine: ``sources_of`` turns each split into a source, so a
union of splits is one workset. A union is a set of samples: parts that overlap hold each
shared sample once, and it freezes into a dataset exactly as it reads.

    DATAVO_SERVER=https://dev.datavo.io python examples/15_member_of.py
"""

from __future__ import annotations

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 (
    DatasetSplitSetCreateRequest,
    KeyGroupSpec,
    MemberOf,
    Not,
    ShardPlan,
    SourceSpec,
    SplitOutputSpec,
    SplitRatioSelector,
    sources_of,
)

SAMPLES = [
    {"__key__": f"u{i:03d}", "audio.wav": f"wav-{i}".encode(), "transcript.txt": f"utterance {i}"}
    for i in range(8)
]


def main() -> None:
    client = make_client()
    store, copy_store = unique("example_member_store"), unique("example_member_copy")
    dataset, split_set = unique("example_member_ds"), unique("example_member_split")
    frozen_union = unique("example_member_union")
    try:
        step(f"seed {store!r}, freeze {dataset!r} over audio.wav only, split it fit/holdout")
        seed_source_store(client, store, samples=SAMPLES)
        build_dataset(
            client,
            dataset,
            sources=[SourceSpec(sample_store=store, keys=["audio.wav"])],
            key_groups=[KeyGroupSpec(name="audio", keys=["audio.wav"])],
            shard_plan=ShardPlan(shard_size=4),
        )
        client.create_dataset_split_set(
            dataset,
            DatasetSplitSetCreateRequest(
                name=split_set,
                assignment_seed=7,
                splits=[
                    SplitOutputSpec(name="fit", selector=SplitRatioSelector(ratio=0.5)),
                    SplitOutputSpec(name="holdout", selector=SplitRatioSelector(ratio=0.5)),
                ],
                shard_plan=ShardPlan(shard_size=4),
            ),
        )
        wait_for_split_set(client, dataset, split_set)
        holdout = client.get_dataset(f"{split_set}/holdout").sample_count
        in_holdout = MemberOf(dataset, split_set=split_set, split_name="holdout")
        print(f"    holdout holds {holdout} of {len(SAMPLES)} samples")

        step("1. the holdout, with a key the dataset never froze (transcript.txt)")
        same_store = client.create_workset(
            sources=[SourceSpec(sample_store=store, keys=["transcript.txt"], filters=[in_holdout])],
            shard_plan=ShardPlan(shard_size=4),
        )
        assert client.get_workset(same_store.workset_id).sample_count == holdout
        print(f"    {holdout} samples, each streaming transcript.txt")

        step(f"2. the holdout in {copy_store!r}, a store its samples were added into")
        client.create_store(name=copy_store)
        client.add_to_store(copy_store, operation_id="copy-all", from_store=store)
        client.wait_for_add(copy_store, "copy-all", timeout_seconds=180)
        copied = client.create_workset(
            sources=[SourceSpec(sample_store=copy_store, keys=["audio.wav"], filters=[in_holdout])],
            shard_plan=ShardPlan(shard_size=4),
        )
        read = client.get_workset(copied.workset_id)
        assert read.sample_count == holdout
        # A store holds one sample per origin, so a full match is matched == member_origins.
        match = read.membership_matches[0]
        print(f"    matched {match.matched} of {match.member_origins} member origins — the copies, under new ids")

        step("3. everything in the store that is NOT in the holdout")
        rest = client.create_workset(
            sources=[SourceSpec(sample_store=store, keys=["audio.wav"], filters=[Not(in_holdout)])],
            shard_plan=ShardPlan(shard_size=4),
        )
        assert client.get_workset(rest.workset_id).sample_count == len(SAMPLES) - holdout
        print(f"    {len(SAMPLES) - holdout} samples — the store minus the holdout")

        step("4. union the splits back together with sources_of")
        fit_split = client.get_dataset(f"{split_set}/fit")
        holdout_split = client.get_dataset(f"{split_set}/holdout")
        both = client.create_workset(
            sources=[*sources_of(fit_split), *sources_of(holdout_split)],
            shard_plan=ShardPlan(shard_size=4),
        )
        assert client.get_workset(both.workset_id).sample_count == len(SAMPLES)
        print(f"    fit + holdout = {len(SAMPLES)} samples, the whole dataset again")

        # The dataset and its own holdout overlap. A union is a set, so the holdout samples
        # are held once — and freezing it holds exactly what reading it streams.
        overlapping = client.create_workset(
            sources=[*sources_of(client.get_dataset(dataset)), *sources_of(holdout_split)],
            shard_plan=ShardPlan(shard_size=4),
        )
        assert client.get_workset(overlapping.workset_id).sample_count == len(SAMPLES)
        client.convert_workset_to_dataset(overlapping.workset_id, frozen_union, shard_plan=ShardPlan(shard_size=4))
        assert client.wait_for_dataset(frozen_union, timeout_seconds=600).sample_count == len(SAMPLES)
        print(f"    dataset + its own holdout = {len(SAMPLES)} samples, live and frozen alike")
    finally:
        cleanup_split_sets(client, dataset, split_set)
        cleanup_datasets(client, frozen_union, dataset)
        cleanup(client, copy_store, store)


if __name__ == "__main__":
    main()