import torch

from pi06mini.model import Pi06Config, Pi06Mini
from pi06mini.tokenizer import MiniTokenizer


def tiny_model() -> Pi06Mini:
    tokenizer = MiniTokenizer()
    return Pi06Mini(
        Pi06Config(
            vocab_size=tokenizer.vocab_size,
            image_size=16,
            patch_size=16,
            d_model=32,
            n_heads=4,
            n_layers=1,
            max_text_tokens=32,
            max_action_tokens=4,
            action_start_token_id=tokenizer.ACTION_START,
        )
    )


def test_flow_reads_subtask_and_advantage_but_not_fast_targets() -> None:
    tokenizer = MiniTokenizer()
    sequence = tokenizer.build_policy_sequence(
        torch.zeros(2, 2), "steer_left", torch.zeros(2, 2), "positive"
    )
    model = tiny_model()
    input_ids = sequence.input_ids[None]
    derived = model._derive_action_prefix_mask(input_ids)
    torch.testing.assert_close(derived[0], sequence.action_prefix_mask)

    text_offset = 1
    prefix, flow = model._masks(
        batch=1,
        prefix_length=1 + input_ids.shape[1],
        action_length=2,
        text_offset=text_offset,
        attention_mask=torch.ones_like(input_ids, dtype=torch.bool),
        action_prefix_mask=derived,
        generated_token_start=sequence.generated_token_start,
        device=torch.device("cpu"),
    )
    generated_start = text_offset + sequence.generated_token_start
    assert prefix[0, 0, :generated_start].all()
    assert not prefix[0, 0, generated_start:].any()
    # Flow sees image plus all text strictly before ACTION_START.
    readable_prefix = text_offset + sequence.action_start_index
    assert flow[0, :, :readable_prefix].all()
    assert not flow[0, :, readable_prefix : 1 + input_ids.shape[1]].any()
    assert flow[0, :, 1 + input_ids.shape[1] :].all()


def test_flow_loss_is_knowledge_insulated_from_every_vlm_parameter() -> None:
    torch.manual_seed(4)
    tokenizer = MiniTokenizer()
    sequence = tokenizer.build_policy_sequence(
        torch.zeros(2, 2), "steer_left", torch.zeros(2, 2), "positive"
    )
    model = tiny_model()
    output = model(
        pixel_values=torch.randn(1, 3, 16, 16),
        input_ids=sequence.input_ids[None],
        noisy_actions=torch.randn(1, 2, 2),
        timesteps=torch.rand(1),
        action_prefix_mask=sequence.action_prefix_mask[None],
        generated_token_start=sequence.generated_token_start,
        target_velocity=torch.randn(1, 2, 2),
    )
    assert output.flow_loss is not None
    output.flow_loss.backward()

    vlm_gradients = [parameter.grad for parameter in model.vlm_parameters()]
    assert vlm_gradients
    assert all(gradient is None or torch.count_nonzero(gradient) == 0 for gradient in vlm_gradients)
    action_gradients = [parameter.grad for parameter in model.action_parameters()]
    assert any(gradient is not None and torch.count_nonzero(gradient) > 0 for gradient in action_gradients)


def test_knowledge_insulation_stops_gradients_not_conditioning() -> None:
    torch.manual_seed(8)
    tokenizer = MiniTokenizer()
    model = tiny_model().eval()
    positive = tokenizer.build_policy_sequence(
        torch.zeros(2, 2), "keep_lane", torch.zeros(2, 2), "positive"
    )
    negative = tokenizer.build_policy_sequence(
        torch.zeros(2, 2), "keep_lane", torch.zeros(2, 2), "negative"
    )
    image = torch.randn(1, 3, 16, 16)
    noise = torch.randn(1, 2, 2)
    timestep = torch.full((1,), 0.5)
    with torch.no_grad():
        pos = model(
            image, positive.input_ids[None], noise, timestep,
            generated_token_start=positive.generated_token_start,
        ).vector_field
        neg = model(
            image, negative.input_ids[None], noise, timestep,
            generated_token_start=negative.generated_token_start,
        ).vector_field
    assert not torch.allclose(pos, neg)


def test_joint_loss_still_trains_vlm_through_token_objective() -> None:
    tokenizer = MiniTokenizer()
    sequence = tokenizer.build_policy_sequence(
        torch.zeros(2, 2), "stop", torch.zeros(2, 2), "negative"
    )
    model = tiny_model()
    output = model(
        torch.randn(1, 3, 16, 16),
        sequence.input_ids[None],
        torch.randn(1, 2, 2),
        torch.rand(1),
        generated_token_start=sequence.generated_token_start,
        text_labels=sequence.labels[None],
        target_velocity=torch.randn(1, 2, 2),
    )
    output.loss.backward()
    assert model.token_embedding.weight.grad is not None
    assert torch.count_nonzero(model.token_embedding.weight.grad) > 0