from __future__ import annotations

import torch

from pi07mini.context import (
    PolicyMetadata,
    build_ambiguous_history_pair,
    drop_prompt_components,
    format_policy_prompt,
    future_observation_proxy,
    make_history_proxy_frames,
    prompt_components,
)


def _history() -> torch.Tensor:
    x = torch.linspace(-8.0, 0.0, 8)
    return torch.stack((x, 0.04 * x.square()), -1)


def test_future_subgoal_has_image_contract_not_coordinate_heatmap() -> None:
    image = torch.linspace(0.0, 1.0, 3 * 24 * 24).reshape(3, 24, 24)
    future = torch.stack((torch.linspace(0.0, 20.0, 16), torch.linspace(0.0, 4.0, 16)), -1)
    subgoal = future_observation_proxy(image, _history(), future)
    assert subgoal.shape == image.shape
    assert torch.isfinite(subgoal).all()
    assert not torch.equal(subgoal, image)


def test_history_proxy_is_oldest_to_newest_and_future_free() -> None:
    image = torch.linspace(0.0, 1.0, 3 * 20 * 20).reshape(3, 20, 20)
    frames = make_history_proxy_frames(image, _history(), frames=6)
    assert frames.shape == (6, 3, 20, 20)
    torch.testing.assert_close(frames[-1], image * 0.98)
    assert not torch.equal(frames[0], frames[-1])


def test_ambiguous_histories_are_distinct_but_share_current_state() -> None:
    pair = build_ambiguous_history_pair(_history(), frames=6)
    assert pair.shape == (2, 6, 2)
    torch.testing.assert_close(pair[0, -1], pair[1, -1])
    assert not torch.allclose(pair[0, :-1], pair[1, :-1])


def test_prompt_component_dropout_is_independent_and_deterministic() -> None:
    metadata = PolicyMetadata("fast", "good", "none", "autonomous")
    components = prompt_components(metadata, task="Drive safely.")
    generator_a = torch.Generator().manual_seed(42)
    generator_b = torch.Generator().manual_seed(42)
    visible_a, mask_a = drop_prompt_components(
        components, drop_probability=0.5, generator=generator_a
    )
    visible_b, mask_b = drop_prompt_components(
        components, drop_probability=0.5, generator=generator_b
    )
    assert visible_a == visible_b
    torch.testing.assert_close(mask_a, mask_b)
    assert visible_a["task"] == "Drive safely."
    assert mask_a.shape == (5,)
    text = format_policy_prompt(visible_a)
    assert "Drive safely." in text
    for dropped in set(components) - set(visible_a):
        assert f"{dropped}=" not in text


def test_full_metadata_dropout_keeps_only_task() -> None:
    components = prompt_components(PolicyMetadata("slow", "poor", "stop", "autonomous"))
    visible, mask = drop_prompt_components(components, drop_probability=1.0)
    assert set(visible) == {"task"}
    assert mask.tolist() == [True, False, False, False, False]