"""
Learned Card Embeddings for Essence Wars.

This module provides infrastructure for replacing raw card IDs in the
observation tensor with learned embedding vectors. Embeddings are learned
end-to-end jointly with the policy network.

The key insight is that raw card IDs (e.g., 1000, 1001) are numerically close
but may represent very different cards. Embeddings allow the network to learn
semantic relationships between cards.

Usage:
    network = EmbeddedPPONetwork(embed_dim=64)
"""

from __future__ import annotations

from dataclasses import dataclass
from typing import Literal, cast

import torch
import torch.nn as nn
from torch import Tensor
from torch.distributions import Categorical

from essence_wars._core import STATE_TENSOR_SIZE
from essence_wars.agents.networks import ResidualBlock, init_network_weights

# =============================================================================
# Tensor Structure Constants
# =============================================================================

# Global state section
GLOBAL_START = 0
GLOBAL_SIZE = 6

# Player state structure (75 floats each)
PLAYER_STATE_SIZE = 75

# Player 1 section
P1_START = 6
P1_BASE_STATS_SIZE = 5
P1_HAND_START = 11  # P1_START + P1_BASE_STATS_SIZE
P1_HAND_SIZE = 10
P1_CREATURES_START = 21  # P1_HAND_START + P1_HAND_SIZE
P1_CREATURES_SIZE = 50  # 5 slots x 10 floats
P1_SUPPORTS_START = 71  # P1_CREATURES_START + P1_CREATURES_SIZE
P1_SUPPORTS_SIZE = 10  # 2 slots x 5 floats

# Player 2 section
P2_START = 81  # P1_START + PLAYER_STATE_SIZE
P2_HAND_START = 86  # P2_START + P1_BASE_STATS_SIZE
P2_CREATURES_START = 96
P2_SUPPORTS_START = 146

# Card embedding section (trailing)
EMBED_SECTION_START = 156
EMBED_SECTION_SIZE = STATE_TENSOR_SIZE - EMBED_SECTION_START  # Commander IDs at 326-327

# Slot sizes
CREATURE_SLOT_SIZE = 10
SUPPORT_SLOT_SIZE = 5

# Card ID positions within slots
# Support slots: [occupied, durability, card_id, reserved, reserved]
SUPPORT_CARD_ID_OFFSET = 2

# Maximum card ID (card IDs range from 1000-4074, but we use 5000 for buffer)
MAX_CARD_ID = 5000


# =============================================================================
# Card ID Position Mapping
# =============================================================================

@dataclass
class CardIdPositions:
    """Positions of card IDs in the 328-float state tensor."""

    # Hand card IDs (10 per player)
    p1_hand: list[int]
    p2_hand: list[int]

    # Support card IDs (2 per player, within 5-float slots)
    p1_supports: list[int]
    p2_supports: list[int]

    # Card embedding section (variable length)
    embed_section: list[int]

    @classmethod
    def default(cls) -> CardIdPositions:
        """Create default card ID positions based on tensor structure."""
        return cls(
            p1_hand=list(range(P1_HAND_START, P1_HAND_START + P1_HAND_SIZE)),
            p2_hand=list(range(P2_HAND_START, P2_HAND_START + P1_HAND_SIZE)),
            p1_supports=[
                P1_SUPPORTS_START + SUPPORT_CARD_ID_OFFSET,  # Support slot 0
                P1_SUPPORTS_START + SUPPORT_SLOT_SIZE + SUPPORT_CARD_ID_OFFSET,  # Support slot 1
            ],
            p2_supports=[
                P2_SUPPORTS_START + SUPPORT_CARD_ID_OFFSET,
                P2_SUPPORTS_START + SUPPORT_SLOT_SIZE + SUPPORT_CARD_ID_OFFSET,
            ],
            embed_section=list(range(EMBED_SECTION_START, STATE_TENSOR_SIZE)),
        )

    def all_positions(self) -> list[int]:
        """Get all card ID positions as a flat list."""
        return (
            self.p1_hand +
            self.p2_hand +
            self.p1_supports +
            self.p2_supports +
            self.embed_section
        )

    def non_embed_section_positions(self) -> list[int]:
        """Get card ID positions excluding the trailing embed section."""
        return (
            self.p1_hand +
            self.p2_hand +
            self.p1_supports +
            self.p2_supports
        )


