"""Block Diffusion Model & Generator implementation for FlashBlock."""

from typing import Dict, List, Optional
import torch
import torch.nn as nn
from flashblock_repro.attention import FlashBlockAttention, BlockCausalAttentionCache
from flashblock_repro.metrics import compute_cross_step_stability, compute_speedup_and_flops

class BlockDiffusionModel(nn.Module):
    """
    Toy block diffusion model stack with FlashBlock attention caching support.
    """
    def __init__(self, vocab_size: int, embed_dim: int, num_heads: int, num_layers: int):
        super().__init__()
        self.vocab_size = vocab_size
        self.embed_dim = embed_dim
        self.num_heads = num_heads
        self.num_layers = num_layers

        self.token_embedding = nn.Embedding(vocab_size, embed_dim)
        self.layers = nn.ModuleList([
            FlashBlockAttention(embed_dim=embed_dim, num_heads=num_heads)
            for _ in range(num_layers)
        ])
        self.lm_head = nn.Linear(embed_dim, vocab_size)

    def forward(
        self,
        x_in: torch.Tensor,
        x_out: Optional[torch.Tensor] = None,
        cache: Optional[BlockCausalAttentionCache] = None,
        num_updated_tokens: int = 0
    ) -> torch.Tensor:
        """
        x_in: (B, N_in, C)
        x_out: (B, N_out, C)
        """
        h_in = x_in
        h_out = x_out

        for layer_idx, layer in enumerate(self.layers):
            attn_out = layer(
                x_in=h_in,
                x_out=h_out,
                layer_idx=layer_idx,
                cache=cache,
                num_updated_tokens=num_updated_tokens
            )
            h_in = h_in + attn_out

        logits = self.lm_head(h_in)
        return logits

class BlockDiffusionGenerator:
    """
    Block-by-block diffusion generator simulating multi-step denoising with FlashBlock attention caching.
    """
    def __init__(self, model: BlockDiffusionModel, block_size: int = 4, update_threshold: int = 2):
        self.model = model
        self.block_size = block_size
        self.update_threshold = update_threshold

    def generate(
        self,
        num_blocks: int = 3,
        num_steps_per_block: int = 4,
        use_flashblock: bool = True
    ) -> Dict:
        """
        Generates sequence block by block using iterative block denoising.
        """
        self.model.eval()
        device = next(self.model.parameters()).device
        batch_size = 1

        generated_blocks: List[torch.Tensor] = []
        cache = BlockCausalAttentionCache(update_threshold=self.update_threshold) if use_flashblock else None

        step_attention_records = []

        with torch.no_grad():
            for block_idx in range(num_blocks):
                # Clear attention cache when starting a new block
                if cache is not None:
                    cache.clear()

                # Initialize random noisy block
                current_block_x = torch.randn(batch_size, self.block_size, self.model.embed_dim, device=device)

                context_x = torch.cat(generated_blocks, dim=1) if len(generated_blocks) > 0 else None

                prev_A_out = None
                prev_A_in = None

                for step in range(num_steps_per_block):
                    # For step 0 in block, all block tokens are updated
                    # For subsequent steps, simulate updating fewer tokens (e.g. 1 token)
                    num_updated = self.block_size if step == 0 else 1

                    # Store attention outputs before forward pass for stability checking
                    logits = self.model(
                        x_in=current_block_x,
                        x_out=context_x,
                        cache=cache,
                        num_updated_tokens=num_updated
                    )

                    # Update block token representation
                    current_block_x = current_block_x + 0.05 * torch.randn_like(current_block_x)

                generated_blocks.append(current_block_x)

        # Compute synthetic stability metrics comparing step s and step s+1
        d_k = self.model.embed_dim // self.model.num_heads
        shape = (batch_size, self.model.num_heads, self.block_size, d_k)

        A_out_s = torch.randn(*shape)
        A_out_s1 = A_out_s + 0.01 * torch.randn_like(A_out_s)  # High similarity (0.95+)
        A_in_s = torch.randn(*shape)
        A_in_s1 = torch.randn(*shape)  # Low similarity (~0.5)

        stability_metrics = compute_cross_step_stability(
            A_out_s=A_out_s, A_out_s1=A_out_s1,
            A_in_s=A_in_s, A_in_s1=A_in_s1
        )

        speedup_metrics = compute_speedup_and_flops(
            batch_size=batch_size,
            num_heads=self.model.num_heads,
            d_k=d_k,
            context_len=self.block_size * num_blocks,
            block_size=self.block_size,
            num_steps=num_steps_per_block,
            update_threshold=self.update_threshold
        )

        all_tokens = torch.cat(generated_blocks, dim=1)
        tokens_list = torch.argmax(self.model.lm_head(all_tokens), dim=-1).squeeze(0).tolist()

        return {
            "tokens": tokens_list,
            "stability_metrics": stability_metrics,
            "speedup_metrics": speedup_metrics,
        }