"""Heuristic baseline agents for each PeptideGym environment family."""
from __future__ import annotations

from typing import Any

import numpy as np

from peptidegym.peptide.properties import AA_TO_INDEX, AMINO_ACIDS


class HeuristicAMPAgent:
    """Biases toward cationic + hydrophobic residues in alternating pattern.

    Designed for antimicrobial peptide environments.
    """

    def __init__(self, env: Any, seed: int | None = None) -> None:
        self.env = env
        self.action_space = env.action_space
        self._rng = np.random.RandomState(seed)
        self._step = 0

    def predict(
        self, obs: Any, deterministic: bool = False
    ) -> tuple[int, None]:
        # Build probability distribution over 20 AAs + STOP (action 20)
        n_actions = self.action_space.n
        probs = np.zeros(n_actions, dtype=np.float64)

        # Preferred residues
        preferred = {
            "K": 0.15, "R": 0.15, "L": 0.15, "A": 0.10,
            "V": 0.10, "I": 0.10,
        }
        preferred_total = sum(preferred.values())  # 0.75
        other_per_aa = (1.0 - preferred_total) / max(1, 20 - len(preferred))  # ~0.025

        for aa, idx in AA_TO_INDEX.items():
            if aa in preferred:
                probs[idx] = preferred[aa]
            else:
                probs[idx] = other_per_aa

        # STOP action probability based on sequence length
        # p_stop = max(0, (len - 15) / 10)
        p_stop = max(0.0, (self._step - 15) / 10.0)
        p_stop = min(p_stop, 0.9)  # cap at 0.9

        if n_actions > 20:
            # Action 20 is STOP
            probs[20] = p_stop
            # Renormalize AA probabilities
            aa_total = probs[:20].sum()
            if aa_total > 0:
                probs[:20] *= (1.0 - p_stop) / aa_total

        probs /= probs.sum()

        if deterministic:
            action = int(np.argmax(probs))
        else:
            action = int(self._rng.choice(n_actions, p=probs))

        self._step += 1
        return action, None

    def learn(self, total_timesteps: int) -> None:
        pass

    def reset(self) -> None:
        self._step = 0


class HeuristicCyclicAgent:
    """Build 8-12 residue sequences with Cys placement for disulfide cyclization.

    Designed for cyclic peptide environments.
    """

    def __init__(self, env: Any, seed: int | None = None) -> None:
        self.env = env
        self.action_space = env.action_space
        self._rng = np.random.RandomState(seed)
        self._step = 0
        self._cys_count = 0
        self._target_length = 10  # will be randomized on reset

    def predict(
        self, obs: Any, deterministic: bool = False
    ) -> tuple[int, None]:
        n_actions = self.action_space.n
        pos = self._step  # 0-indexed current position

        # At positions 3 and target_length-2: always select Cys
        cys_idx = AA_TO_INDEX["C"]  # 4
        if pos == 3 or pos == self._target_length - 2:
            self._step += 1
            self._cys_count += 1
            return cys_idx, None

        # After 8+ residues with 2+ Cys: CYCLIZE_DISULFIDE (action 21) with p=0.3
        if pos >= 8 and self._cys_count >= 2 and n_actions > 21:
            if self._rng.random() < 0.3:
                return 21, None

        # Prefer hydrophobic residues
        probs = np.zeros(n_actions, dtype=np.float64)
        hydrophobic = {"L": 0.15, "I": 0.15, "V": 0.15, "F": 0.15, "A": 0.15}
        other_total = 1.0 - sum(hydrophobic.values())
        other_per_aa = other_total / max(1, 20 - len(hydrophobic))

        for aa, idx in AA_TO_INDEX.items():
            if aa in hydrophobic:
                probs[idx] = hydrophobic[aa]
            else:
                probs[idx] = other_per_aa

        # Zero out non-AA actions for building phase
        for i in range(20, n_actions):
            probs[i] = 0.0

        probs /= probs.sum()

        if deterministic:
            action = int(np.argmax(probs))
        else:
            action = int(self._rng.choice(n_actions, p=probs))

        if action < 20:
            if AMINO_ACIDS[action] == "C":
                self._cys_count += 1
            self._step += 1

        return action, None

    def learn(self, total_timesteps: int) -> None:
        pass

    def reset(self) -> None:
        self._step = 0
        self._cys_count = 0
        self._target_length = int(self._rng.randint(8, 13))  # 8-12


class HeuristicEpitopeAgent:
    """Place HLA-A*02:01 anchors for 9-mer epitope design.

    Designed for epitope/neoantigen environments.
    """

    def __init__(self, env: Any, seed: int | None = None) -> None:
        self.env = env
        self.action_space = env.action_space
        self._rng = np.random.RandomState(seed)
        self._step = 0

    def predict(
        self, obs: Any, deterministic: bool = False
    ) -> tuple[int, None]:
        n_actions = self.action_space.n
        pos = self._step  # 0-indexed

        # Position 1 (P2 anchor): L or M
        if pos == 1:
            aa = "L" if self._rng.random() < 0.6 else "M"
            self._step += 1
            return AA_TO_INDEX[aa], None

        # Position 8 (C-terminal for 9-mer): V or L
        if pos == 8:
            aa = "V" if self._rng.random() < 0.6 else "L"
            self._step += 1
            return AA_TO_INDEX[aa], None

        # STOP logic: at position 8 (after placing it) p=0.7, position 9 p=0.9
        if pos >= 9 and n_actions > 20:
            # STOP action is 20
            if pos == 9:
                if self._rng.random() < 0.7:
                    return 20, None
            elif pos == 10:
                if self._rng.random() < 0.9:
                    return 20, None
            else:
                return 20, None

        # TCR-facing positions (3-6, 0-indexed): prefer hydrophobic
        tcr_hydrophobic = {"F": 0.20, "Y": 0.20, "W": 0.15, "L": 0.15, "I": 0.15}
        if 3 <= pos <= 6:
            probs = np.zeros(n_actions, dtype=np.float64)
            other_total = 1.0 - sum(tcr_hydrophobic.values())
            other_per_aa = max(0.0, other_total / max(1, 20 - len(tcr_hydrophobic)))
            for aa, idx in AA_TO_INDEX.items():
                if aa in tcr_hydrophobic:
                    probs[idx] = tcr_hydrophobic[aa]
                else:
                    probs[idx] = other_per_aa
            for i in range(20, n_actions):
                probs[i] = 0.0
            probs /= probs.sum()
        else:
            # Uniform over amino acids
            probs = np.zeros(n_actions, dtype=np.float64)
            for i in range(20):
                probs[i] = 1.0 / 20.0
            for i in range(20, n_actions):
                probs[i] = 0.0

        if deterministic:
            action = int(np.argmax(probs))
        else:
            action = int(self._rng.choice(n_actions, p=probs))

        self._step += 1
        return action, None

    def learn(self, total_timesteps: int) -> None:
        pass

    def reset(self) -> None:
        self._step = 0
