"""
Module C: RL Intervention Policy — Actor-Critic network for PPO.

Observation vector (flat, 2*H+17 dim, e.g. 2065 for detector_hidden=1024):
  [d_score(1) | l_risk_onehot(5) | c_primary_probs(10) |
   e_H_pool(H) | e_P_pool(H) | t_norm(1)]

The network first parses this flat vector through _encode_obs() →
StateEncoder → 256-dim latent, then feeds the latent to actor/critic heads.

Action: {PASS=0, WARN=1, REWRITE=2, REJECT=3, CRISIS=4}
"""

import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import Tuple, Dict

from src.utils.taxonomy import NUM_ACTIONS, NUM_PRIMARY, NUM_RISK_LEVELS


class StateEncoder(nn.Module):
    """
    Encodes the structured state components for the RL policy.

    Input components (passed separately, not as flat vector):
      d_score       : [B, 1]
      l_risk        : [B]   — integer level index (0-4), embedded
      c_primary_probs: [B, NUM_PRIMARY]
      e_H_pool      : [B, detector_hidden]
      e_P_pool      : [B, detector_hidden]
      t_norm        : [B, 1]
    Output: [B, state_hidden]
    """

    def __init__(
        self,
        detector_hidden: int = 1024,
        level_emb_dim: int = 16,
        state_hidden: int = 256,
        dropout: float = 0.1,
    ):
        super().__init__()
        self.level_emb = nn.Embedding(NUM_RISK_LEVELS, level_emb_dim)

        # d_score(1) + level_emb(16) + c_primary_probs(10) + e_H(H) + e_P(H) + t_norm(1)
        state_dim = 1 + level_emb_dim + NUM_PRIMARY + detector_hidden * 2 + 1

        self.mlp = nn.Sequential(
            nn.Linear(state_dim, state_hidden * 2),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(state_hidden * 2, state_hidden),
            nn.ReLU(),
        )
        self.out_dim = state_hidden

    def forward(
        self,
        d_score: torch.Tensor,           # [B, 1]
        l_risk: torch.Tensor,            # [B]   — integer
        c_primary_probs: torch.Tensor,   # [B, NUM_PRIMARY]
        e_H_pool: torch.Tensor,          # [B, H]
        e_P_pool: torch.Tensor,          # [B, H]
        t_norm: torch.Tensor,            # [B, 1]
    ) -> torch.Tensor:
        level_emb = self.level_emb(l_risk)   # [B, 16]
        state = torch.cat(
            [d_score, level_emb, c_primary_probs, e_H_pool, e_P_pool, t_norm], dim=-1
        )
        return self.mlp(state)               # [B, state_hidden]


