"""Tests for observation history functionality."""

from unittest.mock import Mock

import pytest
import torch
from conftest import get_test_device

from mjlab.managers.observation_manager import (
  ObservationGroupCfg,
  ObservationManager,
  ObservationTermCfg,
)


@pytest.fixture
def device():
  """Test device fixture."""
  return get_test_device()


@pytest.fixture
def mock_env(device):
  """Create a mock environment."""
  env = Mock()
  env.num_envs = 4
  env.device = device
  env.step_dt = 0.02
  return env


@pytest.fixture
def simple_obs_func(device):
  """Create a simple observation function that returns a counter."""
  counter = {"value": 0}

  def obs_func(env):
    counter["value"] += 1
    return torch.full((env.num_envs, 3), float(counter["value"]), device=device)

  return obs_func


# Basic observation history tests.


def test_no_history_by_default(mock_env, simple_obs_func):
  """Test that observations work without history (default behavior)."""

  cfg = {
    "actor": ObservationGroupCfg(
      terms={
        "obs1": ObservationTermCfg(func=simple_obs_func, params={}),
      }
    ),
  }

  manager = ObservationManager(cfg, mock_env)
  assert manager.group_obs_dim["actor"] == (3,)

  obs = manager.compute()
  policy_obs = obs["actor"]
  assert isinstance(policy_obs, torch.Tensor)
  assert policy_obs.shape == (4, 3)


def test_single_step_history(mock_env, simple_obs_func):
  """Test observation with history_length=1."""

  cfg = {
    "actor": ObservationGroupCfg(
      terms={
        "obs1": ObservationTermCfg(func=simple_obs_func, params={}, history_length=1),
      }
    ),
  }

  manager = ObservationManager(cfg, mock_env)
  assert manager.group_obs_dim["actor"] == (3,)

  obs = manager.compute()
  policy_obs = obs["actor"]
  assert isinstance(policy_obs, torch.Tensor)
  assert policy_obs.shape == (4, 3)


def test_multi_step_history_flattened(mock_env, simple_obs_func):
  """Test observation with history_length=3 and flattened."""

  cfg = {
    "actor": ObservationGroupCfg(
      terms={
        "obs1": ObservationTermCfg(
          func=simple_obs_func, params={}, history_length=3, flatten_history_dim=True
        ),
      }
    ),
  }

  manager = ObservationManager(cfg, mock_env)
  assert manager.group_obs_dim["actor"] == (9,)

  obs = manager.compute()
  policy_obs = obs["actor"]
  assert isinstance(policy_obs, torch.Tensor)
  assert policy_obs.shape == (4, 9)


def test_multi_step_history_not_flattened(mock_env, simple_obs_func):
  """Test observation with history_length=3 and not flattened."""

  cfg = {
    "actor": ObservationGroupCfg(
      terms={
        "obs1": ObservationTermCfg(
          func=simple_obs_func, params={}, history_length=3, flatten_history_dim=False
        ),
      }
    ),
  }

  manager = ObservationManager(cfg, mock_env)
  assert manager.group_obs_dim["actor"] == (3, 3)

  obs = manager.compute()
  policy_obs = obs["actor"]
  assert isinstance(policy_obs, torch.Tensor)
  assert policy_obs.shape == (4, 3, 3)


# History accumulation tests.


def test_history_accumulates_correctly(mock_env, simple_obs_func):
  """Test that history buffer accumulates observations in correct order."""

  cfg = {
    "actor": ObservationGroupCfg(
      terms={
        "obs1": ObservationTermCfg(
          func=simple_obs_func, params={}, history_length=3, flatten_history_dim=False
        ),
      }
    ),
  }

  manager = ObservationManager(cfg, mock_env)
  device = mock_env.device
  # Note: counter is incremented during _prepare_terms (value=1).

  # First compute uses value=2 and initializes buffer.
  obs = manager.compute(update_history=False)
  policy_obs = obs["actor"]
  assert isinstance(policy_obs, torch.Tensor)
  assert torch.allclose(policy_obs[0], torch.full((3, 3), 2.0, device=device))

  # Update with value=3.
  obs = manager.compute(update_history=True)
  policy_obs = obs["actor"]
  assert isinstance(policy_obs, torch.Tensor)
  # History: [2, 2, 3] (oldest to newest).
  expected = torch.stack(
    [
      torch.full((3,), 2.0, device=device),
      torch.full((3,), 2.0, device=device),
      torch.full((3,), 3.0, device=device),
    ]
  )
  assert torch.allclose(policy_obs[0], expected)

  # Update with value=4.
  obs = manager.compute(update_history=True)
  policy_obs = obs["actor"]
  assert isinstance(policy_obs, torch.Tensor)
  # History: [2, 3, 4].
  expected = torch.stack(
    [
      torch.full((3,), 2.0, device=device),
      torch.full((3,), 3.0, device=device),
      torch.full((3,), 4.0, device=device),
    ]
  )
  assert torch.allclose(policy_obs[0], expected)

  # Update with value=5, circular overwrite.
  obs = manager.compute(update_history=True)
  policy_obs = obs["actor"]
  assert isinstance(policy_obs, torch.Tensor)
  # History: [3, 4, 5].
  expected = torch.stack(
    [
      torch.full((3,), 3.0, device=device),
      torch.full((3,), 4.0, device=device),
      torch.full((3,), 5.0, device=device),
    ]
  )
  assert torch.allclose(policy_obs[0], expected)


