import numpy as np

from starVLA.dataloader.infoflow_sampler import ValidAnchorManifest
from starVLA.dataloader.temporal_proposal import TemporalProposal
from starVLA.scoring.matched_random import build_matched_random_proposal


def _manifest() -> ValidAnchorManifest:
    return ValidAnchorManifest.from_hierarchy(
        dataset_names=["bridge", "rt1"],
        dataset_weights=[1.0, 1.0],
        episodes=[
            [(10, [0, 1, 2, 3], 1.0), (11, [4, 5, 6], 1.0)],
            [(20, [7, 8, 9, 10, 11], 1.0)],
        ],
    )


def _residual(manifest: ValidAnchorManifest) -> TemporalProposal:
    return TemporalProposal.from_masses(
        manifest=manifest,
        alpha=0.2,
        masses=[
            [[0.1, 0.2, 0.4, 0.8], [0.3, 0.7, 0.9]],
            [[0.2, 0.25, 0.5, 0.75, 1.0]],
        ],
        score_kind="residual",
        scorer_sha256="residual-scorer",
    )


def test_matched_random_preserves_every_episode_multiset() -> None:
    manifest = _manifest()
    residual = _residual(manifest)
    matched = build_matched_random_proposal(
        manifest=manifest,
        residual_proposal=residual,
        permutation_seed=20260728,
    )

    assert matched.score_kind == "random"
    assert matched.alpha == residual.alpha
    assert matched.manifest_sha256 == residual.manifest_sha256
    for dataset_index in range(len(manifest.datasets)):
        for episode_index in range(
            len(manifest.datasets[dataset_index].episodes)
        ):
            np.testing.assert_array_equal(
                np.sort(
                    residual.proposal_masses[dataset_index][episode_index]
                ),
                np.sort(matched.proposal_masses[dataset_index][episode_index]),
            )
            np.testing.assert_array_equal(
                np.sort(
                    residual.anchor_probabilities[dataset_index][episode_index]
                ),
                np.sort(
                    matched.anchor_probabilities[dataset_index][episode_index]
                ),
            )
            np.testing.assert_array_equal(
                np.sort(
                    residual.importance_weights[dataset_index][episode_index]
                ),
                np.sort(
                    matched.importance_weights[dataset_index][episode_index]
                ),
            )


def test_matched_random_is_deterministic_and_seeded() -> None:
    manifest = _manifest()
    residual = _residual(manifest)
    first = build_matched_random_proposal(
        manifest=manifest,
        residual_proposal=residual,
        permutation_seed=1,
    )
    repeated = build_matched_random_proposal(
        manifest=manifest,
        residual_proposal=residual,
        permutation_seed=1,
    )
    other = build_matched_random_proposal(
        manifest=manifest,
        residual_proposal=residual,
        permutation_seed=2,
    )

    assert first.sha256 == repeated.sha256
    assert first.sha256 != other.sha256