import pytest

torch = pytest.importorskip("torch")

from hippotorch.core.episode import Episode
from hippotorch.segmenters.base import (
    ChangePointSegmenter,
    FixedWindowSegmenter,
    GoalAchievementSegmenter,
    TerminalSegmenter,
)


def _episode_with_flags(length: int, done_positions: list[int] | None = None) -> Episode:
    done_positions = done_positions or []
    states = torch.arange(length, dtype=torch.float32).view(length, 1)
    actions = torch.zeros(length, 1)
    rewards = torch.ones(length)
    dones = torch.zeros(length, dtype=torch.bool)
    for idx in done_positions:
        if 0 <= idx < length:
            dones[idx] = True
    return Episode(states=states, actions=actions, rewards=rewards, dones=dones)


def test_terminal_segmenter_splits_on_done_flags():
    ep = _episode_with_flags(6, done_positions=[1, 4])
    seg = TerminalSegmenter()
    parts = seg.segment(ep)
    assert [len(p) for p in parts] == [2, 3, 1]


def test_fixed_window_segmenter_partitions_evenly():
    ep = _episode_with_flags(7)
    seg = FixedWindowSegmenter(window_size=3)
    parts = seg.segment(ep)
    assert [len(p) for p in parts] == [3, 3, 1]


def test_goal_achievement_segmenter_cuts_on_predicate():
    ep = _episode_with_flags(5)
    # goal when state >= 2.0
    seg = GoalAchievementSegmenter(goal_fn=lambda s: bool(s.item() >= 2.0), min_length=2)
    parts = seg.segment(ep)
    # Expect at least one split and all segments >= min_length except possibly the last
    assert len(parts) >= 2
    assert all(len(p) >= 2 for p in parts[:-1])


def test_change_point_segmenter_detects_anomalies_and_limits_length():
    # Create mostly smooth states with a jump and long tail
    vals = torch.tensor([0.0, 0.1, 0.12, 3.0, 3.1, 3.2, 10.0, 10.1, 10.2, 10.3])
    states = vals.view(-1, 1)
    actions = torch.zeros(len(vals), 1)
    rewards = torch.ones(len(vals))
    dones = torch.zeros(len(vals), dtype=torch.bool)
    traj = Episode(states=states, actions=actions, rewards=rewards, dones=dones)

    seg = ChangePointSegmenter(threshold=2.0, min_length=2, max_length=4)
    parts = seg.segment(traj)
    # Expect multiple segments due to anomaly and max_length
    assert len(parts) >= 2
    assert all(len(p) >= 1 for p in parts)
