import torch

from pi06mini.model import Pi06Config, Pi06Mini
from pi06mini.tokenizer import MiniTokenizer
from pi06mini.value import DistributionalValueModel, ValueConfig


def tiny_value_model() -> DistributionalValueModel:
    tokenizer = MiniTokenizer()
    return DistributionalValueModel(
        ValueConfig(
            vocab_size=tokenizer.vocab_size,
            image_size=16,
            patch_size=16,
            d_model=32,
            n_heads=4,
            n_layers=1,
            max_text_tokens=16,
        )
    )


def test_distributional_value_has_201_bins_and_expected_value() -> None:
    tokenizer = MiniTokenizer()
    model = tiny_value_model()
    state = tokenizer.build_value_input(torch.zeros(2, 2))[None].repeat(2, 1)
    output = model(
        torch.randn(2, 3, 16, 16),
        state,
        target_values=torch.tensor([-1.0, 0.0]),
    )

    assert output.logits.shape == (2, 201)
    assert output.probabilities.shape == (2, 201)
    torch.testing.assert_close(output.probabilities.sum(-1), torch.ones(2))
    assert torch.all((-1.0 <= output.expectation) & (output.expectation <= 0.0))
    assert output.target_bins.tolist() == [0, 200]
    assert output.loss is not None and output.loss.ndim == 0


def test_value_discretization_and_expectation_use_configured_support() -> None:
    model = tiny_value_model()
    bins = model.values_to_bins(torch.tensor([-2.0, -0.5, 1.0]))
    assert bins.tolist() == [0, 100, 200]
    torch.testing.assert_close(
        model.bins_to_values(torch.tensor([0, 100, 200])),
        torch.tensor([-1.0, -0.5, 0.0]),
    )
    logits = torch.full((3, 201), -100.0)
    logits[0, 0] = logits[1, 100] = logits[2, 200] = 100.0
    torch.testing.assert_close(
        model.expectation_from_logits(logits),
        torch.tensor([-1.0, -0.5, 0.0]),
    )


def test_value_model_and_policy_share_no_parameters() -> None:
    tokenizer = MiniTokenizer()
    policy = Pi06Mini(
        Pi06Config(
            vocab_size=tokenizer.vocab_size,
            image_size=16,
            patch_size=16,
            d_model=32,
            n_heads=4,
            n_layers=1,
            action_start_token_id=tokenizer.ACTION_START,
        )
    )
    value = tiny_value_model()
    policy_ids = {id(parameter) for parameter in policy.parameters()}
    value_ids = {id(parameter) for parameter in value.parameters()}
    assert policy_ids.isdisjoint(value_ids)