from __future__ import annotations

from typing import Optional

import torch
import torch.nn.functional as F
from diffusers import ConfigMixin, ModelMixin
from diffusers.configuration_utils import register_to_config
from diffusers.models.attention import Attention, FeedForward
from diffusers.models.embeddings import SinusoidalPositionalEmbedding, TimestepEmbedding, Timesteps
from torch import nn


class TimestepEncoder(nn.Module):
    def __init__(self, embedding_dim: int):
        super().__init__()
        self.time_proj = Timesteps(num_channels=256, flip_sin_to_cos=True, downscale_freq_shift=1)
        self.timestep_embedder = TimestepEmbedding(in_channels=256, time_embed_dim=embedding_dim)

    def forward(self, timesteps: torch.Tensor) -> torch.Tensor:
        return self.timestep_embedder(self.time_proj(timesteps).to(next(self.parameters()).dtype))


class AdaLayerNorm(nn.Module):
    def __init__(
        self,
        embedding_dim: int,
        norm_elementwise_affine: bool = False,
        norm_eps: float = 1e-5,
    ):
        super().__init__()
        self.silu = nn.SiLU()
        self.linear = nn.Linear(embedding_dim, embedding_dim * 2)
        self.norm = nn.LayerNorm(embedding_dim, norm_eps, norm_elementwise_affine)

    def forward(self, x: torch.Tensor, temb: torch.Tensor) -> torch.Tensor:
        scale, shift = self.linear(self.silu(temb)).chunk(2, dim=1)
        return self.norm(x) * (1 + scale[:, None]) + shift[:, None]


class BasicTransformerBlock(nn.Module):
    def __init__(
        self,
        dim: int,
        num_attention_heads: int,
        attention_head_dim: int,
        dropout: float = 0.0,
        cross_attention_dim: Optional[int] = None,
        activation_fn: str = "geglu",
        attention_bias: bool = False,
        upcast_attention: bool = False,
        norm_elementwise_affine: bool = True,
        norm_type: str = "layer_norm",
        norm_eps: float = 1e-5,
        final_dropout: bool = False,
        positional_embeddings: Optional[str] = None,
        num_positional_embeddings: Optional[int] = None,
        ff_inner_dim: Optional[int] = None,
        ff_bias: bool = True,
        attention_out_bias: bool = True,
    ):
        super().__init__()
        if positional_embeddings and num_positional_embeddings is None:
            raise ValueError("num_positional_embeddings is required for positional embeddings")
        self.norm_type = norm_type
        self.pos_embed = (
            SinusoidalPositionalEmbedding(dim, max_seq_length=num_positional_embeddings)
            if positional_embeddings == "sinusoidal"
            else None
        )
        self.norm1 = (
            AdaLayerNorm(dim)
            if norm_type == "ada_norm"
            else nn.LayerNorm(dim, elementwise_affine=norm_elementwise_affine, eps=norm_eps)
        )
        self.attn1 = Attention(
            query_dim=dim,
            heads=num_attention_heads,
            dim_head=attention_head_dim,
            dropout=dropout,
            bias=attention_bias,
            cross_attention_dim=cross_attention_dim,
            upcast_attention=upcast_attention,
            out_bias=attention_out_bias,
        )
        self.norm3 = nn.LayerNorm(dim, norm_eps, norm_elementwise_affine)
        self.ff = FeedForward(
            dim,
            dropout=dropout,
            activation_fn=activation_fn,
            final_dropout=final_dropout,
            inner_dim=ff_inner_dim,
            bias=ff_bias,
        )
        self.final_dropout = nn.Dropout(dropout) if final_dropout else None

    def forward(
        self,
        hidden_states: torch.Tensor,
        encoder_hidden_states: Optional[torch.Tensor] = None,
        encoder_attention_mask: Optional[torch.Tensor] = None,
        temb: Optional[torch.Tensor] = None,
    ) -> torch.Tensor:
        norm_hidden_states = (
            self.norm1(hidden_states, temb) if self.norm_type == "ada_norm" else self.norm1(hidden_states)
        )
        if self.pos_embed is not None:
            norm_hidden_states = self.pos_embed(norm_hidden_states)
        attention_output = self.attn1(
            norm_hidden_states,
            encoder_hidden_states=encoder_hidden_states,
            attention_mask=encoder_attention_mask,
        )
        if self.final_dropout is not None:
            attention_output = self.final_dropout(attention_output)
        hidden_states = attention_output + hidden_states
        if hidden_states.ndim == 4:
            hidden_states = hidden_states.squeeze(1)
        hidden_states = self.ff(self.norm3(hidden_states)) + hidden_states
        if hidden_states.ndim == 4:
            hidden_states = hidden_states.squeeze(1)
        return hidden_states


