"""EpitopeEnv — Gymnasium environment for MHC-I epitope 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,
    AA_TO_INDEX,
    INDEX_TO_AA,
    one_hot_encode,
)
from peptidegym.peptide.epitope_scoring import HeuristicEpitopeScorer
from peptidegym.envs.base import PeptideEnv

# Action constants
_NUM_AAS = len(AMINO_ACIDS)  # 20
ACTION_STOP = _NUM_AAS       # 20

# Fixed epitope length constraints (MHC-I typically 8-11mers)
_MIN_LENGTH = 8
_MAX_LENGTH = 11

# ── HLA pseudo-sequence encodings ─────────────────────────────────
# 34 MHC contact residues encoded as normalised floats.
# HLA-A*02:01 is the most common allele in global populations.
_HLA_ENCODINGS: dict[str, np.ndarray] = {
    "HLA-A*02:01": np.array([
        0.72, 0.31, 0.85, 0.44, 0.19, 0.67, 0.53, 0.91, 0.28, 0.76,
        0.62, 0.38, 0.84, 0.15, 0.49, 0.70, 0.22, 0.88, 0.33, 0.57,
        0.41, 0.95, 0.18, 0.63, 0.79, 0.26, 0.54, 0.87, 0.45, 0.11,
        0.68, 0.36, 0.92, 0.50,
    ], dtype=np.float32),
    "HLA-A*01:01": np.array([
        0.65, 0.42, 0.78, 0.33, 0.21, 0.59, 0.47, 0.86, 0.35, 0.71,
        0.58, 0.44, 0.80, 0.19, 0.55, 0.64, 0.28, 0.82, 0.39, 0.51,
        0.37, 0.90, 0.24, 0.69, 0.73, 0.32, 0.48, 0.83, 0.41, 0.16,
        0.62, 0.30, 0.88, 0.46,
    ], dtype=np.float32),
    "HLA-B*07:02": np.array([
        0.58, 0.49, 0.71, 0.27, 0.35, 0.63, 0.40, 0.79, 0.31, 0.66,
        0.52, 0.46, 0.75, 0.22, 0.60, 0.57, 0.34, 0.77, 0.43, 0.48,
        0.33, 0.85, 0.29, 0.64, 0.68, 0.38, 0.45, 0.80, 0.37, 0.20,
        0.56, 0.41, 0.84, 0.53,
    ], dtype=np.float32),
    "HLA-B*44:02": np.array([
        0.61, 0.38, 0.74, 0.30, 0.28, 0.66, 0.43, 0.82, 0.33, 0.69,
        0.55, 0.41, 0.77, 0.17, 0.52, 0.61, 0.31, 0.80, 0.36, 0.54,
        0.39, 0.88, 0.22, 0.67, 0.71, 0.35, 0.50, 0.81, 0.40, 0.14,
        0.59, 0.34, 0.86, 0.49,
    ], dtype=np.float32),
    "HLA-A*03:01": np.array([
        0.69, 0.35, 0.81, 0.40, 0.23, 0.61, 0.50, 0.88, 0.30, 0.74,
        0.60, 0.36, 0.82, 0.18, 0.47, 0.68, 0.25, 0.85, 0.37, 0.55,
        0.43, 0.93, 0.20, 0.65, 0.77, 0.29, 0.52, 0.84, 0.43, 0.13,
        0.66, 0.33, 0.90, 0.48,
    ], dtype=np.float32),
}

# Known MHC-I anchor preferences for HLA-A*02:01 (positions P2, P9/C-terminal)
# Preferred anchor residues at P2 (index 1) and C-terminal
_ANCHOR_P2_PREFERRED = {AA_TO_INDEX[aa] for aa in "LMI" if aa in AA_TO_INDEX}
_ANCHOR_CTER_PREFERRED = {AA_TO_INDEX[aa] for aa in "VLI" if aa in AA_TO_INDEX}


class EpitopeEnv(PeptideEnv):
    """Design MHC-I epitope peptides for vaccine or immunotherapy targets.

    Actions 0-19: append amino acid.
    Action 20 (STOP): terminate and score peptide.
    Peptide length is fixed between 8 and 11 residues.
    """

    def __init__(
        self,
        max_length: int = _MAX_LENGTH,
        difficulty: str = "medium",
        hla_allele: str = "HLA-A*02:01",
        seed: int | None = None,
    ) -> None:
        self.difficulty = difficulty
        self.hla_allele = hla_allele

        # In hard mode, we score across 5 HLA supertypes
        self._alleles: list[str] = (
            list(_HLA_ENCODINGS.keys())
            if difficulty == "hard"
            else [hla_allele]
        )

        reward_weights = {
            "binding_score": 0.45,
            "self_dissimilarity": 0.25,
            "immunogenicity_proxy": 0.20,
            "length_fitness": 0.10,
        }

        scorer = HeuristicEpitopeScorer()

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

        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,
                ),
                "hla_encoding": spaces.Box(
                    low=0.0, high=1.0,
                    shape=(34,),
                    dtype=np.float32,
                ),
                "binding_estimate": spaces.Box(
                    low=-np.inf, high=np.inf,
                    shape=(1,),
                    dtype=np.float32,
                ),
            }
        )

    # ── Observation helpers ───────────────────────────────────────────

    def _get_hla_encoding(self) -> np.ndarray:
        """Return the 34-dim pseudo-sequence encoding for the active allele."""
        return _HLA_ENCODINGS.get(
            self.hla_allele,
            _HLA_ENCODINGS["HLA-A*02:01"],
        ).copy()

    def _get_binding_estimate(self) -> np.ndarray:
        """Rough binding estimate from the scorer for the current partial sequence."""
        if len(self._sequence) < 2:
            return np.zeros(1, dtype=np.float32)
        seq = self._get_sequence_str()
        scores = self.scorer.score(seq, hla_allele=self.hla_allele)
        return np.array([scores.get("binding_score", 0.0)], dtype=np.float32)

    def _get_obs(self) -> dict:
        if self._sequence:
            seq_onehot = one_hot_encode(self._get_sequence_str(), _MAX_LENGTH)
        else:
            seq_onehot = np.zeros((_MAX_LENGTH, _NUM_AAS), dtype=np.float32)

        return {
            "sequence_onehot": seq_onehot,
            "position": np.array([len(self._sequence)], dtype=np.float32),
            "hla_encoding": self._get_hla_encoding(),
            "binding_estimate": self._get_binding_estimate(),
        }

    # ── 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 ─────────────────────────────────────
            self._sequence.append(action)
            self._step_count += 1

            pos = len(self._sequence)

            # Anchor shaping: preferred residues at P2 (pos 2) or C-terminal
            if pos == 2 and action in _ANCHOR_P2_PREFERRED:
                reward += 0.03
            if pos >= _MIN_LENGTH and action in _ANCHOR_CTER_PREFERRED:
                reward += 0.03

            # Auto-stop at max epitope length
            if pos >= _MAX_LENGTH:
                terminated = True
                reward += self._compute_terminal_reward()

        elif action == ACTION_STOP:
            terminated = True
            if len(self._sequence) < _MIN_LENGTH:
                reward = -1.0
            else:
                reward = self._compute_terminal_reward()

        obs = self._get_obs()
        info = self._get_info()
        return obs, float(reward), terminated, truncated, info

    def _compute_terminal_reward(self) -> float:
        """Score across one or more HLA alleles depending on difficulty."""
        seq = self._get_sequence_str()
        if not seq:
            return 0.0

        if self.difficulty == "hard":
            # Average binding across 5 supertypes
            all_rewards: list[float] = []
            for allele in self._alleles:
                scores = self.scorer.score(seq, hla_allele=allele)
                r = sum(
                    self.reward_weights.get(k, 0.0) * scores.get(k, 0.0)
                    for k in self.reward_weights
                )
                all_rewards.append(r)
            return float(np.mean(all_rewards))

        scores = self.scorer.score(seq, hla_allele=self.hla_allele)

        # Easy mode: skip self_dissimilarity component
        if self.difficulty == "easy":
            weights = {k: v for k, v in self.reward_weights.items() if k != "self_dissimilarity"}
            # Redistribute weight
            total = sum(weights.values())
            if total > 0:
                weights = {k: v / total for k, v in weights.items()}
        else:
            weights = self.reward_weights

        reward = sum(weights.get(k, 0.0) * scores.get(k, 0.0) for k in weights)
        return float(reward)

    def _get_info(self) -> dict:
        base = super()._get_info()
        base["hla_allele"] = self.hla_allele
        if self._sequence:
            scores = self.scorer.score(
                self._get_sequence_str(), hla_allele=self.hla_allele
            )
            base["binding_score"] = scores.get("binding_score", 0.0)
        else:
            base["binding_score"] = 0.0
        return base
