from __future__ import annotations

import math
from typing import Optional

import torch
import torch.nn.functional as F
from torch import Tensor, nn

from b1_tensor_core import B1Linear
from .config import DiffusionConfig
from .diffusion import sinusoidal_embedding


class B1LinearScaled(nn.Module):
    """
    1-Bit Tensor Core linear projection normalized by 1/sqrt(in_features)
    to maintain unit activation variance for continuous diffusion stability.
    """

    def __init__(self, in_features: int, out_features: int, bias: bool = True, backend: str = "tc"):
        super().__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.backend = backend
        self.linear = B1Linear(in_features, out_features, bias=bias, backend=backend)
        self.scale = 1.0 / math.sqrt(in_features)

    def freeze_for_inference(self):
        self.linear.freeze_for_inference()

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # B1Linear expects 2D [M, K]
        orig_shape = x.shape
        if x.ndim > 2:
            x_flat = x.reshape(-1, orig_shape[-1])
            out_flat = self.linear(x_flat)
            out = out_flat.reshape(*orig_shape[:-1], self.out_features)
        else:
            out = self.linear(x)
        return out * self.scale


class B1SelfAttention(nn.Module):
    def __init__(self, embed_dim: int, num_heads: int = 8, backend: str = "tc"):
        super().__init__()
        self.embed_dim = embed_dim
        self.num_heads = num_heads
        self.head_dim = embed_dim // num_heads
        self.scale = 1.0 / math.sqrt(self.head_dim)

        self.q_proj = B1LinearScaled(embed_dim, embed_dim, backend=backend)
        self.k_proj = B1LinearScaled(embed_dim, embed_dim, backend=backend)
        self.v_proj = B1LinearScaled(embed_dim, embed_dim, backend=backend)
        self.out_proj = B1LinearScaled(embed_dim, embed_dim, backend=backend)

    def freeze_for_inference(self):
        self.q_proj.freeze_for_inference()
        self.k_proj.freeze_for_inference()
        self.v_proj.freeze_for_inference()
        self.out_proj.freeze_for_inference()

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        B, L, D = x.shape
        q = self.q_proj(x).reshape(B, L, self.num_heads, self.head_dim).transpose(1, 2)
        k = self.k_proj(x).reshape(B, L, self.num_heads, self.head_dim).transpose(1, 2)
        v = self.v_proj(x).reshape(B, L, self.num_heads, self.head_dim).transpose(1, 2)

        scores = torch.matmul(q, k.transpose(-2, -1)) * self.scale
        attn = torch.softmax(scores, dim=-1)
        out = torch.matmul(attn, v)
        out = out.transpose(1, 2).reshape(B, L, D)
        return self.out_proj(out)


class B1CrossAttention(nn.Module):
    def __init__(self, embed_dim: int, num_heads: int = 8, backend: str = "tc"):
        super().__init__()
        self.embed_dim = embed_dim
        self.num_heads = num_heads
        self.head_dim = embed_dim // num_heads
        self.scale = 1.0 / math.sqrt(self.head_dim)

        self.q_proj = B1LinearScaled(embed_dim, embed_dim, backend=backend)
        self.k_proj = B1LinearScaled(embed_dim, embed_dim, backend=backend)
        self.v_proj = B1LinearScaled(embed_dim, embed_dim, backend=backend)
        self.out_proj = B1LinearScaled(embed_dim, embed_dim, backend=backend)

    def freeze_for_inference(self):
        self.q_proj.freeze_for_inference()
        self.k_proj.freeze_for_inference()
        self.v_proj.freeze_for_inference()
        self.out_proj.freeze_for_inference()

    def forward(self, x: torch.Tensor, memory: torch.Tensor) -> torch.Tensor:
        B, L, D = x.shape
        _, S, _ = memory.shape

        q = self.q_proj(x).reshape(B, L, self.num_heads, self.head_dim).transpose(1, 2)
        k = self.k_proj(memory).reshape(B, S, self.num_heads, self.head_dim).transpose(1, 2)
        v = self.v_proj(memory).reshape(B, S, self.num_heads, self.head_dim).transpose(1, 2)

        scores = torch.matmul(q, k.transpose(-2, -1)) * self.scale
        attn = torch.softmax(scores, dim=-1)
        out = torch.matmul(attn, v)
        out = out.transpose(1, 2).reshape(B, L, D)
        return self.out_proj(out)


class B1MLP(nn.Module):
    def __init__(self, embed_dim: int, mlp_dim: int, backend: str = "tc"):
        super().__init__()
        self.fc1 = B1LinearScaled(embed_dim, mlp_dim, backend=backend)
        self.act = nn.GELU()
        self.fc2 = B1LinearScaled(mlp_dim, embed_dim, backend=backend)

    def freeze_for_inference(self):
        self.fc1.freeze_for_inference()
        self.fc2.freeze_for_inference()

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.fc2(self.act(self.fc1(x)))