class DiT(ModelMixin, ConfigMixin):
    _supports_gradient_checkpointing = True

    @register_to_config
    def __init__(
        self,
        num_attention_heads: int = 8,
        attention_head_dim: int = 64,
        output_dim: int = 26,
        num_layers: int = 12,
        dropout: float = 0.1,
        attention_bias: bool = True,
        activation_fn: str = "gelu-approximate",
        num_embeds_ada_norm: Optional[int] = 1000,
        upcast_attention: bool = False,
        norm_type: str = "ada_norm",
        norm_elementwise_affine: bool = False,
        norm_eps: float = 1e-5,
        max_num_positional_embeddings: int = 512,
        compute_dtype: torch.dtype = torch.float32,
        final_dropout: bool = True,
        positional_embeddings: Optional[str] = "sinusoidal",
        interleave_self_attention: bool = False,
        cross_attention_dim: Optional[int] = None,
        **kwargs,
    ):
        super().__init__()
        self.attention_head_dim = attention_head_dim
        self.inner_dim = self.config.num_attention_heads * self.config.attention_head_dim
        self.gradient_checkpointing = False
        self.timestep_encoder = TimestepEncoder(self.inner_dim)
        self.transformer_blocks = nn.ModuleList(
            [
                BasicTransformerBlock(
                    self.inner_dim,
                    self.config.num_attention_heads,
                    self.config.attention_head_dim,
                    dropout=self.config.dropout,
                    activation_fn=self.config.activation_fn,
                    attention_bias=self.config.attention_bias,
                    upcast_attention=self.config.upcast_attention,
                    norm_type=norm_type,
                    norm_elementwise_affine=self.config.norm_elementwise_affine,
                    norm_eps=self.config.norm_eps,
                    positional_embeddings=positional_embeddings,
                    num_positional_embeddings=self.config.max_num_positional_embeddings,
                    final_dropout=final_dropout,
                    cross_attention_dim=(
                        None if index % 2 == 1 and interleave_self_attention else cross_attention_dim
                    ),
                )
                for index in range(self.config.num_layers)
            ]
        )
        self.norm_out = nn.LayerNorm(self.inner_dim, elementwise_affine=False, eps=1e-6)
        self.proj_out_1 = nn.Linear(self.inner_dim, 2 * self.inner_dim)
        self.proj_out_2 = nn.Linear(self.inner_dim, self.config.output_dim)

    def forward(
        self,
        hidden_states: torch.Tensor,
        encoder_hidden_states: torch.Tensor,
        timestep: Optional[torch.LongTensor] = None,
        return_all_hidden_states: bool = False,
        encoder_attention_mask: Optional[torch.Tensor] = None,
    ):
        time_embedding = self.timestep_encoder(timestep)
        hidden_states = hidden_states.contiguous()
        encoder_hidden_states = encoder_hidden_states.contiguous()
        all_hidden_states = [hidden_states]
        for index, block in enumerate(self.transformer_blocks):
            self_attention = index % 2 == 1 and self.config.interleave_self_attention
            hidden_states = block(
                hidden_states,
                encoder_hidden_states=None if self_attention else encoder_hidden_states,
                encoder_attention_mask=None if self_attention else encoder_attention_mask,
                temb=time_embedding,
            )
            all_hidden_states.append(hidden_states)
        shift, scale = self.proj_out_1(F.silu(time_embedding)).chunk(2, dim=1)
        hidden_states = self.norm_out(hidden_states) * (1 + scale[:, None]) + shift[:, None]
        output = self.proj_out_2(hidden_states)
        return (output, all_hidden_states) if return_all_hidden_states else output