class InterventionAgent(nn.Module):
    """
    Actor-Critic network for PPO-based intervention policy.

    Actor:  π(a | s) = softmax(MLP(encoded_state))
    Critic: V(s)     = MLP(encoded_state)

    All public methods (get_action, evaluate_actions, behavior_clone_loss, forward)
    accept the raw flat observation vector (2*H+17 dim) — _encode_obs() is called
    internally to parse and encode it before the actor/critic heads.
    """

    def __init__(
        self,
        detector_hidden: int = 1024,
        state_hidden: int = 256,
        dropout: float = 0.1,
    ):
        super().__init__()
        self.detector_hidden = detector_hidden

        self.state_encoder = StateEncoder(
            detector_hidden=detector_hidden,
            state_hidden=state_hidden,
            dropout=dropout,
        )

        self.actor = nn.Sequential(
            nn.Linear(state_hidden, state_hidden),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(state_hidden, NUM_ACTIONS),
        )

        self.critic = nn.Sequential(
            nn.Linear(state_hidden, state_hidden),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(state_hidden, 1),
        )

    # ── obs parsing ────────────────────────────────────────────────────────

    def _encode_obs(self, obs: torch.Tensor) -> torch.Tensor:
        """
        Parse a flat observation vector and encode it through StateEncoder.

        Observation layout (2*H+17 dim):
          obs[:, 0]                                       → d_score
          obs[:, 1 : 1+NUM_RISK_LEVELS]                  → l_risk one-hot (→ argmax for embedding)
          obs[:, 1+NUM_RISK_LEVELS :
                 1+NUM_RISK_LEVELS+NUM_PRIMARY]           → c_primary_probs
          obs[:, 1+NUM_RISK_LEVELS+NUM_PRIMARY :
                 1+NUM_RISK_LEVELS+NUM_PRIMARY+H]         → e_H_pool
          obs[:, 1+NUM_RISK_LEVELS+NUM_PRIMARY+H :
                 1+NUM_RISK_LEVELS+NUM_PRIMARY+2*H]       → e_P_pool
          obs[:, -1]                                      → t_norm

        Args:
            obs: [B, 2*H+17] float tensor on any device

        Returns:
            encoded: [B, state_hidden]
        """
        H = self.detector_hidden
        ptr = 0

        d_score       = obs[:, ptr : ptr + 1];                     ptr += 1
        l_risk_onehot = obs[:, ptr : ptr + NUM_RISK_LEVELS];        ptr += NUM_RISK_LEVELS
        c_primary     = obs[:, ptr : ptr + NUM_PRIMARY];            ptr += NUM_PRIMARY
        e_H_pool      = obs[:, ptr : ptr + H];                     ptr += H
        e_P_pool      = obs[:, ptr : ptr + H];                     ptr += H
        t_norm        = obs[:, ptr : ptr + 1]

        # Convert one-hot → integer index for the embedding lookup
        l_risk_idx = l_risk_onehot.argmax(dim=-1)   # [B]

        return self.state_encoder(d_score, l_risk_idx, c_primary, e_H_pool, e_P_pool, t_norm)

    # ── public API ─────────────────────────────────────────────────────────

    def forward(self, obs: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
        """
        Args:
            obs: raw flat observation [B, 2*H+17]
        Returns:
            (action_logits [B, NUM_ACTIONS], state_value [B, 1])
        """
        encoded = self._encode_obs(obs)
        return self.actor(encoded), self.critic(encoded)

    def encode_state(
        self,
        d_score: torch.Tensor,
        l_risk: torch.Tensor,
        c_primary_probs: torch.Tensor,
        e_H_pool: torch.Tensor,
        e_P_pool: torch.Tensor,
        t_norm: torch.Tensor,
    ) -> torch.Tensor:
        """Direct StateEncoder call with pre-parsed components (kept for external use)."""
        return self.state_encoder(d_score, l_risk, c_primary_probs, e_H_pool, e_P_pool, t_norm)

    def get_action(
        self, obs: torch.Tensor, deterministic: bool = False
    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
        """
        Sample (or argmax) an action from the policy.

        Args:
            obs: [B, 2*H+17] raw observation
        Returns:
            (action [B], log_prob [B], entropy [B], value [B])
        """
        logits, value = self.forward(obs)
        dist = torch.distributions.Categorical(logits=logits)
        action = logits.argmax(-1) if deterministic else dist.sample()
        return action, dist.log_prob(action), dist.entropy(), value.squeeze(-1)

    def evaluate_actions(
        self, obs: torch.Tensor, actions: torch.Tensor
    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        """
        Re-evaluate stored actions for PPO update.

        Args:
            obs:     [B, 2*H+17] raw observation
            actions: [B] integer action indices
        Returns:
            (log_prob [B], entropy [B], value [B])
        """
        logits, value = self.forward(obs)
        dist = torch.distributions.Categorical(logits=logits)
        return dist.log_prob(actions), dist.entropy(), value.squeeze(-1)

    def behavior_clone_loss(
        self, obs: torch.Tensor, expert_actions: torch.Tensor
    ) -> torch.Tensor:
        """
        Supervised cross-entropy loss for BC warm-up.

        Args:
            obs:            [B, 2*H+17] raw observation
            expert_actions: [B] integer expert action indices
        """
        logits, _ = self.forward(obs)
        return F.cross_entropy(logits, expert_actions)
