"""
Action Heads — G1-specific action decoders.

Converts visual-language features from the video backbone into robot actions.
Supports multiple output heads for different control groups:
- Upper body (14 DoF arms)
- Grippers (2 DoF)
- Locomotion velocity (3 DoF: vx, vy, omega)
- Whole body (all 29 DoF, when enabled)

Architecture inspired by:
- GR00T N1's action tokenizer (VQ-BeT style)
- OmniVLA's multi-head action decoder
- RT-2's action token prediction

v2 enhancements from MicroGPT Parameter Golf winners:
- ReLU² activation (smoother gradients than GELU, cheaper)
- RMSNorm (lighter than LayerNorm, no mean subtraction)
- Learnable residual scales (prevents frozen backbone from dominating)
- U-Net skip connections (gradient highway through deeper MLPs)
- Logit soft-capping (smoother than hard Tanh clamp)

Key insight: Action heads are our "small model" (~2M params). Every
parameter-efficiency trick from the golf competition applies here.
"""

from __future__ import annotations

import logging
import math
from dataclasses import dataclass
from typing import Dict, Optional, Tuple

import torch
import torch.nn as nn
import torch.nn.functional as F

from neon.data.action_space import ControlMode, G1ActionSpace

logger = logging.getLogger(__name__)


# =============================================================================
# Parameter Golf primitives — ported from MicroGPT v2 to PyTorch
# =============================================================================


class ReLUSquared(nn.Module):
    """ReLU² activation: max(0, x)².

    From Parameter Golf winners — smoother gradients than ReLU, cheaper than GELU.
    Gradient: 2 * max(0, x) — continuous at zero unlike ReLU.
    For action heads: produces smoother action trajectories.
    """

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return F.relu(x).square()


class RMSNorm(nn.Module):
    """Root Mean Square Layer Normalization.

    From MicroGPT engine: scale = (mean(x²) + ε)^(-0.5)
    Lighter than LayerNorm — no mean subtraction, no bias.
    Same stabilization benefit for small MLPs.
    """

    def __init__(self, dim: int, eps: float = 1e-5):
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(dim))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        ms = x.pow(2).mean(-1, keepdim=True)
        x_normed = x * torch.rsqrt(ms + self.eps)
        return x_normed * self.weight


def soft_cap(x: torch.Tensor, cap: float = 30.0) -> torch.Tensor:
    """Logit soft-capping: cap * tanh(x / cap).

    From Parameter Golf — smoother than hard Tanh clamp.
    Approaches [-cap, cap] asymptotically instead of hitting a wall.
    For action outputs capped to [-1, 1], use cap=1.0.
    """
    return cap * torch.tanh(x / cap)


@dataclass
class ActionHeadConfig:
    """Configuration for the action decoder heads."""

    hidden_size: int = 2048  # Input feature dim from backbone
    action_dim: int = 17  # Total action output dim (depends on ControlMode)
    num_action_steps: int = 16  # Predict this many future steps (action chunking)
    mlp_hidden: int = 512  # Hidden dim of MLP layers
    num_layers: int = 3  # Depth of the MLP decoder
    dropout: float = 0.1
    use_separate_heads: bool = True  # Separate MLP per group (arms, grip, loco)
    use_action_chunking: bool = True  # Predict multiple future steps
    # Parameter Golf v2 options
    use_relu_squared: bool = True  # ReLU² instead of GELU
    use_rmsnorm: bool = True  # RMSNorm before each MLP layer
    use_residual_scale: bool = True  # Learnable residual connection scales
    use_skip_connections: bool = True  # U-Net skip from layer 0 → last
    soft_cap_value: float = 1.0  # Soft-cap output (0 = use hard Tanh)