class SinusoidalPositionalEncoding(nn.Module):
    def __init__(self, embedding_dim: int):
        super().__init__()
        self.embedding_dim = embedding_dim

    def forward(self, timesteps: torch.Tensor) -> torch.Tensor:
        timesteps = timesteps.float()
        half_dim = self.embedding_dim // 2
        exponent = -torch.arange(half_dim, dtype=torch.float, device=timesteps.device) * (
            torch.log(torch.tensor(10000.0, device=timesteps.device)) / half_dim
        )
        frequencies = timesteps.unsqueeze(-1) * exponent.exp()
        return torch.cat([torch.sin(frequencies), torch.cos(frequencies)], dim=-1)


class CategorySpecificLinear(nn.Module):
    def __init__(self, num_categories: int, input_dim: int, output_dim: int):
        super().__init__()
        self.num_categories = num_categories
        self.W = nn.Parameter(0.02 * torch.randn(num_categories, input_dim, output_dim))
        self.b = nn.Parameter(torch.zeros(num_categories, output_dim))

    def forward(self, x: torch.Tensor, category_ids: torch.Tensor) -> torch.Tensor:
        if category_ids is None:
            raise ValueError("embodiment_id (B,) is required")
        return torch.bmm(x, self.W[category_ids]) + self.b[category_ids].unsqueeze(1)


class CategorySpecificMLP(nn.Module):
    def __init__(self, num_categories: int, input_dim: int, hidden_dim: int, output_dim: int):
        super().__init__()
        self.layer1 = CategorySpecificLinear(num_categories, input_dim, hidden_dim)
        self.layer2 = CategorySpecificLinear(num_categories, hidden_dim, output_dim)

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


class MultiEmbodimentActionEncoder(nn.Module):
    def __init__(self, action_dim: int, hidden_size: int, num_embodiments: int):
        super().__init__()
        self.W1 = CategorySpecificLinear(num_embodiments, action_dim, hidden_size)
        self.W2 = CategorySpecificLinear(num_embodiments, 2 * hidden_size, hidden_size)
        self.W3 = CategorySpecificLinear(num_embodiments, hidden_size, hidden_size)
        self.pos_encoding = SinusoidalPositionalEncoding(hidden_size)

    def forward(
        self,
        actions: torch.Tensor,
        timesteps: torch.Tensor,
        embodiment_id: torch.Tensor,
    ) -> torch.Tensor:
        batch_size, horizon, _ = actions.shape
        if timesteps.dim() != 1 or timesteps.shape[0] != batch_size:
            raise ValueError("timesteps must have shape (B,)")
        timesteps = timesteps.unsqueeze(1).expand(-1, horizon)
        action_embedding = self.W1(actions, embodiment_id)
        time_embedding = self.pos_encoding(timesteps).to(dtype=action_embedding.dtype)
        hidden = self.W2(torch.cat([action_embedding, time_embedding], dim=-1), embodiment_id)
        return self.W3(hidden * torch.sigmoid(hidden), embodiment_id)