def get_non_card_positions(card_positions: CardIdPositions) -> list[int]:
    """Get all tensor positions that are NOT card IDs."""
    all_card_positions = set(card_positions.all_positions())
    return [i for i in range(STATE_TENSOR_SIZE) if i not in all_card_positions]


# =============================================================================
# Observation Transformer
# =============================================================================

class ObservationTransformer(nn.Module):
    """
    Transform flat observations by replacing card IDs with learned embeddings.

    This module extracts card IDs from specific positions in the state tensor,
    looks up their embeddings, and concatenates them with the non-card features.

    Args:
        card_embedding: The embedding layer to use for card IDs
        include_embed_section: Whether to embed the trailing card ID section

    Example:
        card_embed = nn.Embedding(5000, 64)
        transformer = ObservationTransformer(card_embed)

        obs = torch.randn(32, 328)  # Batch of flat observations
        embedded = transformer(obs)  # Shape: (32, new_dim)
    """

    def __init__(
        self,
        card_embedding: nn.Embedding,
        include_embed_section: bool = True,
    ):
        super().__init__()

        self.card_embedding = card_embedding
        self.include_embed_section = include_embed_section
        self.embed_dim = card_embedding.embedding_dim

        # Get card ID positions
        self.positions = CardIdPositions.default()

        if include_embed_section:
            self.card_id_positions = self.positions.all_positions()
        else:
            self.card_id_positions = self.positions.non_embed_section_positions()

        # Get non-card positions (excludes ALL card ID positions including embed section)
        self.non_card_positions = get_non_card_positions(self.positions)
        # NOTE: When include_embed_section=False, we do NOT add embed section back
        # as raw features. Those positions contain card IDs (1000-4074) which would
        # dominate the input and break learning. Instead, we only embed the 24 card
        # IDs from hands and supports, and exclude the embed section entirely.

        # Pre-compute dimensions
        self.num_card_slots = len(self.card_id_positions)
        self.num_non_card_features = len(self.non_card_positions)
        self.output_dim = self.num_non_card_features + self.num_card_slots * self.embed_dim

        # Register positions as buffers for device handling
        self.register_buffer(
            'card_id_indices',
            torch.tensor(self.card_id_positions, dtype=torch.long)
        )
        self.register_buffer(
            'non_card_indices',
            torch.tensor(self.non_card_positions, dtype=torch.long)
        )

    def forward(self, obs: torch.Tensor) -> torch.Tensor:
        """
        Transform flat observation to embedded observation.

        Args:
            obs: Flat observation tensor of shape (batch, 328) or (328,)

        Returns:
            Embedded observation of shape (batch, output_dim) or (output_dim,)
        """
        squeeze = obs.dim() == 1
        if squeeze:
            obs = obs.unsqueeze(0)

        batch_size = obs.shape[0]

        # Extract non-card features
        non_card_idx = cast('Tensor', self.non_card_indices)
        non_card_features = obs[:, non_card_idx]  # (batch, num_non_card)

        # Extract card IDs and convert to long indices
        card_id_idx = cast('Tensor', self.card_id_indices)
        card_ids = obs[:, card_id_idx].long()  # (batch, num_cards)

        # Clamp card IDs to valid range (0 = padding, 1000-4074 = valid cards)
        card_ids = card_ids.clamp(0, MAX_CARD_ID - 1)

        # Look up embeddings
        card_embeds = self.card_embedding(card_ids)  # (batch, num_cards, embed_dim)

        # Flatten card embeddings
        card_embeds_flat = card_embeds.view(batch_size, -1)  # (batch, num_cards * embed_dim)

        # Concatenate
        embedded_obs = torch.cat([non_card_features, card_embeds_flat], dim=1)

        if squeeze:
            embedded_obs = embedded_obs.squeeze(0)

        return embedded_obs


# =============================================================================
# Embedding Factory
# =============================================================================