def test_update_history_false_doesnt_modify_buffer(mock_env, simple_obs_func):
  """Test that update_history=False doesn't modify the buffer."""

  cfg = {
    "actor": ObservationGroupCfg(
      terms={
        "obs1": ObservationTermCfg(
          func=simple_obs_func, params={}, history_length=2, flatten_history_dim=False
        ),
      }
    ),
  }

  manager = ObservationManager(cfg, mock_env)

  # Initialize (value=1).
  obs1 = manager.compute(update_history=False)
  policy_obs1 = obs1["actor"]
  assert isinstance(policy_obs1, torch.Tensor)

  # Call without update (value=2, but buffer unchanged).
  obs2 = manager.compute(update_history=False)
  policy_obs2 = obs2["actor"]
  assert isinstance(policy_obs2, torch.Tensor)

  # History should still be [1, 1].
  assert torch.allclose(policy_obs1, policy_obs2)


# Group-level history override tests.


def test_group_history_overrides_term(mock_env, simple_obs_func):
  """Test group history_length overrides term history_length."""

  cfg = {
    "actor": ObservationGroupCfg(
      history_length=5,  # Group level.
      flatten_history_dim=False,
      terms={
        "obs1": ObservationTermCfg(
          func=simple_obs_func,
          params={},
          history_length=2,  # Overridden.
        ),
      },
    ),
  }

  manager = ObservationManager(cfg, mock_env)
  assert manager.group_obs_dim["actor"] == (5, 3)

  obs = manager.compute()
  policy_obs = obs["actor"]
  assert isinstance(policy_obs, torch.Tensor)
  assert policy_obs.shape == (4, 5, 3)


# History reset tests.


def test_reset_clears_all_envs(mock_env, simple_obs_func):
  """Test that reset without env_ids clears all environments."""

  cfg = {
    "actor": ObservationGroupCfg(
      terms={
        "obs1": ObservationTermCfg(
          func=simple_obs_func, params={}, history_length=2, flatten_history_dim=False
        ),
      }
    ),
  }

  manager = ObservationManager(cfg, mock_env)

  # Build up history.
  manager.compute(update_history=True)
  manager.compute(update_history=True)

  # Reset all.
  manager.reset()

  # Buffer should be zeroed.
  obs = manager.compute(update_history=False)
  policy_obs = obs["actor"]
  assert isinstance(policy_obs, torch.Tensor)
  assert torch.allclose(policy_obs, torch.zeros((4, 2, 3), device=mock_env.device))


def test_reset_partial_envs(mock_env, simple_obs_func):
  """Test that reset with specific env_ids only resets those envs."""

  cfg = {
    "actor": ObservationGroupCfg(
      terms={
        "obs1": ObservationTermCfg(
          func=simple_obs_func, params={}, history_length=3, flatten_history_dim=False
        ),
      }
    ),
  }

  manager = ObservationManager(cfg, mock_env)

  # Build up history.
  manager.compute(update_history=True)
  manager.compute(update_history=True)
  obs_before = manager.compute(update_history=True)
  policy_obs_before = obs_before["actor"]
  assert isinstance(policy_obs_before, torch.Tensor)

  # Reset only envs 0 and 2.
  manager.reset(env_ids=torch.tensor([0, 2], device=mock_env.device))

  obs_after = manager.compute(update_history=False)
  policy_obs_after = obs_after["actor"]
  assert isinstance(policy_obs_after, torch.Tensor)

  # Envs 0 and 2 should be reset (zeros), 1 and 3 unchanged.
  assert torch.allclose(
    policy_obs_after[0], torch.zeros((3, 3), device=mock_env.device)
  )
  assert torch.allclose(policy_obs_after[1], policy_obs_before[1])
  assert torch.allclose(
    policy_obs_after[2], torch.zeros((3, 3), device=mock_env.device)
  )
  assert torch.allclose(policy_obs_after[3], policy_obs_before[3])


