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


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


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


def test_step_adds_residue(epitope_env):
    """Stepping with an AA action should extend the sequence."""
    epitope_env.reset(seed=42)
    _, _, _, _, info = epitope_env.step(0)  # A
    assert info["sequence"] == "A"
    assert info["length"] == 1


def test_stop_at_position_9(epitope_env):
    """STOP after building a 9-mer should terminate with a reward."""
    epitope_env.reset(seed=42)
    for aa in "ALAAAAAAV":  # 9 residues
        epitope_env.step(AA_TO_INDEX[aa])
    _, reward, terminated, _, _ = epitope_env.step(20)  # STOP
    assert terminated is True
    assert reward > 0.0  # Should get positive reward for a 9-mer


def test_stop_before_position_8(epitope_env):
    """STOP before 8 residues should give reward = -1."""
    epitope_env.reset(seed=42)
    for aa in "ALAAA":  # 5 residues, below minimum
        epitope_env.step(AA_TO_INDEX[aa])
    _, reward, terminated, _, _ = epitope_env.step(20)
    assert terminated is True
    assert reward == pytest.approx(-1.0)


def test_auto_stop_at_position_11(epitope_env):
    """Building 11 residues (max_length) should auto-terminate."""
    epitope_env.reset(seed=42)
    done = False
    for i in range(11):
        if done:
            break
        _, _, terminated, truncated, info = epitope_env.step(i % 20)
        done = terminated or truncated
    assert done
    assert info["length"] == 11


def test_obs_has_hla_encoding(epitope_env):
    """Observation should contain hla_encoding."""
    obs, _ = epitope_env.reset(seed=42)
    assert "hla_encoding" in obs


def test_hla_encoding_shape_34(epitope_env):
    """hla_encoding should be a 34-dim vector."""
    obs, _ = epitope_env.reset(seed=42)
    assert obs["hla_encoding"].shape == (34,)


def test_binding_estimate_updates(epitope_env):
    """binding_estimate should change as the sequence grows."""
    obs_initial, _ = epitope_env.reset(seed=42)
    initial_be = obs_initial["binding_estimate"][0]
    # Build several residues
    for aa in "ALAAA":
        obs, _, _, _, _ = epitope_env.step(AA_TO_INDEX[aa])
    updated_be = obs["binding_estimate"][0]
    # After adding residues, the binding estimate should differ from empty
    assert initial_be != updated_be or initial_be == 0.0


def test_anchor_shaping_reward(epitope_env):
    """Placing L at P2 (position 2) should yield anchor shaping reward."""
    epitope_env.reset(seed=42)
    epitope_env.step(0)  # A at position 1
    _, reward, _, _, _ = epitope_env.step(AA_TO_INDEX["L"])  # L at position 2 (P2)
    # P2 = index 1 in 0-indexed, but the env checks pos==2 (after append, len==2)
    assert reward > 0.0


def test_obs_sequence_shape(epitope_env):
    """sequence_onehot should be (11, 20) for epitope env."""
    obs, _ = epitope_env.reset(seed=42)
    assert obs["sequence_onehot"].shape == (11, 20)


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

    obs2, _ = epitope_env.reset(seed=77)
    epitope_env.step(0)
    obs_b, _, _, _, _ = epitope_env.step(1)

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


def test_info_contains_binding_score(epitope_env):
    """info should contain binding_score after stepping."""
    epitope_env.reset(seed=42)
    for aa in "ALA":
        epitope_env.step(AA_TO_INDEX[aa])
    _, _, _, _, info = epitope_env.step(AA_TO_INDEX["A"])
    assert "binding_score" in info


def test_hard_mode_multi_allele(epitope_hard_env):
    """Hard mode should score across multiple HLA alleles."""
    epitope_hard_env.reset(seed=42)
    env = epitope_hard_env.unwrapped
    assert len(env._alleles) > 1, "Hard mode should have multiple alleles"
    # Build a 9-mer and stop
    for aa in "ALAAAAAAV":
        epitope_hard_env.step(AA_TO_INDEX[aa])
    _, reward, terminated, _, _ = epitope_hard_env.step(20)
    assert terminated is True
    # Reward should still be a valid float
    assert isinstance(reward, float)