def create_card_embedding(
    num_cards: int = MAX_CARD_ID,
    embed_dim: int = 64,
    pretrained_embeds: torch.Tensor | None = None,
    freeze_embeds: bool = False,
) -> tuple[nn.Embedding, int]:
    """Create card embedding layer.

    Factory function for creating card embedding layers with optional
    pretrained weights.

    Args:
        num_cards: Maximum number of cards in vocabulary
        embed_dim: Embedding dimension (ignored if pretrained)
        pretrained_embeds: Optional pretrained embedding weights
        freeze_embeds: Whether to freeze embeddings during training

    Returns:
        (embedding_layer, actual_embed_dim) tuple
    """
    if pretrained_embeds is not None:
        return nn.Embedding.from_pretrained(
            pretrained_embeds,
            freeze=freeze_embeds,
            padding_idx=0,
        ), pretrained_embeds.shape[1]

    return nn.Embedding(num_cards, embed_dim, padding_idx=0), embed_dim


# =============================================================================
# Embedded Network Architectures
# =============================================================================


class EmbeddedPPONetwork(nn.Module):
    """
    PPO network with learned card embeddings.

    This network replaces raw card IDs in the observation with learned
    embedding vectors before processing through the policy-value network.

    Args:
        embed_dim: Dimension of card embeddings (default: 64)
        hidden_dim: Hidden layer dimension (default: 256)
        num_cards: Number of cards in embedding table (default: 5000)
        include_embed_section: Whether to embed the trailing card section
        pretrained_embeds: Optional pre-trained embedding weights
        freeze_embeds: Whether to freeze embeddings during training

    Example:
        # End-to-end learning
        network = EmbeddedPPONetwork(embed_dim=64)

        # Pre-trained embeddings
        pretrained = torch.load("card2vec_embeds.pt")
        network = EmbeddedPPONetwork(pretrained_embeds=pretrained)
    """

    def __init__(
        self,
        embed_dim: int = 64,
        hidden_dim: int = 256,
        num_cards: int = MAX_CARD_ID,
        include_embed_section: bool = False,  # Start with simpler version
        pretrained_embeds: torch.Tensor | None = None,
        freeze_embeds: bool = False,
    ):
        super().__init__()

        self.hidden_dim = hidden_dim
        self.num_cards = num_cards
        self.include_embed_section = include_embed_section
        self._using_pretrained_embeds = pretrained_embeds is not None

        # Create or load card embedding layer
        self.card_embedding, self.embed_dim = create_card_embedding(
            num_cards=num_cards,
            embed_dim=embed_dim,
            pretrained_embeds=pretrained_embeds,
            freeze_embeds=freeze_embeds,
        )

        # Create observation transformer
        self.obs_transformer = ObservationTransformer(
            self.card_embedding,
            include_embed_section=include_embed_section,
        )

        # Calculate input dimension after transformation
        input_dim = self.obs_transformer.output_dim

        # Shared trunk
        self.trunk = nn.Sequential(
            nn.Linear(input_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, hidden_dim),
            nn.ReLU(),
        )

        # Policy head
        self.policy_head = nn.Sequential(
            nn.Linear(hidden_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, 422),  # ACTION_SPACE_SIZE
        )

        # Value head
        self.value_head = nn.Sequential(
            nn.Linear(hidden_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, 1),
        )

        # Initialize weights
        self._init_weights()

    def _init_weights(self) -> None:
        """Initialize network weights using orthogonal initialization."""
        # Only initialize embeddings if not using pretrained
        card_embed = None if self._using_pretrained_embeds else self.card_embedding
        init_network_weights(self, policy_head=self.policy_head, card_embedding=card_embed)

    def forward(
        self,
        obs: torch.Tensor,
        action_mask: torch.Tensor | None = None,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """
        Forward pass through the network.

        Args:
            obs: Flat observations of shape (batch, 328) or (328,)
            action_mask: Boolean mask of shape (batch, 422) where True = legal

        Returns:
            logits: Action logits of shape (batch, 422)
            value: State values of shape (batch,)
        """
        # Transform observation (replace card IDs with embeddings)
        embedded_obs = self.obs_transformer(obs)

        # Process through trunk
        features = self.trunk(embedded_obs)

        # Get outputs
        logits = self.policy_head(features)
        value = self.value_head(features).squeeze(-1)

        # Apply action mask
        if action_mask is not None:
            logits = logits.masked_fill(~action_mask, -1e8)

        return logits, value

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

        Args:
            obs: Observations of shape (batch, 328)
            action_mask: Boolean mask of shape (batch, 422)
            deterministic: If True, return argmax action

        Returns:
            action: Sampled actions of shape (batch,)
            log_prob: Log probabilities of shape (batch,)
            entropy: Policy entropy of shape (batch,)
        """
        logits, _ = self.forward(obs, action_mask)
        dist = Categorical(logits=logits)

        action = logits.argmax(dim=-1) if deterministic else dist.sample()

        log_prob = dist.log_prob(action)
        entropy = dist.entropy()

        return action, log_prob, entropy

    def get_value(self, obs: torch.Tensor) -> torch.Tensor:
        """Get value estimate for observations."""
        embedded_obs = self.obs_transformer(obs)
        features = self.trunk(embedded_obs)
        return self.value_head(features).squeeze(-1)

    def evaluate_actions(
        self,
        obs: torch.Tensor,
        action_mask: torch.Tensor,
        actions: torch.Tensor,
    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        """
        Evaluate log probability and entropy of given actions.

        Used during PPO update to compute policy loss.
        """
        logits, value = self.forward(obs, action_mask)
        dist = Categorical(logits=logits)

        log_prob = dist.log_prob(actions)
        entropy = dist.entropy()

        return log_prob, entropy, value


class EmbeddedAlphaZeroNetwork(nn.Module):
    """
    AlphaZero-style network with learned card embeddings.

    Uses residual blocks for deeper feature extraction and tanh activation
    on the value head for output in [-1, 1].

    Args:
        embed_dim: Dimension of card embeddings (default: 64)
        hidden_dim: Hidden layer dimension (default: 256)
        num_blocks: Number of residual blocks (default: 4)
        num_cards: Number of cards in embedding table (default: 5000)
        include_embed_section: Whether to embed the trailing card section
        pretrained_embeds: Optional pre-trained embedding weights
        freeze_embeds: Whether to freeze embeddings during training
    """

    def __init__(
        self,
        embed_dim: int = 64,
        hidden_dim: int = 256,
        num_blocks: int = 4,
        num_cards: int = MAX_CARD_ID,
        include_embed_section: bool = False,
        pretrained_embeds: torch.Tensor | None = None,
        freeze_embeds: bool = False,
    ):
        super().__init__()

        self.hidden_dim = hidden_dim
        self.num_blocks = num_blocks
        self._using_pretrained_embeds = pretrained_embeds is not None

        # Create or load card embedding layer
        self.card_embedding, self.embed_dim = create_card_embedding(
            num_cards=num_cards,
            embed_dim=embed_dim,
            pretrained_embeds=pretrained_embeds,
            freeze_embeds=freeze_embeds,
        )

        # Create observation transformer
        self.obs_transformer = ObservationTransformer(
            self.card_embedding,
            include_embed_section=include_embed_section,
        )

        # Calculate input dimension
        input_dim = self.obs_transformer.output_dim

        # Input projection
        self.input_proj = nn.Sequential(
            nn.Linear(input_dim, hidden_dim),
            nn.ReLU(),
        )

        # Residual tower
        self.residual_tower = nn.ModuleList([
            ResidualBlock(hidden_dim) for _ in range(num_blocks)
        ])

        # Policy head
        self.policy_head = nn.Sequential(
            nn.Linear(hidden_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, 422),
        )

        # Value head with tanh for [-1, 1] output
        self.value_head = nn.Sequential(
            nn.Linear(hidden_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, 1),
            nn.Tanh(),
        )

        self._init_weights()

    def _init_weights(self) -> None:
        """Initialize network weights."""
        # Only initialize embeddings if not using pretrained
        card_embed = None if self._using_pretrained_embeds else self.card_embedding
        init_network_weights(self, policy_head=self.policy_head, card_embedding=card_embed)

    def forward(
        self,
        obs: torch.Tensor,
        action_mask: torch.Tensor | None = None,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """
        Forward pass through the network.

        Returns:
            logits: Action logits of shape (batch, 422)
            value: State values of shape (batch,) in [-1, 1]
        """
        # Transform observation
        embedded_obs = self.obs_transformer(obs)

        # Input projection
        x = self.input_proj(embedded_obs)

        # Residual tower
        for block in self.residual_tower:
            x = block(x)

        # Heads
        logits = self.policy_head(x)
        value = self.value_head(x).squeeze(-1)

        # Apply action mask
        if action_mask is not None:
            logits = logits.masked_fill(~action_mask, -1e8)

        return logits, value

    def get_policy(
        self,
        obs: torch.Tensor,
        action_mask: torch.Tensor,
    ) -> torch.Tensor:
        """Get policy probabilities for MCTS."""
        logits, _ = self.forward(obs, action_mask)
        return torch.softmax(logits, dim=-1)

    def get_value(self, obs: torch.Tensor) -> torch.Tensor:
        """Get value estimate for observations."""
        embedded_obs = self.obs_transformer(obs)
        x = self.input_proj(embedded_obs)
        for block in self.residual_tower:
            x = block(x)
        return self.value_head(x).squeeze(-1)

    def evaluate(
        self,
        obs: torch.Tensor,
        action_mask: torch.Tensor,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """Evaluate position for MCTS (returns policy probs and value)."""
        logits, value = self.forward(obs, action_mask)
        policy = torch.softmax(logits, dim=-1)
        return policy, value


# =============================================================================
# Factory Functions
# =============================================================================

ObservationMode = Literal["flat", "embedded", "embedded_pretrained"]


def create_network(
    observation_mode: ObservationMode = "flat",
    network_type: Literal["ppo", "alphazero"] = "ppo",
    embed_dim: int = 64,
    hidden_dim: int = 256,
    num_blocks: int = 4,
    pretrained_path: str | None = None,
    freeze_embeds: bool = False,
    include_embed_section: bool = False,
) -> nn.Module:
    """
    Factory function to create networks with different observation modes.

    Args:
        observation_mode: "flat" for original, "embedded" for end-to-end,
                         "embedded_pretrained" for pre-trained embeddings
        network_type: "ppo" or "alphazero"
        embed_dim: Embedding dimension (ignored if pretrained)
        hidden_dim: Hidden layer dimension
        num_blocks: Number of residual blocks (alphazero only)
        pretrained_path: Path to pre-trained embeddings (for embedded_pretrained)
        freeze_embeds: Whether to freeze embeddings
        include_embed_section: Whether to embed the trailing card section

    Returns:
        Configured network module
    """
    # Load pre-trained embeddings if specified
    pretrained_embeds = None
    if observation_mode == "embedded_pretrained" and pretrained_path:
        pretrained_embeds = torch.load(pretrained_path, weights_only=False)
        if isinstance(pretrained_embeds, dict):
            pretrained_embeds = pretrained_embeds.get("weight", pretrained_embeds.get("embeddings"))

    if observation_mode == "flat":
        # Use original networks
        if network_type == "ppo":
            from essence_wars.agents.networks import EssenceWarsNetwork
            return EssenceWarsNetwork(hidden_dim=hidden_dim)
        else:
            from essence_wars.agents.networks import AlphaZeroNetwork
            return AlphaZeroNetwork(hidden_dim=hidden_dim, num_blocks=num_blocks)

    # Embedded modes
    if network_type == "ppo":
        return EmbeddedPPONetwork(
            embed_dim=embed_dim,
            hidden_dim=hidden_dim,
            pretrained_embeds=pretrained_embeds,
            freeze_embeds=freeze_embeds,
            include_embed_section=include_embed_section,
        )
    else:
        return EmbeddedAlphaZeroNetwork(
            embed_dim=embed_dim,
            hidden_dim=hidden_dim,
            num_blocks=num_blocks,
            pretrained_embeds=pretrained_embeds,
            freeze_embeds=freeze_embeds,
            include_embed_section=include_embed_section,
        )


def get_embedding_info() -> dict[str, int | dict[str, list[int] | int]]:
    """Get information about the embedding configuration."""
    positions = CardIdPositions.default()
    non_card = get_non_card_positions(positions)

    return {
        "state_tensor_size": STATE_TENSOR_SIZE,
        "num_card_positions": len(positions.all_positions()),
        "num_non_card_positions": len(non_card),
        "card_positions": {
            "p1_hand": positions.p1_hand,
            "p2_hand": positions.p2_hand,
            "p1_supports": positions.p1_supports,
            "p2_supports": positions.p2_supports,
            "embed_section_size": len(positions.embed_section),
        },
        "max_card_id": MAX_CARD_ID,
    }