class MiniCPMV_VLA_ActionHead(nn.Module):
    """80-D, 32-embodiment action head for MiniCPM-VLA."""

    def __init__(
        self,
        hidden_size: int = 1024,
        action_dim: int = 80,
        state_dim: int = 80,
        action_horizon: int = 30,
        num_inference_timesteps: int = 4,
        num_target_vision_tokens: int = 32,
        max_seq_len: int = 1024,
        num_timestep_buckets: int = 1000,
        max_num_embodiments: int = 32,
    ):
        super().__init__()
        self.hidden_size = hidden_size
        self.input_embedding_dim = 768
        self.model = DiT(
            input_embedding_dim=768,
            attention_head_dim=64,
            num_attention_heads=12,
            cross_attention_dim=1024,
            dropout=0.2,
            final_dropout=True,
            interleave_self_attention=True,
            norm_type="ada_norm",
            num_layers=16,
            output_dim=1024,
            positional_embeddings=None,
        )
        self.action_dim = action_dim
        self.state_dim = state_dim
        self.action_horizon = action_horizon
        self.num_inference_timesteps = num_inference_timesteps
        self.max_num_embodiments = max_num_embodiments
        self.multi_embodiment = True
        self.proprio_inject = "concat"
        self.state_encoder = None
        self.action_encoder = MultiEmbodimentActionEncoder(
            action_dim=action_dim + state_dim,
            hidden_size=self.input_embedding_dim,
            num_embodiments=max_num_embodiments,
        )
        self.action_decoder = CategorySpecificMLP(
            num_categories=max_num_embodiments,
            input_dim=self.model.config.output_dim,
            hidden_dim=hidden_size,
            output_dim=action_dim,
        )
        self.future_tokens = nn.Embedding(num_target_vision_tokens, self.input_embedding_dim)
        nn.init.normal_(self.future_tokens.weight, mean=0.0, std=0.02)
        self.position_embedding = nn.Embedding(max_seq_len, self.input_embedding_dim)
        nn.init.normal_(self.position_embedding.weight, mean=0.0, std=0.02)
        self.num_timestep_buckets = num_timestep_buckets

    def _encode_action_tokens(
        self,
        noisy_actions: torch.Tensor,
        state: torch.Tensor,
        timesteps: torch.Tensor,
        embodiment_id: torch.Tensor,
    ) -> torch.Tensor:
        if state is None:
            raise ValueError("state is required because PROPRIO_INJECT=concat")
        state = state.expand(-1, noisy_actions.shape[1], -1)
        inputs = torch.cat([noisy_actions, state], dim=-1)
        return self.action_encoder(inputs, timesteps, embodiment_id)

    def _build_sequence(self, action_features: torch.Tensor) -> torch.Tensor:
        future_tokens = self.future_tokens.weight.unsqueeze(0).expand(action_features.shape[0], -1, -1)
        return torch.cat((future_tokens, action_features), dim=1)

    def _predict(
        self,
        noisy_actions: torch.Tensor,
        vl_embs: torch.Tensor,
        state: torch.Tensor,
        timesteps: torch.Tensor,
        embodiment_id: torch.Tensor,
    ) -> torch.Tensor:
        features = self._encode_action_tokens(noisy_actions, state, timesteps, embodiment_id)
        position_ids = torch.arange(features.shape[1], dtype=torch.long, device=features.device)
        features = features + self.position_embedding(position_ids).unsqueeze(0)
        output = self.model(
            hidden_states=self._build_sequence(features),
            encoder_hidden_states=vl_embs,
            timestep=timesteps,
        )
        return self.action_decoder(output, embodiment_id)[:, -self.action_horizon :]

    @torch.no_grad()
    def predict_action(
        self,
        vl_embs: torch.Tensor,
        state: torch.Tensor,
        embodiment_id: torch.Tensor,
    ) -> torch.Tensor:
        actions = torch.zeros(
            (vl_embs.shape[0], self.action_horizon, self.action_dim),
            dtype=vl_embs.dtype,
            device=vl_embs.device,
        )
        noise = torch.randn_like(actions)
        for step in range(self.num_inference_timesteps, 0, -1):
            time = step / float(self.num_inference_timesteps)
            timestep = min(int(time * self.num_timestep_buckets), self.num_timestep_buckets - 1)
            timesteps = torch.full(
                (vl_embs.shape[0],),
                timestep,
                device=vl_embs.device,
                dtype=torch.long,
            )
            noisy_actions = time * noise + (1 - time) * actions
            actions = self._predict(noisy_actions, vl_embs, state, timesteps, embodiment_id)
        return actions