"""CPU-only contracts for symmetry augmentation."""
from __future__ import annotations

import sys
from pathlib import Path

import gymnasium as gym
import numpy as np
import pytest

_TRAIN_DIR = Path(__file__).resolve().parents[1]
if str(_TRAIN_DIR) not in sys.path:
    sys.path.insert(0, str(_TRAIN_DIR))

from symmetry_wrapper import (  # noqa: E402
    SymmetryAugmentWrapper,
    apply_G_action,
    apply_G_obs,
)


class RecordingEnv(gym.Env):
    metadata = {}

    def __init__(self, obs_dim: int):
        self.observation_space = gym.spaces.Box(-np.inf, np.inf, (obs_dim,), dtype=np.float32)
        self.action_space = gym.spaces.Box(-1.0, 1.0, (3,), dtype=np.float32)
        self.obs = np.arange(1, obs_dim + 1, dtype=np.float32)
        self.actions: list[np.ndarray] = []

    def reset(self, *, seed=None, options=None):
        super().reset(seed=seed)
        return self.obs.copy(), {}

    def step(self, action):
        self.actions.append(np.asarray(action).copy())
        self.obs = self.obs + 1
        return self.obs.copy(), 1.0, False, False, {}


@pytest.mark.parametrize("dim", [12, 14])
def test_observation_transform_is_an_involution(dim):
    obs = np.arange(1, dim + 1, dtype=np.float32)
    np.testing.assert_array_equal(apply_G_obs(apply_G_obs(obs)), obs)


def test_action_transform_and_involution():
    action = np.array([1.0, 2.0, 3.0], dtype=np.float32)
    np.testing.assert_array_equal(apply_G_action(action), [-1.0, -3.0, -2.0])
    np.testing.assert_array_equal(apply_G_action(apply_G_action(action)), action)


def test_exact_12d_karman_transform():
    obs = np.arange(1, 13, dtype=np.float32)
    expected = np.array([1, -2, 5, -6, 3, -4, 11, -12, 9, -10, 7, -8])
    np.testing.assert_array_equal(apply_G_obs(obs), expected)


def test_exact_14d_illusion_transform_preserves_cd_and_negates_cl():
    obs = np.arange(1, 15, dtype=np.float32)
    expected = np.array([1, -2, 5, -6, 3, -4, 11, -12, 9, -10, 7, -8, 13, -14])
    np.testing.assert_array_equal(apply_G_obs(obs), expected)


@pytest.mark.parametrize("dim", [0, 11, 13, 15])
def test_unsupported_observation_dimensions_are_rejected(dim):
    with pytest.raises(ValueError, match="12D Karman or 14D Illusion"):
        apply_G_obs(np.zeros(dim, dtype=np.float32))
    with pytest.raises(ValueError, match="12D Karman or 14D Illusion"):
        SymmetryAugmentWrapper(RecordingEnv(dim))


def test_mirrored_frame_is_selected_before_and_fixed_through_rollout():
    raw = RecordingEnv(12)
    wrapped = SymmetryAugmentWrapper(raw, prob=1.0, seed=7, rollout_len=2)

    initial_raw = raw.obs.copy()
    obs, _ = wrapped.reset()
    np.testing.assert_array_equal(obs, apply_G_obs(initial_raw))

    action = np.array([0.25, -0.5, 0.75], dtype=np.float32)
    for _ in range(5):
        obs, reward, terminated, truncated, _ = wrapped.step(action)
        np.testing.assert_array_equal(raw.actions[-1], apply_G_action(action))
        np.testing.assert_array_equal(obs, apply_G_obs(raw.obs))
        assert reward == 1.0
        assert not terminated and not truncated


def test_probability_zero_preserves_evaluation_frame():
    raw = RecordingEnv(14)
    wrapped = SymmetryAugmentWrapper(raw, prob=0.0, seed=7)
    obs, _ = wrapped.reset()
    np.testing.assert_array_equal(obs, raw.obs)
    action = np.array([0.25, -0.5, 0.75], dtype=np.float32)
    obs, *_ = wrapped.step(action)
    np.testing.assert_array_equal(raw.actions[-1], action)
    np.testing.assert_array_equal(obs, raw.obs)
