"""Policy-matched proxy architectures for data quality assessment.

Each architecture mirrors the learning dynamics of its full-scale counterpart
so the proxy's GO/NO-GO signal is relevant to the actual policy:

- ProxyMLP: Basic BC (2-layer MLP) — tests basic learnability
- ProxyACT: Tiny Transformer with action chunking — tests ACT-style learning
- ProxyDiffusion: 1D UNet with DDPM — tests multimodal action distributions
- ProxyRNN: Single-layer LSTM — tests temporal sequence learning

References:
- ACT (Zhao et al. RSS 2023): only first decoder layer was used (known bug)
- Diffusion Policy (Chi et al. RSS 2023): ConditionalUnet1D, ~330 lines
- DataMIL (2025): BC validation loss as proxy for downstream performance
"""

from __future__ import annotations

import math

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


# ---------------------------------------------------------------------------
# ProxyMLP — the existing 2-layer MLP (kept for BC / BC-RNN fallback)
# ---------------------------------------------------------------------------


class ProxyMLP(nn.Module):
    """2-layer MLP for basic behavioral cloning proxy."""

    def __init__(self, state_dim: int, action_dim: int, hidden_dim: int = 128):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(state_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, action_dim),
        )

    def forward(self, states: torch.Tensor) -> torch.Tensor:
        return self.net(states)


# ---------------------------------------------------------------------------
# ProxyACT — tiny Transformer encoder with action chunking
# ---------------------------------------------------------------------------


class ProxyACT(nn.Module):
    """Tiny Transformer for ACT-style action chunking proxy.

    Architecture matches ACT's learning dynamics:
    - Encoder-only (the original ACT decoder bug means only 1 layer was used)
    - Action chunking: predicts chunk_size future actions at once
    - Small: dim=64, 2 encoder layers, 4 heads

    ~200-500K params depending on state/action dims.
    """

    def __init__(
        self,
        state_dim: int,
        action_dim: int,
        d_model: int = 64,
        n_heads: int = 4,
        n_layers: int = 2,
        chunk_size: int = 10,
    ):
        super().__init__()
        self.chunk_size = chunk_size
        self.d_model = d_model

        self.state_proj = nn.Linear(state_dim, d_model)
        self.pos_embed = nn.Parameter(torch.randn(1, chunk_size, d_model) * 0.02)

        encoder_layer = nn.TransformerEncoderLayer(
            d_model=d_model,
            nhead=n_heads,
            dim_feedforward=d_model * 4,
            dropout=0.1,
            batch_first=True,
        )
        self.encoder = nn.TransformerEncoder(encoder_layer, num_layers=n_layers)
        self.action_head = nn.Linear(d_model, action_dim)

    def forward(self, states: torch.Tensor) -> torch.Tensor:
        """Predict chunk_size actions from a single state.

        Args:
            states: (B, state_dim)

        Returns:
            (B, chunk_size, action_dim) predicted action chunks.
        """
        B = states.shape[0]
        # Project state and tile across chunk positions
        s = self.state_proj(states).unsqueeze(1).expand(-1, self.chunk_size, -1)
        s = s + self.pos_embed[:, : self.chunk_size, :]

        h = self.encoder(s)
        return self.action_head(h)


# ---------------------------------------------------------------------------
# ProxyDiffusion — minimal 1D UNet with DDPM
# ---------------------------------------------------------------------------


class _SinusoidalPosEmb(nn.Module):
    """Sinusoidal timestep embedding for diffusion."""

    def __init__(self, dim: int):
        super().__init__()
        self.dim = dim

    def forward(self, t: torch.Tensor) -> torch.Tensor:
        half = self.dim // 2
        emb = math.log(10000) / (half - 1)
        emb = torch.exp(torch.arange(half, device=t.device, dtype=torch.float32) * -emb)
        emb = t.float().unsqueeze(1) * emb.unsqueeze(0)
        return torch.cat([emb.sin(), emb.cos()], dim=-1)


class _ResBlock1D(nn.Module):
    """1D residual block conditioned on timestep embedding."""

    def __init__(self, in_ch: int, out_ch: int, time_dim: int):
        super().__init__()
        self.conv1 = nn.Conv1d(in_ch, out_ch, 3, padding=1)
        self.conv2 = nn.Conv1d(out_ch, out_ch, 3, padding=1)
        self.time_proj = nn.Linear(time_dim, out_ch)
        self.skip = nn.Conv1d(in_ch, out_ch, 1) if in_ch != out_ch else nn.Identity()
        self.norm1 = nn.GroupNorm(min(8, out_ch), out_ch)
        self.norm2 = nn.GroupNorm(min(8, out_ch), out_ch)

    def forward(self, x: torch.Tensor, t_emb: torch.Tensor) -> torch.Tensor:
        h = F.silu(self.norm1(self.conv1(x)))
        h = h + self.time_proj(t_emb).unsqueeze(-1)
        h = F.silu(self.norm2(self.conv2(h)))
        return h + self.skip(x)


