from collections import Counter
import json

import pytest

from starVLA.dataloader.infoflow_sampler import LogicalDraw, ValidAnchorManifest
from starVLA.dataloader.temporal_proposal import TemporalProposal


def _manifest() -> ValidAnchorManifest:
    return ValidAnchorManifest.from_hierarchy(
        ["dataset"],
        [1.0],
        [[(5, [0, 1, 2], 1.0)]],
        identity={"fixture": "proposal"},
    )


def test_alpha_zero_preserves_exact_anchor_path() -> None:
    manifest = _manifest()
    proposal = TemporalProposal.from_masses(manifest, 0.0, [[[100.0, 1.0, 1.0]]])

    for draw_id in range(100):
        draw = LogicalDraw(2, draw_id)
        baseline = manifest.sample(11, draw)
        proposed = proposal.sample(manifest, 11, draw)
        assert (
            proposed.dataset_index,
            proposed.trajectory_id,
            proposed.base_index,
        ) == (
            baseline.dataset_index,
            baseline.trajectory_id,
            baseline.base_index,
        )
        assert proposed.importance_weight == pytest.approx(1.0)


def test_safe_temporal_proposal_probability_and_weight_bound() -> None:
    manifest = _manifest()
    proposal = TemporalProposal.from_masses(manifest, 0.5, [[[8.0, 1.0, 1.0]]])

    probabilities = proposal.anchor_probabilities[0][0]
    weights = proposal.importance_weights[0][0]

    assert probabilities.sum() == pytest.approx(1.0)
    assert max(weights) <= 2.0 + 1e-12
    assert all(weight > 0 for weight in weights)


def test_empirical_anchor_frequency_matches_temporal_proposal() -> None:
    manifest = _manifest()
    proposal = TemporalProposal.from_masses(manifest, 0.6, [[[8.0, 1.0, 1.0]]])
    counts = Counter(
        proposal.sample(manifest, 3, LogicalDraw(0, draw_id)).base_index
        for draw_id in range(30_000)
    )

    for anchor, probability in enumerate(proposal.anchor_probabilities[0][0]):
        assert counts[anchor] / 30_000 == pytest.approx(probability, abs=0.01)


def test_likelihood_label_requires_validation_certificate() -> None:
    with pytest.raises(ValueError, match="validation certificate"):
        TemporalProposal.from_masses(
            _manifest(),
            0.2,
            [[[1.0, 1.0, 1.0]]],
            score_kind="likelihood",
            scorer_sha256="model-hash",
        )


def test_proposal_serialization_round_trip_and_hash_validation() -> None:
    manifest = _manifest()
    proposal = TemporalProposal.from_masses(manifest, 0.4, [[[3.0, 2.0, 1.0]]])
    payload = json.loads(json.dumps(proposal.to_dict()))

    loaded = TemporalProposal.from_dict(manifest, payload)

    assert loaded.sha256 == proposal.sha256
    assert loaded.to_dict() == proposal.to_dict()

    payload["proposal_masses"][0][0][0] = 999.0
    with pytest.raises(ValueError, match="hash mismatch"):
        TemporalProposal.from_dict(manifest, payload)


def test_proposal_rejects_schema_and_manifest_mismatch_but_accepts_legacy_metadata() -> None:
    manifest = _manifest()
    proposal = TemporalProposal.from_masses(manifest, 0.2, [[[1.0, 2.0, 3.0]]])

    bad_schema = proposal.to_dict()
    bad_schema["schema_version"] = 999
    with pytest.raises(ValueError, match="unsupported temporal proposal schema"):
        TemporalProposal.from_dict(manifest, bad_schema)

    other_manifest = ValidAnchorManifest.from_hierarchy(
        ["dataset"],
        [1.0],
        [[(5, [0, 1, 2], 1.0)]],
        identity={"fixture": "different"},
    )
    with pytest.raises(ValueError, match="different valid-anchor manifest"):
        TemporalProposal.from_dict(other_manifest, proposal.to_dict())

    legacy = proposal.to_dict()
    legacy.pop("sha256")
    legacy["implementation_sha256"] = "0" * 64
    loaded = TemporalProposal.from_dict(manifest, legacy)
    assert loaded.legacy_implementation_sha256 == "0" * 64