class MLPDecoder(nn.Module):
    """MLP action decoder with Parameter Golf v2 enhancements.

    Enhancements over vanilla MLP:
    - RMSNorm before each hidden layer (training stability)
    - ReLU² activation (smoother action outputs)
    - Learnable residual scales (prevents backbone feature domination)
    - U-Net skip connection from first hidden to last hidden
    - Soft-capping instead of hard Tanh for output range
    """

    def __init__(
        self,
        input_dim: int,
        output_dim: int,
        hidden_dim: int = 512,
        num_layers: int = 3,
        dropout: float = 0.1,
        use_relu_squared: bool = True,
        use_rmsnorm: bool = True,
        use_residual_scale: bool = True,
        use_skip_connections: bool = True,
        soft_cap_value: float = 1.0,
    ):
        super().__init__()
        self.use_skip_connections = use_skip_connections and num_layers >= 3
        self.soft_cap_value = soft_cap_value

        activation = ReLUSquared if use_relu_squared else nn.GELU

        # Build layers individually for skip connection support
        self.input_proj = nn.Linear(input_dim, hidden_dim)
        self.input_norm = RMSNorm(hidden_dim) if use_rmsnorm else nn.Identity()
        self.input_act = activation()
        self.input_drop = nn.Dropout(dropout)

        # Hidden layers
        self.hidden_layers = nn.ModuleList()
        self.hidden_norms = nn.ModuleList()
        self.hidden_acts = nn.ModuleList()
        self.hidden_drops = nn.ModuleList()
        for i in range(num_layers - 2):
            # Skip connection doubles input dim for the last hidden layer
            in_d = hidden_dim * 2 if (self.use_skip_connections and i == num_layers - 3) else hidden_dim
            self.hidden_layers.append(nn.Linear(in_d, hidden_dim))
            self.hidden_norms.append(RMSNorm(hidden_dim) if use_rmsnorm else nn.Identity())
            self.hidden_acts.append(activation())
            self.hidden_drops.append(nn.Dropout(dropout))

        # Output projection — no activation, we use soft-cap
        self.output_proj = nn.Linear(hidden_dim, output_dim)

        # Learnable residual scales (one per hidden layer + input)
        if use_residual_scale and num_layers >= 3:
            # α starts at 1.0 — network learns to scale residuals
            self.residual_scales = nn.ParameterList([
                nn.Parameter(torch.ones(1)) for _ in range(num_layers - 1)
            ])
        else:
            self.residual_scales = None

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # Input projection
        h = self.input_proj(x)
        h = self.input_norm(h)
        h = self.input_act(h)
        h = self.input_drop(h)
        skip = h  # Save for U-Net skip

        # Hidden layers
        for i, (linear, norm, act, drop) in enumerate(
            zip(self.hidden_layers, self.hidden_norms, self.hidden_acts, self.hidden_drops)
        ):
            h_in = h
            # U-Net skip: concat first hidden features into last hidden layer
            if self.use_skip_connections and i == len(self.hidden_layers) - 1:
                h = torch.cat([h, skip], dim=-1)

            h = linear(h)
            h = norm(h)
            h = act(h)
            h = drop(h)

            # Residual with learnable scale (only if dims match)
            if self.residual_scales is not None and h.shape == h_in.shape:
                scale = self.residual_scales[min(i, len(self.residual_scales) - 1)]
                h = h + scale * h_in

        # Output with soft-capping
        out = self.output_proj(h)
        if self.soft_cap_value > 0:
            out = soft_cap(out, self.soft_cap_value)
        else:
            out = torch.tanh(out)  # Legacy hard clamp
        return out


