"""Tests for the CyclicPeptide 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(cyclic_env):
    """CyclicPeptideEnv should be constructible via gym.make."""
    assert cyclic_env is not None


def test_reset_returns_obs(cyclic_env):
    """reset() returns (obs_dict, info_dict)."""
    obs, info = cyclic_env.reset(seed=42)
    assert isinstance(obs, dict)
    assert isinstance(info, dict)


def test_action_space_is_discrete_24(cyclic_env):
    """Cyclic env has 20 AAs + 3 cyclization + 1 stop_linear = 24 actions."""
    assert cyclic_env.action_space.n == 24


def test_step_with_amino_acid(cyclic_env):
    """Stepping with a valid AA should extend the sequence."""
    cyclic_env.reset(seed=42)
    _, _, _, _, info = cyclic_env.step(0)  # A
    assert info["sequence"] == "A"


def test_cyclize_head_to_tail_valid(cyclic_env):
    """Head-to-tail cyclization with 5+ residues should succeed."""
    cyclic_env.reset(seed=42)
    for aa in "AGLVF":
        cyclic_env.step(AA_TO_INDEX[aa])
    _, reward, terminated, _, info = cyclic_env.step(20)  # cyclize head-to-tail
    assert terminated is True
    assert info["cyclization_type"] == "head_to_tail"
    assert reward > -0.5  # Should not be the penalty


def test_cyclize_head_to_tail_too_short(cyclic_env):
    """Head-to-tail with only 2 residues should give negative reward."""
    cyclic_env.reset(seed=42)
    cyclic_env.step(0)  # A
    cyclic_env.step(1)  # R
    _, reward, terminated, _, _ = cyclic_env.step(20)
    assert terminated is True
    assert reward == pytest.approx(-0.5)


def test_cyclize_disulfide_with_cys(cyclic_env):
    """Disulfide cyclization with 2+ Cys should succeed."""
    cyclic_env.reset(seed=42)
    cys = AA_TO_INDEX["C"]
    for aa_idx in [0, cys, 2, 3, cys]:  # A, C, N, D, C
        cyclic_env.step(aa_idx)
    _, reward, terminated, _, info = cyclic_env.step(21)  # cyclize disulfide
    assert terminated is True
    assert info["cyclization_type"] == "disulfide"


def test_cyclize_disulfide_without_cys(cyclic_env):
    """Disulfide cyclization without Cys should give penalty."""
    cyclic_env.reset(seed=42)
    for aa in "AGLVF":
        cyclic_env.step(AA_TO_INDEX[aa])
    _, reward, terminated, _, _ = cyclic_env.step(21)
    assert terminated is True
    assert reward == pytest.approx(-0.5)


def test_stop_linear(cyclic_env):
    """STOP_LINEAR (action 23) should terminate as linear peptide."""
    cyclic_env.reset(seed=42)
    for aa in "AGLV":
        cyclic_env.step(AA_TO_INDEX[aa])
    _, _, terminated, _, info = cyclic_env.step(23)  # stop linear
    assert terminated is True
    assert info["cyclization_type"] == "none"


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


def test_obs_has_cyclization_valid(cyclic_env):
    """Observation should include cyclization_valid key."""
    obs, _ = cyclic_env.reset(seed=42)
    assert "cyclization_valid" in obs
    assert obs["cyclization_valid"].shape == (3,)


def test_cyclization_valid_updates(cyclic_env):
    """cyclization_valid should change as sequence is built."""
    cyclic_env.reset(seed=42)
    obs_initial, _ = cyclic_env.reset(seed=42)
    # Empty sequence: head_to_tail should be invalid (length < 4)
    assert obs_initial["cyclization_valid"][0] == 0.0

    # Build 5-residue sequence
    for aa in "AGLVF":
        obs, _, _, _, _ = cyclic_env.step(AA_TO_INDEX[aa])
    # Now head_to_tail should be valid
    assert obs["cyclization_valid"][0] == 1.0


def test_seed_reproducibility(cyclic_env):
    """Same seed should produce identical observations."""
    obs1, _ = cyclic_env.reset(seed=99)
    cyclic_env.step(0)
    obs_a, _, _, _, _ = cyclic_env.step(1)

    obs2, _ = cyclic_env.reset(seed=99)
    cyclic_env.step(0)
    obs_b, _, _, _, _ = cyclic_env.step(1)

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


def test_info_contains_cyclization_type(cyclic_env):
    """info should contain cyclization_type."""
    cyclic_env.reset(seed=42)
    _, _, _, _, info = cyclic_env.step(0)
    assert "cyclization_type" in info


def test_easy_mode_restricts_actions():
    """Easy mode should disable disulfide and sidechain cyclization."""
    env = gym.make("PeptideGym/CyclicPeptide-Easy-v0")
    obs, _ = env.reset(seed=42)
    # Build enough residues with Cys to test
    cys = AA_TO_INDEX["C"]
    for aa_idx in [0, cys, 2, 3, cys]:
        obs, _, _, _, _ = env.step(aa_idx)
    # In easy mode, disulfide and sidechain should show as invalid
    assert obs["cyclization_valid"][1] == 0.0  # disulfide disabled
    assert obs["cyclization_valid"][2] == 0.0  # sidechain disabled
    env.close()
