from __future__ import annotations

import torch

from pi06mini.recap import (
    apply_condition_dropout,
    classifier_free_guidance,
    condition_dropout_mask,
)


def test_default_condition_dropout_is_statistically_thirty_percent() -> None:
    generator = torch.Generator().manual_seed(11)
    mask = condition_dropout_mask(20_000, generator=generator)

    assert mask.dtype == torch.bool
    assert abs(float(mask.float().mean()) - 0.3) < 0.015


def test_condition_dropout_replaces_whole_examples() -> None:
    conditions = torch.tensor([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]])
    dropped, mask = apply_condition_dropout(
        conditions,
        null_condition=-1.0,
        mask=torch.tensor([False, True, False]),
    )

    assert mask.tolist() == [False, True, False]
    torch.testing.assert_close(dropped[0], conditions[0])
    torch.testing.assert_close(dropped[1], torch.tensor([-1.0, -1.0]))
    torch.testing.assert_close(dropped[2], conditions[2])


def test_cfg_interpolates_and_extrapolates_flow_predictions() -> None:
    conditional = torch.tensor([[2.0, 4.0]])
    unconditional = torch.tensor([[0.0, 2.0]])

    torch.testing.assert_close(
        classifier_free_guidance(conditional, unconditional, guidance_scale=0.0),
        unconditional,
    )
    torch.testing.assert_close(
        classifier_free_guidance(conditional, unconditional, guidance_scale=1.0),
        conditional,
    )
    torch.testing.assert_close(
        classifier_free_guidance(conditional, unconditional, guidance_scale=1.5),
        torch.tensor([[3.0, 5.0]]),
    )