def test_reset_partial_envs_with_backfill(mock_env, simple_obs_func):
  """Test that reset envs get backfilled on next update."""

  cfg = {
    "actor": ObservationGroupCfg(
      terms={
        "obs1": ObservationTermCfg(
          func=simple_obs_func, params={}, history_length=3, flatten_history_dim=False
        ),
      }
    ),
  }

  manager = ObservationManager(cfg, mock_env)
  device = mock_env.device

  # Build history: [2, 3, 4] for all envs (counter starts at 1 after _prepare_terms).
  manager.compute(update_history=True)  # value=2
  manager.compute(update_history=True)  # value=3
  manager.compute(update_history=True)  # value=4

  # Reset only envs 0 and 2.
  manager.reset(env_ids=torch.tensor([0, 2], device=device))

  # Next update with value=5.
  obs = manager.compute(update_history=True)
  policy_obs = obs["actor"]
  assert isinstance(policy_obs, torch.Tensor)

  # Env 0: [5, 5, 5] (backfilled after reset).
  expected_env0 = torch.stack(
    [
      torch.full((3,), 5.0, device=device),
      torch.full((3,), 5.0, device=device),
      torch.full((3,), 5.0, device=device),
    ]
  )
  assert torch.allclose(policy_obs[0], expected_env0)

  # Env 1: [3, 4, 5] (continues normally).
  expected_env1 = torch.stack(
    [
      torch.full((3,), 3.0, device=device),
      torch.full((3,), 4.0, device=device),
      torch.full((3,), 5.0, device=device),
    ]
  )
  assert torch.allclose(policy_obs[1], expected_env1)

  # Env 2: [5, 5, 5] (backfilled after reset).
  assert torch.allclose(policy_obs[2], expected_env0)

  # Env 3: [3, 4, 5] (continues normally).
  assert torch.allclose(policy_obs[3], expected_env1)


# History with other features tests.


def test_history_with_clip(mock_env, simple_obs_func):
  """Test that clipping is applied before history."""

  cfg = {
    "actor": ObservationGroupCfg(
      terms={
        "obs1": ObservationTermCfg(
          func=simple_obs_func,
          params={},
          history_length=2,
          flatten_history_dim=False,
          clip=(-0.5, 0.5),
        ),
      }
    ),
  }

  manager = ObservationManager(cfg, mock_env)

  obs = manager.compute(update_history=True)
  policy_obs = obs["actor"]
  assert isinstance(policy_obs, torch.Tensor)
  # Values should be clipped.
  assert torch.all(policy_obs >= -0.5)
  assert torch.all(policy_obs <= 0.5)


def test_history_with_scale(mock_env, simple_obs_func):
  """Test that scaling is applied before history."""

  cfg = {
    "actor": ObservationGroupCfg(
      terms={
        "obs1": ObservationTermCfg(
          func=simple_obs_func,
          params={},
          history_length=2,
          flatten_history_dim=False,
          scale=2.0,
        ),
      }
    ),
  }

  manager = ObservationManager(cfg, mock_env)
  device = mock_env.device
  # Counter at 1 after _prepare_terms.

  # First call uses value=2, scaled to 4, initializes buffer.
  manager.compute(update_history=False)
  # Second call uses value=3, scaled to 6, and updates history.
  obs = manager.compute(update_history=True)
  policy_obs = obs["actor"]
  assert isinstance(policy_obs, torch.Tensor)

  # History: [4, 6].
  expected = torch.stack(
    [torch.full((3,), 4.0, device=device), torch.full((3,), 6.0, device=device)]
  )
  assert torch.allclose(policy_obs[0], expected)


# Mixed history terms tests.


def test_mixed_terms_concatenated(mock_env, simple_obs_func, device):
  """Test group with both history and non-history terms concatenated."""

  counter = {"value": 0}

  def obs_func2(env):
    counter["value"] += 1
    return torch.full((env.num_envs, 2), float(counter["value"]) * 10, device=device)

  cfg = {
    "actor": ObservationGroupCfg(
      terms={
        "obs_with_history": ObservationTermCfg(
          func=simple_obs_func, params={}, history_length=2, flatten_history_dim=True
        ),
        "obs_no_history": ObservationTermCfg(func=obs_func2, params={}),
      }
    ),
  }

  manager = ObservationManager(cfg, mock_env)

  # Should concatenate: (3*2) + 2 = 8.
  assert manager.group_obs_dim["actor"] == (8,)

  obs = manager.compute()
  policy_obs = obs["actor"]
  assert isinstance(policy_obs, torch.Tensor)
  assert policy_obs.shape == (4, 8)


