"""Tests for the AMP (antimicrobial peptide) Gymnasium environment."""
import gymnasium as gym
import numpy as np
import pytest

import peptidegym  # noqa: F401
from peptidegym.peptide.properties import AA_TO_INDEX


def test_env_creates(amp_env):
    """AMPEnv should be constructible via gym.make."""
    assert amp_env is not None


def test_reset_returns_obs_info(amp_env):
    """reset() should return (obs_dict, info_dict)."""
    obs, info = amp_env.reset(seed=42)
    assert isinstance(obs, dict)
    assert isinstance(info, dict)


def test_obs_space_contains_obs(amp_env):
    """Observation returned by reset should be within observation_space."""
    obs, _ = amp_env.reset(seed=42)
    assert amp_env.observation_space.contains(obs)


def test_action_space_is_discrete_21(amp_env):
    """AMP env has 20 AAs + 1 STOP = 21 actions."""
    assert amp_env.action_space.n == 21


def test_step_with_amino_acid(amp_env):
    """Stepping with a valid AA action should return 5-tuple."""
    amp_env.reset(seed=42)
    obs, reward, terminated, truncated, info = amp_env.step(0)  # Alanine
    assert isinstance(obs, dict)
    assert isinstance(reward, float)
    assert isinstance(terminated, bool)
    assert isinstance(truncated, bool)
    assert isinstance(info, dict)


def test_step_sequence_grows(amp_env):
    """Each AA action should extend the sequence by one residue."""
    amp_env.reset(seed=42)
    amp_env.step(0)  # A
    amp_env.step(1)  # R
    _, _, _, _, info = amp_env.step(2)  # N
    assert info["sequence"] == "ARN"
    assert info["length"] == 3


def test_stop_action_terminates(amp_env):
    """STOP (action 20) should terminate the episode."""
    amp_env.reset(seed=42)
    amp_env.step(0)
    _, reward, terminated, _, _ = amp_env.step(20)
    assert terminated is True


def test_max_length_truncates(amp_env):
    """Reaching max_length should truncate the episode."""
    amp_env.reset(seed=42)
    max_len = amp_env.unwrapped.max_length
    done = False
    for i in range(max_len + 5):
        if done:
            break
        _, _, terminated, truncated, info = amp_env.step(i % 20)
        done = terminated or truncated
    assert done, "Episode should end at max_length"


def test_reward_on_stop(amp_env):
    """Stopping a multi-residue peptide should produce a non-zero reward."""
    amp_env.reset(seed=42)
    # Build a cationic peptide
    for aa in "KRLKRL":
        amp_env.step(AA_TO_INDEX[aa])
    _, reward, terminated, _, _ = amp_env.step(20)
    assert terminated is True
    assert reward != 0.0


def test_obs_keys(amp_env):
    """Observation dict should have expected keys."""
    obs, _ = amp_env.reset(seed=42)
    assert "sequence_onehot" in obs
    assert "position" in obs
    assert "properties" in obs


def test_obs_sequence_shape(amp_env):
    """sequence_onehot should be (max_length, 20)."""
    obs, _ = amp_env.reset(seed=42)
    max_len = amp_env.unwrapped.max_length
    assert obs["sequence_onehot"].shape == (max_len, 20)


def test_obs_properties_shape(amp_env):
    """properties should be 6-dim for medium difficulty."""
    obs, _ = amp_env.reset(seed=42)
    assert obs["properties"].shape == (6,)


def test_seed_reproducibility(amp_env):
    """Same seed should produce same observations."""
    obs1, _ = amp_env.reset(seed=123)
    amp_env.step(0)
    obs_a, _, _, _, _ = amp_env.step(1)

    obs2, _ = amp_env.reset(seed=123)
    amp_env.step(0)
    obs_b, _, _, _, _ = amp_env.step(1)

    np.testing.assert_array_equal(obs_a["sequence_onehot"], obs_b["sequence_onehot"])


def test_info_contains_sequence(amp_env):
    """info dict should contain 'sequence' key."""
    amp_env.reset(seed=42)
    _, _, _, _, info = amp_env.step(0)
    assert "sequence" in info
    assert isinstance(info["sequence"], str)


def test_multiple_episodes(amp_env):
    """Running multiple episodes should work cleanly."""
    for _ in range(3):
        obs, info = amp_env.reset(seed=42)
        done = False
        steps = 0
        while not done and steps < 10:
            obs, _, terminated, truncated, info = amp_env.step(steps % 20)
            done = terminated or truncated
            steps += 1
        assert info["length"] > 0