class ProxyDiffusion(nn.Module):
    """Minimal 1D UNet for diffusion policy proxy.

    2-level UNet with DDPM noise schedule. Tests whether the data
    has learnable multimodal action distributions.

    Architecture: channels=(64, 128), 20 DDPM steps at inference.
    ~300K params.
    """

    def __init__(
        self,
        state_dim: int,
        action_dim: int,
        horizon: int = 16,
        channels: tuple[int, ...] = (64, 128),
        n_diffusion_steps: int = 20,
    ):
        super().__init__()
        self.horizon = horizon
        self.action_dim = action_dim
        self.n_steps = n_diffusion_steps
        time_dim = channels[0] * 4

        # Timestep embedding
        self.time_mlp = nn.Sequential(
            _SinusoidalPosEmb(channels[0]),
            nn.Linear(channels[0], time_dim),
            nn.SiLU(),
            nn.Linear(time_dim, time_dim),
        )

        # State conditioning — project to time_dim so it can be added
        self.state_proj = nn.Linear(state_dim, time_dim)

        # Input projection: action_dim -> channels[0]
        self.input_proj = nn.Conv1d(action_dim, channels[0], 1)

        # Downsampling
        self.down1 = _ResBlock1D(channels[0], channels[0], time_dim)
        self.down2 = _ResBlock1D(channels[0], channels[1], time_dim)
        self.downsample = nn.Conv1d(channels[1], channels[1], 2, stride=2)

        # Middle
        self.mid = _ResBlock1D(channels[1], channels[1], time_dim)

        # Upsampling
        self.upsample = nn.ConvTranspose1d(channels[1], channels[1], 2, stride=2)
        self.up1 = _ResBlock1D(channels[1] * 2, channels[0], time_dim)  # skip connection
        self.up2 = _ResBlock1D(channels[0], channels[0], time_dim)

        # Output
        self.output_proj = nn.Conv1d(channels[0], action_dim, 1)

        # DDPM schedule
        betas = torch.linspace(1e-4, 0.02, n_diffusion_steps)
        alphas = 1.0 - betas
        alphas_cumprod = torch.cumprod(alphas, dim=0)
        self.register_buffer("betas", betas)
        self.register_buffer("alphas", alphas)
        self.register_buffer("alphas_cumprod", alphas_cumprod)
        self.register_buffer(
            "sqrt_alphas_cumprod", torch.sqrt(alphas_cumprod)
        )
        self.register_buffer(
            "sqrt_one_minus_alphas_cumprod", torch.sqrt(1.0 - alphas_cumprod)
        )

    def forward(
        self, noisy_actions: torch.Tensor, timesteps: torch.Tensor, states: torch.Tensor
    ) -> torch.Tensor:
        """Predict noise given noisy actions, timestep, and conditioning state.

        Args:
            noisy_actions: (B, action_dim, horizon) noisy action trajectory
            timesteps: (B,) integer diffusion timesteps
            states: (B, state_dim) conditioning states

        Returns:
            (B, action_dim, horizon) predicted noise
        """
        t_emb = self.time_mlp(timesteps)

        # Add state conditioning to time embedding
        s_emb = F.silu(self.state_proj(states))
        t_emb = t_emb + s_emb

        # UNet forward
        x = self.input_proj(noisy_actions)
        h1 = self.down1(x, t_emb)
        h2 = self.down2(h1, t_emb)
        h2_down = self.downsample(h2)

        h_mid = self.mid(h2_down, t_emb)

        h_up = self.upsample(h_mid)
        # Pad if sizes don't match exactly
        if h_up.shape[-1] != h2.shape[-1]:
            h_up = F.pad(h_up, (0, h2.shape[-1] - h_up.shape[-1]))
        h_up = torch.cat([h_up, h2], dim=1)
        h_up = self.up1(h_up, t_emb)
        h_up = self.up2(h_up, t_emb)

        return self.output_proj(h_up)

    def compute_loss(
        self, actions: torch.Tensor, states: torch.Tensor
    ) -> torch.Tensor:
        """Compute DDPM training loss (predict noise).

        Args:
            actions: (B, action_dim, horizon) clean action trajectories
            states: (B, state_dim) conditioning states

        Returns:
            Scalar MSE loss.
        """
        B = actions.shape[0]
        t = torch.randint(0, self.n_steps, (B,), device=actions.device)
        noise = torch.randn_like(actions)

        sqrt_alpha = self.sqrt_alphas_cumprod[t].view(B, 1, 1)
        sqrt_one_minus = self.sqrt_one_minus_alphas_cumprod[t].view(B, 1, 1)
        noisy = sqrt_alpha * actions + sqrt_one_minus * noise

        pred_noise = self(noisy, t, states)
        return F.mse_loss(pred_noise, noise)


# ---------------------------------------------------------------------------
# ProxyRNN — single-layer LSTM
# ---------------------------------------------------------------------------


class ProxyRNN(nn.Module):
    """Single-layer LSTM for temporal sequence learning proxy.

    Tests whether the data has learnable temporal dependencies.
    ~30K params.
    """

    def __init__(self, state_dim: int, action_dim: int, hidden_dim: int = 64):
        super().__init__()
        self.lstm = nn.LSTM(
            input_size=state_dim,
            hidden_size=hidden_dim,
            num_layers=1,
            batch_first=True,
        )
        self.head = nn.Linear(hidden_dim, action_dim)

    def forward(self, states: torch.Tensor) -> torch.Tensor:
        """Predict actions from a sequence of states.

        Args:
            states: (B, T, state_dim) state sequence

        Returns:
            (B, T, action_dim) predicted actions
        """
        h, _ = self.lstm(states)
        return self.head(h)


# ---------------------------------------------------------------------------
# Architecture registry
# ---------------------------------------------------------------------------

PROXY_ARCHITECTURES = {
    "bc": "mlp",
    "bc_rnn": "rnn",
    "act": "act",
    "diffusion_policy": "diffusion",
    "dp3": "diffusion",
    "smolvla": "mlp",
}


def get_proxy_architecture(policy_type: str) -> str:
    """Return the proxy architecture name for a given policy type."""
    return PROXY_ARCHITECTURES.get(policy_type, "mlp")
