"""CyclicPeptideEnv — Gymnasium environment for cyclic 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,
    MOLECULAR_WEIGHT,
    CHARGE_PH74,
    HYDROPHOBICITY,
    molecular_weight,
)
from peptidegym.peptide.cyclic_scoring import HeuristicCyclicScorer
from peptidegym.envs.base import PeptideEnv

# Action constants
_NUM_AAS = len(AMINO_ACIDS)                  # 20
ACTION_CYCLIZE_HEAD_TO_TAIL = _NUM_AAS       # 20
ACTION_CYCLIZE_DISULFIDE = _NUM_AAS + 1      # 21
ACTION_CYCLIZE_SIDECHAIN = _NUM_AAS + 2      # 22
ACTION_STOP_LINEAR = _NUM_AAS + 3            # 23


class CyclicPeptideEnv(PeptideEnv):
    """Design cyclic peptides via sequential residue placement and cyclization.

    Actions 0-19: append amino acid.
    Action 20: cyclize head-to-tail.
    Action 21: cyclize via disulfide bridge.
    Action 22: cyclize via sidechain lactam.
    Action 23: stop as linear peptide (no cyclization).
    """

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

        reward_weights = {
            "cyclization_bonus": 0.30,
            "permeability_score": 0.25,
            "structural_score": 0.25,
            "synthetic_score": 0.20,
        }

        scorer = HeuristicCyclicScorer()

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

        self.action_space = spaces.Discrete(ACTION_STOP_LINEAR + 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=(5,),
                    dtype=np.float32,
                ),
                "cyclization_valid": spaces.Box(
                    low=0.0, high=1.0,
                    shape=(3,),
                    dtype=np.float32,
                ),
            }
        )

    # ── Property and observation helpers ──────────────────────────────

    def _compute_properties(self) -> np.ndarray:
        """[MW, n_hbond_donors, logP_estimate, backbone_flexibility, cys_count]."""
        if not self._sequence:
            return np.zeros(5, dtype=np.float32)

        seq = self._get_sequence_str()
        mw = molecular_weight(seq)

        # Rough H-bond donor count: backbone NH (one per residue except first)
        # plus sidechain donors for S,T,Y,N,Q,K,R,H,W
        hbond_sidechain = set("STYNQKRHW")
        n_hbond = float(max(0, len(seq) - 1))
        n_hbond += sum(1.0 for aa in seq if aa in hbond_sidechain)

        # Crude logP estimate: sum of AA hydrophobicity / length
        logp_est = sum(HYDROPHOBICITY.get(aa, 0.0) for aa in seq) / max(len(seq), 1)

        # Backbone flexibility heuristic: G and P destabilise / restrict
        gly_pro_frac = sum(1.0 for aa in seq if aa in ("G", "P")) / max(len(seq), 1)
        backbone_flex = 1.0 - gly_pro_frac  # higher = more flexible

        cys_count = float(seq.count("C"))

        return np.array(
            [mw, n_hbond, logp_est, backbone_flex, cys_count],
            dtype=np.float32,
        )

    def _compute_cyclization_valid(self) -> np.ndarray:
        """Return a 3-element binary mask: [head_to_tail_ok, disulfide_ok, lactam_ok]."""
        seq = self._get_sequence_str()
        head_to_tail_ok = float(len(seq) >= 4)
        disulfide_ok = float(seq.count("C") >= 2)
        lactam_ok = float(
            ("K" in seq) and any(aa in seq for aa in ("D", "E"))
        )

        # Difficulty gating
        if self.difficulty == "easy":
            disulfide_ok = 0.0
            lactam_ok = 0.0
        elif self.difficulty == "medium":
            lactam_ok = 0.0

        return np.array(
            [head_to_tail_ok, disulfide_ok, lactam_ok],
            dtype=np.float32,
        )

    def _get_obs(self) -> dict:
        obs = super()._get_obs()
        obs["cyclization_valid"] = self._compute_cyclization_valid()
        return obs

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

    def reset(
        self,
        *,
        seed: int | None = None,
        options: dict[str, Any] | None = None,
    ) -> tuple[dict, dict]:
        self._cyclization_type = "none"
        return super().reset(seed=seed, options=options)

    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

            if len(self._sequence) >= self.max_length:
                truncated = True
                self._cyclization_type = "none"
                reward = self._compute_terminal_reward()

        elif action == ACTION_CYCLIZE_HEAD_TO_TAIL:
            valid = self._compute_cyclization_valid()
            if valid[0] > 0.5:
                self._cyclization_type = "head_to_tail"
                terminated = True
                reward = self._compute_terminal_reward()
            else:
                terminated = True
                reward = -0.5

        elif action == ACTION_CYCLIZE_DISULFIDE:
            valid = self._compute_cyclization_valid()
            if valid[1] > 0.5:
                self._cyclization_type = "disulfide"
                terminated = True
                reward = self._compute_terminal_reward()
            else:
                terminated = True
                reward = -0.5

        elif action == ACTION_CYCLIZE_SIDECHAIN:
            valid = self._compute_cyclization_valid()
            if valid[2] > 0.5:
                self._cyclization_type = "sidechain"
                terminated = True
                reward = self._compute_terminal_reward()
            else:
                terminated = True
                reward = -0.5

        elif action == ACTION_STOP_LINEAR:
            self._cyclization_type = "none"
            terminated = True
            reward = self._compute_terminal_reward() if self._sequence else 0.0

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

    def _compute_terminal_reward(self) -> float:
        """Score with cyclization context passed to the scorer."""
        seq = self._get_sequence_str()
        scores: dict[str, float] = self.scorer.score(
            seq, cyclization_type=self._cyclization_type
        )
        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:
        base = super()._get_info()
        base["cyclization_type"] = self._cyclization_type
        base["ring_size"] = len(self._sequence) if self._cyclization_type != "none" else 0
        return base
