"""Base PeptideEnv — abstract Gymnasium environment for peptide design."""
from __future__ import annotations

from abc import abstractmethod
from typing import Any

import gymnasium
import numpy as np
from gymnasium import spaces

from peptidegym.peptide.properties import (
    AMINO_ACIDS,
    AA_TO_INDEX,
    INDEX_TO_AA,
    net_charge,
    mean_hydrophobicity,
    aromaticity_fraction,
    instability_index,
    molecular_weight,
    amphipathicity_score,
    one_hot_encode,
)
from peptidegym.peptide.scoring import RewardBackend


class PeptideEnv(gymnasium.Env):
    """Abstract base class for peptide-design RL environments.

    Subclasses must set ``action_space``, ``observation_space``, and implement
    ``step()`` to handle special actions beyond amino-acid appending.
    """

    metadata = {"render_modes": ["human"]}

    def __init__(
        self,
        max_length: int,
        scorer: RewardBackend,
        reward_weights: dict[str, float],
        seed: int | None = None,
    ) -> None:
        super().__init__()
        self.max_length = max_length
        self.scorer = scorer
        self.reward_weights = reward_weights

        # Internal state
        self._sequence: list[int] = []
        self._step_count: int = 0

        # Seed through gymnasium's API — actual spaces set by subclass
        if seed is not None:
            self._np_random, _ = gymnasium.utils.seeding.np_random(seed)

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

    def reset(
        self,
        *,
        seed: int | None = None,
        options: dict[str, Any] | None = None,
    ) -> tuple[dict, dict]:
        super().reset(seed=seed)
        self._sequence = []
        self._step_count = 0
        return self._get_obs(), self._get_info()

    @abstractmethod
    def step(self, action: int) -> tuple[dict, float, bool, bool, dict]:
        """Execute *action* and return (obs, reward, terminated, truncated, info)."""
        ...

    def render(self) -> str:
        """Return the current peptide sequence as a string."""
        return self._get_sequence_str()

    # ── Helpers ───────────────────────────────────────────────────────

    def _get_obs(self) -> dict:
        """Build the observation dictionary for the current state."""
        if self._sequence:
            seq_onehot = one_hot_encode(self._get_sequence_str(), self.max_length)
        else:
            seq_onehot = np.zeros(
                (self.max_length, len(AMINO_ACIDS)), dtype=np.float32
            )

        position = np.array([len(self._sequence)], dtype=np.float32)
        properties = self._compute_properties()

        return {
            "sequence_onehot": seq_onehot,
            "position": position,
            "properties": properties,
        }

    def _get_sequence_str(self) -> str:
        """Convert the internal index list to a one-letter amino-acid string."""
        return "".join(INDEX_TO_AA[int(i)] for i in self._sequence)

    def _compute_properties(self) -> np.ndarray:
        """Compute a property vector for the current sequence.

        Default: [MW, net_charge, mean_hydro, amphipathicity,
                  aromaticity, instability]
        """
        if not self._sequence:
            return np.zeros(6, dtype=np.float32)
        seq = self._get_sequence_str()
        return np.array(
            [
                molecular_weight(seq),
                net_charge(seq),
                mean_hydrophobicity(seq),
                amphipathicity_score(seq),
                aromaticity_fraction(seq),
                instability_index(seq),
            ],
            dtype=np.float32,
        )

    def _compute_terminal_reward(self) -> float:
        """Score the completed peptide with the scorer and apply weights."""
        seq = self._get_sequence_str()
        scores: dict[str, float] = self.scorer.score(seq)
        reward = 0.0
        for key, weight in self.reward_weights.items():
            reward += weight * scores.get(key, 0.0)
        return float(reward)

    def _get_info(self) -> dict:
        """Return informational dict for the current state."""
        return {
            "sequence": self._get_sequence_str(),
            "length": len(self._sequence),
        }