class ActionChunkingHead(nn.Module):
    """Action chunking decoder — predicts multiple future action steps.

    Instead of predicting a single action, predicts a sequence of future actions.
    This reduces compounding errors and improves policy smoothness.
    Reference: ACT (Action Chunking with Transformers), Zhao et al. 2023.

    v2: Uses Parameter Golf primitives (ReLU², RMSNorm, soft-capping).

    Output shape: (batch, num_steps, action_dim)
    """

    def __init__(
        self,
        input_dim: int,
        action_dim: int,
        num_steps: int = 16,
        hidden_dim: int = 512,
        num_layers: int = 3,
        dropout: float = 0.1,
        use_relu_squared: bool = True,
        use_rmsnorm: bool = True,
        soft_cap_value: float = 1.0,
    ):
        super().__init__()
        self.num_steps = num_steps
        self.action_dim = action_dim
        self.soft_cap_value = soft_cap_value

        activation = ReLUSquared if use_relu_squared else nn.GELU

        # Temporal embedding for each step
        self.step_embed = nn.Embedding(num_steps, hidden_dim)

        # Feature projection with RMSNorm
        self.feature_proj = nn.Sequential(
            nn.Linear(input_dim, hidden_dim),
            RMSNorm(hidden_dim) if use_rmsnorm else nn.Identity(),
            activation(),
            nn.Dropout(dropout),
        )

        # Per-step decoder with v2 enhancements
        self.decoder = nn.Sequential(
            nn.Linear(hidden_dim * 2, hidden_dim),  # Features + step embedding
            RMSNorm(hidden_dim) if use_rmsnorm else nn.Identity(),
            activation(),
            nn.Dropout(dropout),
            nn.Linear(hidden_dim, hidden_dim),
            RMSNorm(hidden_dim) if use_rmsnorm else nn.Identity(),
            activation(),
            nn.Dropout(dropout),
            nn.Linear(hidden_dim, action_dim),
            # No activation — soft-cap applied in forward()
        )

        # Learnable residual scale for decoder skip
        self.residual_scale = nn.Parameter(torch.ones(1))

    def forward(self, features: torch.Tensor) -> torch.Tensor:
        """Predict a chunk of future actions.

        Args:
            features: (batch, hidden_size) pooled backbone features

        Returns:
            (batch, num_steps, action_dim) predicted action chunk
        """
        batch_size = features.shape[0]
        device = features.device

        # Project features
        feat = self.feature_proj(features)  # (batch, hidden_dim)

        # Generate step embeddings
        step_ids = torch.arange(self.num_steps, device=device)
        step_emb = self.step_embed(step_ids)  # (num_steps, hidden_dim)

        # Expand and concatenate
        feat_expanded = feat.unsqueeze(1).expand(-1, self.num_steps, -1)  # (batch, num_steps, hidden_dim)
        step_expanded = step_emb.unsqueeze(0).expand(batch_size, -1, -1)  # (batch, num_steps, hidden_dim)
        combined = torch.cat([feat_expanded, step_expanded], dim=-1)  # (batch, num_steps, hidden_dim*2)

        # Decode each step with soft-capping
        actions = self.decoder(combined)  # (batch, num_steps, action_dim)
        if self.soft_cap_value > 0:
            actions = soft_cap(actions, self.soft_cap_value)
        else:
            actions = torch.tanh(actions)

        return actions


