"""AMPEnv — Gymnasium environment for antimicrobial peptide design."""
from __future__ import annotations

from typing import Any

import numpy as np
from gymnasium import spaces

from peptidegym.peptide.properties import (
    AMINO_ACIDS,
    CHARGE_PH74,
    INDEX_TO_AA,
    net_charge,
    mean_hydrophobicity,
    amphipathicity_score,
    aromaticity_fraction,
    instability_index,
    molecular_weight,
)
from peptidegym.peptide.amp_scoring import HeuristicAMPScorer
from peptidegym.envs.base import PeptideEnv

# Action constants
_NUM_AAS = len(AMINO_ACIDS)  # 20
ACTION_STOP = _NUM_AAS       # 20
_MIN_LENGTH = 10             # AMPs need ≥10 residues to be functional


class AMPEnv(PeptideEnv):
    """Design antimicrobial peptides via sequential amino-acid selection.

    Actions 0-19 append the corresponding amino acid.
    Action 20 (STOP) terminates the episode and scores the peptide.
    """

    def __init__(
        self,
        max_length: int = 35,
        difficulty: str = "medium",
        seed: int | None = None,
    ) -> None:
        self.difficulty = difficulty

        reward_weights = {
            "activity_score": 0.4,
            "toxicity_penalty": -0.3,
            "amphipathicity_bonus": 0.2,
            "length_fitness": 0.1,
        }

        scorer = HeuristicAMPScorer()

        super().__init__(
            max_length=max_length,
            scorer=scorer,
            reward_weights=reward_weights,
            seed=seed,
        )

        # Spaces
        n_properties = self._properties_dim()
        self.action_space = spaces.Discrete(ACTION_STOP + 1)
        self.observation_space = spaces.Dict(
            {
                "sequence_onehot": spaces.Box(
                    low=0.0, high=1.0,
                    shape=(max_length, _NUM_AAS),
                    dtype=np.float32,
                ),
                "position": spaces.Box(
                    low=0.0, high=float(max_length),
                    shape=(1,),
                    dtype=np.float32,
                ),
                "properties": spaces.Box(
                    low=-np.inf, high=np.inf,
                    shape=(n_properties,),
                    dtype=np.float32,
                ),
            }
        )

    # ── Difficulty helpers ────────────────────────────────────────────

    def _properties_dim(self) -> int:
        if self.difficulty == "easy":
            return 2  # net_charge, mean_hydro
        return 6       # MW, net_charge, mean_hydro, amphipathicity, aromaticity, instability

    def _compute_properties(self) -> np.ndarray:
        if not self._sequence:
            return np.zeros(self._properties_dim(), dtype=np.float32)
        seq = self._get_sequence_str()
        if self.difficulty == "easy":
            return np.array(
                [net_charge(seq), mean_hydrophobicity(seq)],
                dtype=np.float32,
            )
        return np.array(
            [
                molecular_weight(seq),
                net_charge(seq),
                mean_hydrophobicity(seq),
                amphipathicity_score(seq),
                aromaticity_fraction(seq),
                instability_index(seq),
            ],
            dtype=np.float32,
        )

    # ── Gymnasium API ─────────────────────────────────────────────────

    def step(self, action: int) -> tuple[dict, float, bool, bool, dict]:
        action = int(action)
        assert self.action_space.contains(action), f"Invalid action {action}"

        terminated = False
        truncated = False
        reward = 0.0

        if action < _NUM_AAS:
            # ── Append amino acid ─────────────────────────────────────
            prev_charge = net_charge(self._get_sequence_str()) if self._sequence else 0.0
            self._sequence.append(action)
            self._step_count += 1
            new_charge = net_charge(self._get_sequence_str())

            # Intermediate shaping: reward cationic bias (small to avoid dominating terminal)
            if new_charge > prev_charge:
                reward += 0.005

            # Diversity shaping: penalize low unique-residue ratio
            if len(self._sequence) >= 5:
                diversity = len(set(self._sequence)) / len(self._sequence)
                if diversity < 0.3:
                    reward -= 0.01

            # Truncation: reached max_length
            if not terminated and len(self._sequence) >= self.max_length:
                truncated = True
                reward += self._compute_terminal_reward()

        elif action == ACTION_STOP:
            # ── STOP: score and terminate ──────────────────────────────
            terminated = True
            if len(self._sequence) < _MIN_LENGTH:
                reward = -1.0  # too short to be a functional AMP
            elif self._sequence:
                reward = self._compute_terminal_reward()

        obs = self._get_obs()
        info = self._get_info()
        info["activity_score"] = self.scorer.score(self._get_sequence_str()).get(
            "activity_score", 0.0
        ) if self._sequence else 0.0

        return obs, float(reward), terminated, truncated, info

    def _get_info(self) -> dict:
        base = super()._get_info()
        base["activity_score"] = 0.0  # overwritten in step()
        return base
