from __future__ import annotations

import torch

from pi07mini.context import (
    PolicyMetadata,
    build_ambiguous_history_pair,
    drop_prompt_components,
    endpoint_subgoal_heatmap,
    format_policy_prompt,
    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_subgoal_heatmap_depends_on_endpoint_only() -> None:
    endpoint = torch.tensor([20.0, 4.0])
    first = endpoint_subgoal_heatmap(endpoint, size=24)
    # Two future paths with radically different interiors produce exactly the
    # same condition because only their shared endpoint is passed to the API.
    path_a = torch.stack((torch.linspace(0.0, 20.0, 16), torch.linspace(0.0, 4.0, 16)), -1)
    path_b = path_a.clone()
    path_b[1:-1, 1] += 100.0
    second = endpoint_subgoal_heatmap(path_b[-1], size=24)
    torch.testing.assert_close(first, second, atol=0.0, rtol=0.0)
    assert first.shape == (1, 24, 24)
    assert float(first.max()) == 1.0
    assert torch.isfinite(first).all()


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]