def test_no_double_append_on_first_call(mock_env, simple_obs_func):
  """Test that first call with update_history=True only appends once, not twice."""

  cfg = {
    "actor": ObservationGroupCfg(
      terms={
        "obs1": ObservationTermCfg(
          func=simple_obs_func, params={}, history_length=3, flatten_history_dim=False
        ),
      }
    ),
  }

  manager = ObservationManager(cfg, mock_env)
  device = mock_env.device
  # Counter is at 1 after _prepare_terms.

  # First call with update_history=True (value=2).
  # This should initialize the buffer AND append once (not twice).
  obs = manager.compute(update_history=True)
  policy_obs = obs["actor"]
  assert isinstance(policy_obs, torch.Tensor)

  # Verify buffer was initialized and backfilled correctly.
  # All slots should be filled with value 2.
  expected_first = torch.stack(
    [
      torch.full((3,), 2.0, device=device),
      torch.full((3,), 2.0, device=device),
      torch.full((3,), 2.0, device=device),
    ]
  )
  assert torch.allclose(policy_obs[0], expected_first)

  # Get the circular buffer and check pointer position.
  circular_buffer = manager._group_obs_term_history_buffer["actor"]["obs1"]
  # After one append, pointer should be at 0 (not 1 which would indicate double-append).
  assert circular_buffer._pointer == 0
  # And we should have exactly 1 push recorded.
  assert torch.all(circular_buffer._num_pushes == 1)

  # Second call with update_history=True (value=3).
  obs = manager.compute(update_history=True)
  policy_obs = obs["actor"]
  assert isinstance(policy_obs, torch.Tensor)

  # History should be [2, 2, 3] (oldest to newest).
  expected_second = torch.stack(
    [
      torch.full((3,), 2.0, device=device),
      torch.full((3,), 2.0, device=device),
      torch.full((3,), 3.0, device=device),
    ]
  )
  assert torch.allclose(policy_obs[0], expected_second)

  # Pointer should now be at 1.
  assert circular_buffer._pointer == 1
  # And we should have exactly 2 pushes.
  assert torch.all(circular_buffer._num_pushes == 2)


def test_term_major_ordering(mock_env, device):
  """Test that history flattening uses term-major ordering.

  Term-major: [A_t0, A_t1, ..., A_tH-1, B_t0, B_t1, ..., B_tH-1]
  Time-major: [A_t0, B_t0, ..., A_t1, B_t1, ..., A_tH-1, B_tH-1]

  This test verifies mjlab uses term-major ordering.
  """

  # Create observation functions with distinct values per term.
  def obs_A(env):
    # Returns [100, 101] for easy identification.
    return torch.tensor([[100.0, 101.0]] * env.num_envs, device=device)

  def obs_B(env):
    # Returns [200, 201, 202] for easy identification.
    return torch.tensor([[200.0, 201.0, 202.0]] * env.num_envs, device=device)

  cfg = {
    "actor": ObservationGroupCfg(
      concatenate_terms=True,
      terms={
        "term_A": ObservationTermCfg(
          func=obs_A, params={}, history_length=3, flatten_history_dim=True
        ),
        "term_B": ObservationTermCfg(
          func=obs_B, params={}, history_length=3, flatten_history_dim=True
        ),
      },
    ),
  }

  manager = ObservationManager(cfg, mock_env)

  # Compute observations (history will be backfilled with same values).
  obs = manager.compute(update_history=False)
  policy_obs = obs["actor"]
  assert isinstance(policy_obs, torch.Tensor)

  # Expected shape: (4 envs, 2*3 + 3*3) = (4, 15).
  assert policy_obs.shape == (4, 15)

  # Check ordering for first environment.
  result = policy_obs[0].cpu().tolist()

  # Term-major: all A history, then all B history.
  # A has 2 dims, B has 3 dims, history_length=3.
  # fmt: off
  expected_term_major = [
    100.0, 101.0, 100.0, 101.0, 100.0, 101.0,  # All A history
    200.0, 201.0, 202.0, 200.0, 201.0, 202.0, 200.0, 201.0, 202.0,  # All B
  ]
  # fmt: on

  # Time-major would interleave: [A_t0, B_t0, A_t1, B_t1, A_t2, B_t2].
  # fmt: off
  expected_time_major = [
    100.0, 101.0, 200.0, 201.0, 202.0,  # Frame t0
    100.0, 101.0, 200.0, 201.0, 202.0,  # Frame t1
    100.0, 101.0, 200.0, 201.0, 202.0,  # Frame t2
  ]
  # fmt: on

  # Verify term-major ordering.
  assert result == expected_term_major, f"Expected term-major ordering, got {result}"
  assert result != expected_time_major, "Should not match time-major"