"""Tests for PeptideGym environment wrappers."""
import gymnasium as gym
import numpy as np
import pytest

import peptidegym  # noqa: F401
from peptidegym.envs.wrappers import FlattenObservation, RewardShapingWrapper
from peptidegym.peptide.properties import AA_TO_INDEX


def test_flatten_observation_creates():
    """FlattenObservation wrapper should be constructible."""
    env = gym.make("PeptideGym/AMP-v0")
    wrapped = FlattenObservation(env)
    assert wrapped is not None
    wrapped.close()


def test_flatten_observation_box_space():
    """Wrapped observation space should be a flat Box."""
    env = gym.make("PeptideGym/AMP-v0")
    wrapped = FlattenObservation(env)
    assert isinstance(wrapped.observation_space, gym.spaces.Box)
    assert len(wrapped.observation_space.shape) == 1
    wrapped.close()


def test_flatten_observation_reset():
    """reset() should return a flat numpy array."""
    env = gym.make("PeptideGym/AMP-v0")
    wrapped = FlattenObservation(env)
    obs, info = wrapped.reset(seed=42)
    assert isinstance(obs, np.ndarray)
    assert obs.ndim == 1
    assert wrapped.observation_space.contains(obs)
    wrapped.close()


def test_flatten_observation_step():
    """step() should return a flat observation."""
    env = gym.make("PeptideGym/AMP-v0")
    wrapped = FlattenObservation(env)
    wrapped.reset(seed=42)
    obs, reward, terminated, truncated, info = wrapped.step(0)
    assert isinstance(obs, np.ndarray)
    assert obs.ndim == 1
    wrapped.close()


def test_reward_shaping_wrapper():
    """RewardShapingWrapper should apply tanh squashing."""
    env = gym.make("PeptideGym/AMP-v0")
    wrapped = RewardShapingWrapper(env, scale=1.0)
    wrapped.reset(seed=42)
    # Build a short peptide and stop to get a terminal reward
    for aa in "KRLKRL":
        wrapped.step(AA_TO_INDEX[aa])
    _, reward, _, _, _ = wrapped.step(20)  # STOP
    # tanh maps any real to (-1, 1)
    assert -1.0 <= reward <= 1.0
    wrapped.close()


def test_reward_shaping_range():
    """All rewards from RewardShapingWrapper should be in [-1, 1]."""
    env = gym.make("PeptideGym/AMP-v0")
    wrapped = RewardShapingWrapper(env, scale=2.0)
    wrapped.reset(seed=42)
    rewards = []
    for i in range(10):
        _, reward, terminated, truncated, _ = wrapped.step(i % 20)
        rewards.append(reward)
        if terminated or truncated:
            break
    for r in rewards:
        assert -1.0 <= r <= 1.0
    wrapped.close()


def test_wrapper_stacking():
    """Flatten + RewardShaping should compose correctly."""
    env = gym.make("PeptideGym/AMP-v0")
    wrapped = RewardShapingWrapper(FlattenObservation(env), scale=1.0)
    obs, _ = wrapped.reset(seed=42)
    assert isinstance(obs, np.ndarray)
    assert obs.ndim == 1
    obs2, reward, _, _, _ = wrapped.step(0)
    assert isinstance(obs2, np.ndarray)
    assert -1.0 <= reward <= 1.0
    wrapped.close()


def test_wrapper_with_all_envs():
    """FlattenObservation should work with AMP, Cyclic, and Epitope envs."""
    env_ids = [
        "PeptideGym/AMP-v0",
        "PeptideGym/CyclicPeptide-v0",
        "PeptideGym/Epitope-v0",
    ]
    for env_id in env_ids:
        env = gym.make(env_id)
        wrapped = FlattenObservation(env)
        obs, _ = wrapped.reset(seed=42)
        assert isinstance(obs, np.ndarray), f"Failed for {env_id}"
        assert obs.ndim == 1, f"Not flat for {env_id}"
        wrapped.close()
