import torch

from pi07mini.model import (
    HighLevelPolicy,
    Pi07Config,
    Pi07Mini,
    PrefixLayout,
    TinyGoalGenerator,
)
from pi07mini.tokenizer import EpisodeMetadata, MiniTokenizer


def tiny_model() -> Pi07Mini:
    tokenizer = MiniTokenizer()
    return Pi07Mini(
        Pi07Config(
            vocab_size=tokenizer.vocab_size,
            image_size=16,
            patch_size=16,
            d_model=32,
            n_heads=4,
            n_layers=1,
            memory_layers=1,
            max_history_frames=3,
            max_text_tokens=64,
            max_action_tokens=4,
            action_start_token_id=tokenizer.ACTION_START,
        )
    )


def sample_sequence():
    tokenizer = MiniTokenizer()
    metadata = EpisodeMetadata(2000, 5, False, "ee", "demonstration")
    return tokenizer, tokenizer.build_policy_sequence(
        "drive", "steer_left", torch.zeros(2, 2), metadata
    )


def common_inputs(sequence):
    return dict(
        pixel_values=torch.randn(1, 3, 16, 16),
        history_pixel_values=torch.randn(1, 2, 3, 16, 16),
        history_states=torch.randn(1, 2, 2),
        current_state=torch.randn(1, 2),
        subgoal_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],
    )


def test_block_causal_attention_matches_pi07_modal_order() -> None:
    layout = PrefixLayout(observation_tokens=2, goal_tokens=2, text_tokens=3)
    valid = torch.ones(1, layout.total_tokens, dtype=torch.bool)
    edges = Pi07Mini.build_prefix_attention_mask(layout, valid)[0]
    assert edges[0, :2].all()
    assert not edges[0, 2:].any()  # observation cannot read goal or text
    assert edges[2, :4].all()
    assert not edges[2, 4:].any()  # goal reads observation + goal, not text
    assert edges[4, :5].all()
    assert not edges[4, 5:].any()  # first text token is causal
    assert edges[6].all()


def test_prefix_has_fixed_mem_tokens_and_one_token_per_state() -> None:
    _, sequence = sample_sequence()
    model = tiny_model()
    inputs = common_inputs(sequence)
    prefix, _, layout = model.encode_prefix(
        inputs["pixel_values"],
        inputs["history_pixel_values"],
        inputs["history_states"],
        inputs["current_state"],
        inputs["input_ids"],
        inputs["subgoal_pixel_values"],
    )
    assert layout.observation_tokens == 1 + 2 + 1  # MEM + history states + current
    assert layout.goal_tokens == 1
    assert prefix.shape[1] == layout.total_tokens


def test_flow_cannot_read_teacher_forced_fast_tokens() -> None:
    torch.manual_seed(12)
    tokenizer, sequence = sample_sequence()
    model = tiny_model().eval()
    changed_ids = sequence.input_ids.clone()
    changed_ids[sequence.action_start_index + 1 : -1] = tokenizer.ACTION_TOKEN_START
    inputs = common_inputs(sequence)
    with torch.no_grad():
        first = model(**inputs).vector_field
        inputs["input_ids"] = changed_ids[None]
        second = model(**inputs).vector_field
    torch.testing.assert_close(first, second)


def test_flow_loss_is_knowledge_insulated_from_all_vlm_parameters() -> None:
    torch.manual_seed(13)
    _, sequence = sample_sequence()
    model = tiny_model()
    output = model(
        **common_inputs(sequence), target_velocity=torch.randn(1, 2, 2)
    )
    assert output.flow_loss is not None
    output.flow_loss.backward()
    assert all(
        parameter.grad is None or torch.count_nonzero(parameter.grad) == 0
        for parameter in model.vlm_parameters()
    )
    assert any(
        parameter.grad is not None and torch.count_nonzero(parameter.grad) > 0
        for parameter in model.action_parameters()
    )


def test_fast_cross_entropy_still_trains_vlm() -> None:
    _, sequence = sample_sequence()
    model = tiny_model()
    output = model(
        **common_inputs(sequence), text_labels=sequence.labels[None]
    )
    assert output.text_loss is not None
    output.text_loss.backward()
    assert model.token_embedding.weight.grad is not None
    assert torch.count_nonzero(model.token_embedding.weight.grad) > 0


def test_visual_subgoal_conditions_flow_expert() -> None:
    torch.manual_seed(14)
    _, sequence = sample_sequence()
    model = tiny_model().eval()
    inputs = common_inputs(sequence)
    with torch.no_grad():
        first = model(**inputs).vector_field
        inputs["subgoal_pixel_values"] = inputs["subgoal_pixel_values"] + 3.0
        second = model(**inputs).vector_field
    assert not torch.allclose(first, second)


def test_metadata_cfg_uses_paper_extrapolation_formula() -> None:
    conditioned = torch.tensor([1.0, 2.0])
    unconditional = torch.tensor([0.5, 1.5])
    torch.testing.assert_close(
        Pi07Mini.combine_metadata_cfg(conditioned, unconditional, 0.0), conditioned
    )
    torch.testing.assert_close(
        Pi07Mini.combine_metadata_cfg(conditioned, unconditional, 1.0),
        torch.tensor([1.5, 2.5]),
    )


def test_high_level_policy_predicts_next_subtask_token() -> None:
    model = HighLevelPolicy(tiny_model().config)
    output = model(
        pixel_values=torch.randn(2, 3, 16, 16),
        history_pixel_values=torch.randn(2, 2, 3, 16, 16),
        history_states=torch.randn(2, 2, 2),
        current_state=torch.randn(2, 2),
        input_ids=torch.tensor([[256, 257, 264], [256, 257, 264]]),
        target_token=torch.tensor([271, 272]),
    )
    assert output.logits.shape == (2, tiny_model().config.vocab_size)
    assert output.loss is not None


def test_tiny_goal_generator_predicts_image_flow() -> None:
    tokenizer = MiniTokenizer()
    generator = TinyGoalGenerator(tokenizer.vocab_size)
    current = torch.randn(2, 3, 16, 16)
    noisy = torch.randn_like(current)
    output = generator(
        current,
        torch.tensor([[256, 257, 264], [256, 257, 264]]),
        noisy,
        torch.rand(2),
        target_velocity=torch.randn_like(current),
    )
    assert output.vector_field.shape == current.shape
    assert output.loss is not None