class G1ActionHead(nn.Module):
    """Multi-head action decoder for the G1 humanoid.

    Can use either:
    1. Single MLP for all actions
    2. Separate heads per group (arms, grippers, locomotion)

    With optional action chunking for temporal smoothness.
    """

    def __init__(self, config: ActionHeadConfig, action_space: G1ActionSpace):
        super().__init__()
        self.config = config
        self.action_space = action_space

        if config.use_separate_heads:
            self._build_separate_heads()
        else:
            self._build_single_head()

        # Count parameters
        total_params = sum(p.numel() for p in self.parameters())
        trainable_params = sum(p.numel() for p in self.parameters() if p.requires_grad)
        logger.info(f"ActionHead: {trainable_params:,} trainable / {total_params:,} total parameters")

    def _build_single_head(self):
        """Single MLP/chunking head for all actions."""
        if self.config.use_action_chunking:
            self.head = ActionChunkingHead(
                input_dim=self.config.hidden_size,
                action_dim=self.action_space.action_dim,
                num_steps=self.config.num_action_steps,
                hidden_dim=self.config.mlp_hidden,
                dropout=self.config.dropout,
                use_relu_squared=self.config.use_relu_squared,
                use_rmsnorm=self.config.use_rmsnorm,
                soft_cap_value=self.config.soft_cap_value,
            )
        else:
            self.head = MLPDecoder(
                input_dim=self.config.hidden_size,
                output_dim=self.action_space.action_dim,
                hidden_dim=self.config.mlp_hidden,
                num_layers=self.config.num_layers,
                dropout=self.config.dropout,
                use_relu_squared=self.config.use_relu_squared,
                use_rmsnorm=self.config.use_rmsnorm,
                use_residual_scale=self.config.use_residual_scale,
                use_skip_connections=self.config.use_skip_connections,
                soft_cap_value=self.config.soft_cap_value,
            )

    def _build_separate_heads(self):
        """Separate decoder head per action group."""
        self.heads = nn.ModuleDict()

        # Common v2 kwargs for all sub-heads
        v2_mlp = dict(
            use_relu_squared=self.config.use_relu_squared,
            use_rmsnorm=self.config.use_rmsnorm,
            use_residual_scale=self.config.use_residual_scale,
            use_skip_connections=self.config.use_skip_connections,
            soft_cap_value=self.config.soft_cap_value,
        )
        v2_chunk = dict(
            use_relu_squared=self.config.use_relu_squared,
            use_rmsnorm=self.config.use_rmsnorm,
            soft_cap_value=self.config.soft_cap_value,
        )

        # Arms head (left + right)
        arm_indices = (
            self.action_space.get_group_indices("left_arm")
            + self.action_space.get_group_indices("right_arm")
        )
        if arm_indices:
            arm_dim = len(arm_indices)
            if self.config.use_action_chunking:
                self.heads["arms"] = ActionChunkingHead(
                    self.config.hidden_size, arm_dim,
                    self.config.num_action_steps, self.config.mlp_hidden,
                    dropout=self.config.dropout, **v2_chunk,
                )
            else:
                self.heads["arms"] = MLPDecoder(
                    self.config.hidden_size, arm_dim,
                    self.config.mlp_hidden, self.config.num_layers,
                    self.config.dropout, **v2_mlp,
                )

        # Locomotion head
        if self.action_space.include_locomotion:
            if self.config.use_action_chunking:
                self.heads["locomotion"] = ActionChunkingHead(
                    self.config.hidden_size, 3,
                    self.config.num_action_steps, self.config.mlp_hidden // 2,
                    dropout=self.config.dropout, **v2_chunk,
                )
            else:
                self.heads["locomotion"] = MLPDecoder(
                    self.config.hidden_size, 3,
                    self.config.mlp_hidden // 2, 2,
                    self.config.dropout, **v2_mlp,
                )

        # Torso + head joints (if in upper_body or whole_body mode)
        if self.action_space.mode in (ControlMode.UPPER_BODY, ControlMode.WHOLE_BODY):
            torso_head_indices = (
                self.action_space.get_group_indices("torso")
                + self.action_space.get_group_indices("head")
            )
            if torso_head_indices:
                dim = len(torso_head_indices)
                self.heads["torso_head"] = MLPDecoder(
                    self.config.hidden_size, dim,
                    self.config.mlp_hidden // 2, 2,
                    self.config.dropout, **v2_mlp,
                )

        # Leg joints (if whole_body mode)
        if self.action_space.mode == ControlMode.WHOLE_BODY:
            leg_indices = (
                self.action_space.get_group_indices("left_leg")
                + self.action_space.get_group_indices("right_leg")
            )
            if leg_indices:
                dim = len(leg_indices)
                self.heads["legs"] = MLPDecoder(
                    self.config.hidden_size, dim,
                    self.config.mlp_hidden, self.config.num_layers,
                    self.config.dropout, **v2_mlp,
                )

        # Record index mapping for assembly
        self._arm_indices = arm_indices if arm_indices else []
        self._torso_head_indices = (
            self.action_space.get_group_indices("torso")
            + self.action_space.get_group_indices("head")
        ) if self.action_space.mode != ControlMode.ARMS_ONLY else []
        self._leg_indices = (
            self.action_space.get_group_indices("left_leg")
            + self.action_space.get_group_indices("right_leg")
        ) if self.action_space.mode == ControlMode.WHOLE_BODY else []

    def forward(self, features: torch.Tensor) -> torch.Tensor:
        """Decode features into actions.

        Args:
            features: (batch, hidden_size) pooled features from backbone

        Returns:
            If action_chunking: (batch, num_steps, action_dim)
            Else: (batch, action_dim)
        """
        if not self.config.use_separate_heads:
            return self.head(features)

        # Separate heads — assemble output
        parts = []

        if "arms" in self.heads:
            parts.append(self.heads["arms"](features))

        if "torso_head" in self.heads:
            parts.append(self.heads["torso_head"](features))

        if "legs" in self.heads:
            parts.append(self.heads["legs"](features))

        if "locomotion" in self.heads:
            parts.append(self.heads["locomotion"](features))

        # Concatenate along action dimension
        if self.config.use_action_chunking:
            # All parts: (batch, num_steps, dim_i)
            return torch.cat(parts, dim=-1)
        else:
            # All parts: (batch, dim_i)
            return torch.cat(parts, dim=-1)

    def compute_loss(
        self,
        predicted: torch.Tensor,
        target: torch.Tensor,
        mask: Optional[torch.Tensor] = None,
    ) -> torch.Tensor:
        """Compute MSE loss between predicted and target actions.

        Args:
            predicted: Predicted actions (batch, [num_steps,] action_dim)
            target: Ground truth actions (batch, [num_steps,] action_dim)
            mask: Optional mask for valid timesteps (batch, [num_steps])

        Returns:
            Scalar loss
        """
        loss = F.mse_loss(predicted, target, reduction="none")  # Per-element loss

        if mask is not None:
            # Expand mask to action dim
            if loss.dim() == 3:  # (batch, steps, dim)
                mask = mask.unsqueeze(-1).expand_as(loss)
            loss = loss * mask

        return loss.mean()