class B1TransformerDecoderLayer(nn.Module):
    def __init__(self, embed_dim: int, num_heads: int, mlp_dim: int, backend: str = "tc"):
        super().__init__()
        self.norm1 = nn.LayerNorm(embed_dim)
        self.self_attn = B1SelfAttention(embed_dim, num_heads, backend=backend)
        self.norm2 = nn.LayerNorm(embed_dim)
        self.cross_attn = B1CrossAttention(embed_dim, num_heads, backend=backend)
        self.norm3 = nn.LayerNorm(embed_dim)
        self.mlp = B1MLP(embed_dim, mlp_dim, backend=backend)

    def freeze_for_inference(self):
        self.self_attn.freeze_for_inference()
        self.cross_attn.freeze_for_inference()
        self.mlp.freeze_for_inference()

    def forward(self, x: torch.Tensor, memory: torch.Tensor) -> torch.Tensor:
        # Pre-LN architecture
        x = x + self.self_attn(self.norm1(x))
        x = x + self.cross_attn(self.norm2(x), memory)
        x = x + self.mlp(self.norm3(x))
        return x


class B1CoherentTransformer(nn.Module):
    """
    Sub-Byte 1-Bit Tensor Core continuous transformer backbone.
    All internal attention and feedforward projections execute via inline PTX MMA.
    """

    def __init__(self, config: DiffusionConfig, backend: str = "tc", pure_1bit: bool = False) -> None:
        super().__init__()
        self.config = config
        self.pure_1bit = pure_1bit
        if pure_1bit:
            self.input_projection = B1LinearScaled(config.embedding_dim, config.hidden_dim, backend=backend)
            self.output_projection = B1LinearScaled(config.hidden_dim, config.embedding_dim, backend=backend)
            self.query_projection = B1LinearScaled(config.embedding_dim, config.hidden_dim, backend=backend)
            self.time_mlp = nn.Sequential(
                B1LinearScaled(config.hidden_dim, config.hidden_dim, backend=backend),
                nn.SiLU(),
                B1LinearScaled(config.hidden_dim, config.hidden_dim, backend=backend),
            )
        else:
            self.input_projection = nn.Linear(config.embedding_dim, config.hidden_dim)
            self.output_projection = nn.Linear(config.hidden_dim, config.embedding_dim)
            self.query_projection = nn.Linear(config.embedding_dim, config.hidden_dim)
            self.time_mlp = nn.Sequential(
                nn.Linear(config.hidden_dim, config.hidden_dim),
                nn.SiLU(),
                nn.Linear(config.hidden_dim, config.hidden_dim),
            )
        self.position = nn.Parameter(torch.zeros(1, config.sequence_length, config.hidden_dim))
        self.layers = nn.ModuleList([
            B1TransformerDecoderLayer(
                embed_dim=config.hidden_dim,
                num_heads=config.heads,
                mlp_dim=config.mlp_dim,
                backend=backend,
            )
            for _ in range(config.layers)
        ])
        self.final_norm = nn.LayerNorm(config.hidden_dim)
        nn.init.normal_(self.position, std=0.02)

    def freeze_for_inference(self):
        if self.pure_1bit:
            self.input_projection.freeze_for_inference()
            self.output_projection.freeze_for_inference()
            self.query_projection.freeze_for_inference()
            self.time_mlp[0].freeze_for_inference()
            self.time_mlp[2].freeze_for_inference()
        for layer in self.layers:
            layer.freeze_for_inference()

    def forward(
        self,
        noisy_targets: Tensor,
        noise_embedding: Tensor,
        query_embedding: Tensor,
        condition_drop_mask: Tensor | None = None,
    ) -> Tensor:
        batch, length, _ = noisy_targets.shape
        if length > self.config.sequence_length:
            raise ValueError("input sequence exceeds configured sequence_length")
        if condition_drop_mask is not None:
            query_embedding = query_embedding.masked_fill(condition_drop_mask[:, None], 0.0)
        tokens = self.input_projection(noisy_targets) + self.position[:, :length]
        tokens = tokens + self.time_mlp(noise_embedding)[:, None]
        memory = self.query_projection(query_embedding)[:, None]

        for layer in self.layers:
            tokens = layer(tokens, memory)

        return self.output_projection(self.final_norm(tokens))


class B1EDMDenoiser(nn.Module):
    """1-Bit Hardware Tensor Core EDM-preconditioned denoiser."""

    def __init__(self, config: DiffusionConfig, backend: str = "tc", pure_1bit: bool = False) -> None:
        super().__init__()
        self.config = config
        self.pure_1bit = pure_1bit
        self.backbone = B1CoherentTransformer(config, backend=backend, pure_1bit=pure_1bit)

    def freeze_for_inference(self):
        self.backbone.freeze_for_inference()

    def forward(
        self,
        noisy_targets: Tensor,
        sigma: Tensor,
        query_embedding: Tensor,
        condition_drop_mask: Tensor | None = None,
    ) -> Tensor:
        sigma_data = self.config.sigma_data
        sigma_sq = sigma.square()
        denominator = sigma_sq + sigma_data**2
        c_skip = sigma_data**2 / denominator
        c_out = sigma * sigma_data / denominator.sqrt()
        c_in = denominator.rsqrt()

        c_skip = c_skip.reshape(sigma.shape[0], 1, 1)
        c_out = c_out.reshape(sigma.shape[0], 1, 1)
        c_in = c_in.reshape(sigma.shape[0], 1, 1)

        noise_embedding = sinusoidal_embedding(sigma, self.config.hidden_dim)
        effective_query = (
            query_embedding
            if condition_drop_mask is None
            else query_embedding.masked_fill(condition_drop_mask[:, None], 0.0)
        )
        scaled_input = noisy_targets * c_in
        model_output = self.backbone(
            scaled_input,
            noise_embedding,
            effective_query,
            condition_drop_mask=None,
        )
        return c_skip * noisy_targets + c_out * model_output