def standard_attention(Q, K, V):
    """
    Q, K, V: [batch, seq_len, dim]
    복잡도: O(n² × d)
    """
    d_k = Q.shape[-1]

    # Step 1: QK^T 계산 - O(n² × d)
    scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(d_k)
    # scores: [batch, n, n] - n²개의 원소!

    # Step 2: Softmax - O(n²)
    attn_weights = F.softmax(scores, dim=-1)

    # Step 3: V와 곱 - O(n² × d)
    output = torch.matmul(attn_weights, V)

    return output

# 1024×1024 이미지의 경우:
# n = (1024/16)² = 4096 패치
# n² = 16,777,216 연산!

class LinearAttention(nn.Module):
    """
    SANA의 Linear Attention 구현
    """
    def __init__(self, dim, num_heads=8, qk_dim=64):
        super().__init__()
        self.num_heads = num_heads
        self.qk_dim = qk_dim
        self.scale = qk_dim ** -0.5

        # Q, K는 낮은 차원으로 projection
        self.q_proj = nn.Linear(dim, num_heads * qk_dim)
        self.k_proj = nn.Linear(dim, num_heads * qk_dim)
        self.v_proj = nn.Linear(dim, dim)
        self.out_proj = nn.Linear(dim, dim)

        # Feature map (kernel function)
        self.feature_map = nn.Sequential(
            nn.Linear(qk_dim, qk_dim),
            nn.ReLU()  # φ(x) = ReLU(Wx)
        )

    def forward(self, x):
        B, N, C = x.shape

        # Q, K, V 계산
        q = self.q_proj(x).view(B, N, self.num_heads, self.qk_dim)
        k = self.k_proj(x).view(B, N, self.num_heads, self.qk_dim)
        v = self.v_proj(x).view(B, N, self.num_heads, -1)

        # Feature map 적용 (커널 근사)
        q = self.feature_map(q)  # φ(Q)
        k = self.feature_map(k)  # φ(K)

        # Linear Attention: Q × (K^T × V)
        # Step 1: K^T × V - O(n × d_k × d_v)
        kv = torch.einsum('bnhk,bnhv->bhkv', k, v)

        # Step 2: Q × (K^T × V) - O(n × d_k × d_v)
        out = torch.einsum('bnhk,bhkv->bnhv', q, kv)

        # Normalization (numerical stability)
        normalizer = torch.einsum('bnhk,bhk->bnh', q, k.sum(dim=1))
        out = out / (normalizer.unsqueeze(-1) + 1e-6)

        # Reshape and project
        out = out.reshape(B, N, -1)
        out = self.out_proj(out)

        return out

class DeepCompressionAutoEncoder(nn.Module):
    """
    SANA의 32× 압축 AutoEncoder
    """
    def __init__(
        self,
        in_channels=3,
        latent_channels=32,
        base_channels=128
    ):
        super().__init__()

        # Encoder: 32× 다운샘플링 (5번의 2× 다운샘플)
        self.encoder = nn.Sequential(
            # 1024 → 512
            ConvBlock(in_channels, base_channels, stride=2),
            ResBlock(base_channels),

            # 512 → 256
            ConvBlock(base_channels, base_channels * 2, stride=2),
            ResBlock(base_channels * 2),

            # 256 → 128
            ConvBlock(base_channels * 2, base_channels * 4, stride=2),
            ResBlock(base_channels * 4),

            # 128 → 64
            ConvBlock(base_channels * 4, base_channels * 8, stride=2),
            ResBlock(base_channels * 8),

            # 64 → 32
            ConvBlock(base_channels * 8, latent_channels, stride=2),
        )

        # Decoder: 32× 업샘플링
        self.decoder = nn.Sequential(
            # 32 → 64
            UpConvBlock(latent_channels, base_channels * 8),
            ResBlock(base_channels * 8),

            # 64 → 128
            UpConvBlock(base_channels * 8, base_channels * 4),
            ResBlock(base_channels * 4),

            # 128 → 256
            UpConvBlock(base_channels * 4, base_channels * 2),
            ResBlock(base_channels * 2),

            # 256 → 512
            UpConvBlock(base_channels * 2, base_channels),
            ResBlock(base_channels),

            # 512 → 1024
            UpConvBlock(base_channels, in_channels),
        )

    def encode(self, x):
        return self.encoder(x)

    def decode(self, z):
        return self.decoder(z)

    def forward(self, x):
        z = self.encode(x)
        recon = self.decode(z)
        return recon, z

class DCAutoEncoderLoss(nn.Module):
    """
    DC-AE 학습을 위한 다중 손실 함수
    """
    def __init__(self):
        super().__init__()
        self.perceptual = LPIPS()
        self.discriminator = PatchGAN()

    def forward(self, x, recon, z):
        # 1. Reconstruction Loss
        l1_loss = F.l1_loss(recon, x)

        # 2. Perceptual Loss (더 중요!)
        perceptual_loss = self.perceptual(recon, x)

        # 3. Adversarial Loss
        real_pred = self.discriminator(x)
        fake_pred = self.discriminator(recon)
        adv_loss = F.binary_cross_entropy_with_logits(
            fake_pred, torch.ones_like(fake_pred)
        )

        # 4. Latent Regularization (KL divergence)
        kl_loss = 0.5 * (z.pow(2) - 1).mean()

        # 가중치 조합
        total_loss = (
            l1_loss * 1.0 +
            perceptual_loss * 0.5 +
            adv_loss * 0.1 +
            kl_loss * 0.0001
        )

        return total_loss

class MixFFN(nn.Module):
    """
    Mix-FFN: FFN에 Depthwise Convolution 추가
    지역적 정보와 전역적 정보를 동시에 처리
    """
    def __init__(self, dim, hidden_dim=None, kernel_size=3):
        super().__init__()
        hidden_dim = hidden_dim or dim * 4

        self.fc1 = nn.Linear(dim, hidden_dim)

        # Depthwise Convolution: 지역적 정보 처리
        self.dwconv = nn.Conv2d(
            hidden_dim, hidden_dim,
            kernel_size=kernel_size,
            padding=kernel_size // 2,
            groups=hidden_dim  # Depthwise!
        )

        self.act = nn.GELU()
        self.fc2 = nn.Linear(hidden_dim, dim)

    def forward(self, x, H, W):
        """
        x: [B, N, C] where N = H × W
        """
        B, N, C = x.shape

        # Linear projection
        x = self.fc1(x)

        # Reshape for conv: [B, N, C] → [B, C, H, W]
        x = x.transpose(1, 2).view(B, -1, H, W)

        # Depthwise convolution (지역적 패턴)
        x = self.dwconv(x)
        x = self.act(x)

        # Reshape back: [B, C, H, W] → [B, N, C]
        x = x.flatten(2).transpose(1, 2)

        # Final projection
        x = self.fc2(x)

        return x

class SANABlock(nn.Module):
    """
    SANA Transformer Block:
    Linear Attention + Mix-FFN + AdaLN
    """
    def __init__(self, dim, num_heads, mlp_ratio=4, qk_dim=64):
        super().__init__()

        # Normalization
        self.norm1 = nn.LayerNorm(dim, elementwise_affine=False)
        self.norm2 = nn.LayerNorm(dim, elementwise_affine=False)

        # Linear Attention
        self.attn = LinearAttention(dim, num_heads, qk_dim)

        # Mix-FFN
        self.ffn = MixFFN(dim, dim * mlp_ratio)

        # AdaLN modulation
        self.adaLN_modulation = nn.Sequential(
            nn.SiLU(),
            nn.Linear(dim, 6 * dim)
        )

    def forward(self, x, c, H, W):
        """
        x: [B, N, C] - 패치 토큰
        c: [B, C] - 조건 임베딩 (timestep + text)
        """
        # AdaLN parameters
        shift_attn, scale_attn, gate_attn, shift_ffn, scale_ffn, gate_ffn = \
            self.adaLN_modulation(c).unsqueeze(1).chunk(6, dim=-1)

        # Linear Attention with AdaLN
        x_norm = self.norm1(x) * (1 + scale_attn) + shift_attn
        x = x + gate_attn * self.attn(x_norm)

        # Mix-FFN with AdaLN
        x_norm = self.norm2(x) * (1 + scale_ffn) + shift_ffn
        x = x + gate_ffn * self.ffn(x_norm, H, W)

        return x

class SANA(nn.Module):
    """
    SANA: Linear Diffusion Transformer for High-Resolution Image Synthesis
    """
    def __init__(
        self,
        image_size=1024,
        patch_size=32,  # DC-AE compression
        latent_channels=32,
        hidden_dim=1152,
        depth=28,
        num_heads=16,
        mlp_ratio=4,
        text_dim=4096  # T5-XXL
    ):
        super().__init__()

        self.image_size = image_size
        self.patch_size = patch_size
        self.num_patches = (image_size // patch_size) ** 2  # 1024개

        # Patch embedding
        self.patch_embed = nn.Linear(latent_channels * patch_size * patch_size, hidden_dim)

        # Positional embedding
        self.pos_embed = nn.Parameter(torch.zeros(1, self.num_patches, hidden_dim))

        # Timestep embedding
        self.time_embed = nn.Sequential(
            SinusoidalPosEmb(hidden_dim),
            nn.Linear(hidden_dim, hidden_dim),
            nn.SiLU(),
            nn.Linear(hidden_dim, hidden_dim)
        )

        # Text projection
        self.text_proj = nn.Linear(text_dim, hidden_dim)

        # SANA Blocks
        self.blocks = nn.ModuleList([
            SANABlock(hidden_dim, num_heads, mlp_ratio)
            for _ in range(depth)
        ])

        # Cross-Attention blocks (매 4번째 블록)
        self.cross_attn_blocks = nn.ModuleList([
            CrossAttention(hidden_dim, num_heads)
            for _ in range(depth // 4)
        ])

        # Final layer
        self.final_norm = nn.LayerNorm(hidden_dim)
        self.final_proj = nn.Linear(hidden_dim, latent_channels * patch_size * patch_size)

    def forward(self, z, t, text_emb, text_mask=None):
        """
        z: [B, C, H, W] - 잠재 표현
        t: [B] - timestep
        text_emb: [B, L, D] - 텍스트 임베딩
        """
        B, C, H, W = z.shape

        # Patchify
        x = z.view(B, C, H // self.patch_size, self.patch_size,
                      W // self.patch_size, self.patch_size)
        x = x.permute(0, 2, 4, 1, 3, 5).reshape(B, -1, C * self.patch_size ** 2)
        x = self.patch_embed(x)

        # Add positional embedding
        x = x + self.pos_embed

        # Timestep conditioning
        t_emb = self.time_embed(t)

        # Text pooling for AdaLN
        text_pooled = text_emb.mean(dim=1)
        c = t_emb + self.text_proj(text_pooled)

        # Process through blocks
        patch_H, patch_W = H // self.patch_size, W // self.patch_size
        cross_attn_idx = 0

        for i, block in enumerate(self.blocks):
            x = block(x, c, patch_H, patch_W)

            # Cross-attention every 4 blocks
            if (i + 1) % 4 == 0 and cross_attn_idx < len(self.cross_attn_blocks):
                x = self.cross_attn_blocks[cross_attn_idx](x, text_emb, text_mask)
                cross_attn_idx += 1

        # Final projection
        x = self.final_norm(x)
        x = self.final_proj(x)

        # Unpatchify
        x = x.view(B, patch_H, patch_W, C, self.patch_size, self.patch_size)
        x = x.permute(0, 3, 1, 4, 2, 5).reshape(B, C, H, W)

        return x

class CrossAttention(nn.Module):
    """
    텍스트 조건을 위한 Cross-Attention
    (Linear Attention이 아닌 표준 Attention 사용)
    """
    def __init__(self, dim, num_heads):
        super().__init__()
        self.num_heads = num_heads
        self.head_dim = dim // num_heads
        self.scale = self.head_dim ** -0.5

        self.norm = nn.LayerNorm(dim)
        self.q_proj = nn.Linear(dim, dim)
        self.kv_proj = nn.Linear(dim, dim * 2)
        self.out_proj = nn.Linear(dim, dim)

    def forward(self, x, text_emb, text_mask=None):
        B, N, C = x.shape
        _, L, _ = text_emb.shape

        # Normalize
        x_norm = self.norm(x)

        # Q from image, K,V from text
        q = self.q_proj(x_norm).view(B, N, self.num_heads, self.head_dim)
        kv = self.kv_proj(text_emb).view(B, L, 2, self.num_heads, self.head_dim)
        k, v = kv.unbind(dim=2)

        # Attention
        q, k, v = q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2)
        attn = (q @ k.transpose(-2, -1)) * self.scale

        if text_mask is not None:
            attn = attn.masked_fill(~text_mask.unsqueeze(1).unsqueeze(2), float('-inf'))

        attn = attn.softmax(dim=-1)
        out = (attn @ v).transpose(1, 2).reshape(B, N, C)

        return x + self.out_proj(out)

import triton
import triton.language as tl

@triton.jit
def linear_attention_kernel(
    Q_ptr, K_ptr, V_ptr, O_ptr,
    stride_qb, stride_qh, stride_qn, stride_qk,
    stride_kb, stride_kh, stride_kn, stride_kk,
    stride_vb, stride_vh, stride_vn, stride_vd,
    stride_ob, stride_oh, stride_on, stride_od,
    N, D_K, D_V,
    BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr, BLOCK_V: tl.constexpr
):
    """
    Triton 커널로 구현한 Linear Attention
    메모리 효율적인 online computation
    """
    # Block indices
    batch_idx = tl.program_id(0)
    head_idx = tl.program_id(1)

    # Initialize accumulators for KV
    kv_acc = tl.zeros([BLOCK_K, BLOCK_V], dtype=tl.float32)
    k_sum = tl.zeros([BLOCK_K], dtype=tl.float32)

    # First pass: compute K^T @ V
    for n_start in range(0, N, BLOCK_N):
        n_offs = n_start + tl.arange(0, BLOCK_N)
        mask_n = n_offs < N

        # Load K and V
        k_ptrs = K_ptr + batch_idx * stride_kb + head_idx * stride_kh + \
                 n_offs[:, None] * stride_kn + tl.arange(0, BLOCK_K)[None, :] * stride_kk
        v_ptrs = V_ptr + batch_idx * stride_vb + head_idx * stride_vh + \
                 n_offs[:, None] * stride_vn + tl.arange(0, BLOCK_V)[None, :] * stride_vd

        k = tl.load(k_ptrs, mask=mask_n[:, None], other=0.0)
        v = tl.load(v_ptrs, mask=mask_n[:, None], other=0.0)

        # Apply feature map (ReLU)
        k = tl.maximum(k, 0)

        # Accumulate KV
        kv_acc += tl.dot(k.T, v)
        k_sum += tl.sum(k, axis=0)

    # Second pass: Q @ (K^T @ V)
    for n_start in range(0, N, BLOCK_N):
        n_offs = n_start + tl.arange(0, BLOCK_N)
        mask_n = n_offs < N

        # Load Q
        q_ptrs = Q_ptr + batch_idx * stride_qb + head_idx * stride_qh + \
                 n_offs[:, None] * stride_qn + tl.arange(0, BLOCK_K)[None, :] * stride_qk
        q = tl.load(q_ptrs, mask=mask_n[:, None], other=0.0)
        q = tl.maximum(q, 0)  # Feature map

        # Compute output
        out = tl.dot(q, kv_acc)

        # Normalize
        normalizer = tl.dot(q, k_sum[:, None]) + 1e-6
        out = out / normalizer

        # Store
        o_ptrs = O_ptr + batch_idx * stride_ob + head_idx * stride_oh + \
                 n_offs[:, None] * stride_on + tl.arange(0, BLOCK_V)[None, :] * stride_od
        tl.store(o_ptrs, out, mask=mask_n[:, None])

def benchmark_attention():
    """
    Linear Attention 구현별 성능 비교
    """
    B, N, H, D = 4, 4096, 16, 64  # 1024×1024 이미지

    results = {}

    # 1. Standard Attention (baseline)
    standard_time = measure_time(standard_attention, B, N, H, D)
    results["Standard Attention"] = standard_time

    # 2. PyTorch Linear Attention
    pytorch_linear_time = measure_time(pytorch_linear_attention, B, N, H, D)
    results["PyTorch Linear"] = pytorch_linear_time

    # 3. Triton Linear Attention
    triton_linear_time = measure_time(triton_linear_attention, B, N, H, D)
    results["Triton Linear"] = triton_linear_time

    return results

# 결과 (A100 기준):
# Standard Attention: 15.2ms
# PyTorch Linear: 4.8ms (3.2x faster)
# Triton Linear: 2.1ms (7.2x faster)

class SANATrainer:
    def __init__(self, config):
        self.model = SANA(**config.model)
        self.dc_ae = DeepCompressionAutoEncoder()
        self.text_encoder = T5EncoderModel.from_pretrained("google/flan-t5-xxl")

        self.optimizer = torch.optim.AdamW(
            self.model.parameters(),
            lr=config.lr,
            betas=(0.9, 0.999),
            weight_decay=0.01
        )

        self.scheduler = DDPMScheduler(
            num_train_timesteps=1000,
            beta_schedule="scaled_linear"
        )

    def train_step(self, images, captions):
        # 1. Encode images with DC-AE
        with torch.no_grad():
            latents = self.dc_ae.encode(images)

        # 2. Encode text
        with torch.no_grad():
            text_emb = self.text_encoder(captions).last_hidden_state

        # 3. Sample noise and timesteps
        noise = torch.randn_like(latents)
        timesteps = torch.randint(0, 1000, (images.shape[0],))

        # 4. Add noise
        noisy_latents = self.scheduler.add_noise(latents, noise, timesteps)

        # 5. Predict noise
        pred_noise = self.model(noisy_latents, timesteps, text_emb)

        # 6. Compute loss
        loss = F.mse_loss(pred_noise, noise)

        # 7. Backprop
        self.optimizer.zero_grad()
        loss.backward()
        torch.nn.utils.clip_grad_norm_(self.model.parameters(), 1.0)
        self.optimizer.step()

        return loss.item()

class SANAPipeline:
    def __init__(self, model_path):
        self.model = SANA.from_pretrained(model_path)
        self.dc_ae = DeepCompressionAutoEncoder.from_pretrained(model_path)
        self.text_encoder = T5EncoderModel.from_pretrained("google/flan-t5-xxl")
        self.scheduler = DDIMScheduler(num_train_timesteps=1000)

        # Compile for speed (PyTorch 2.0)
        self.model = torch.compile(self.model, mode="max-autotune")

    @torch.no_grad()
    def generate(
        self,
        prompt: str,
        negative_prompt: str = "",
        num_inference_steps: int = 20,  # DDIM은 적은 스텝 가능
        guidance_scale: float = 7.5,
        height: int = 1024,
        width: int = 1024,
        seed: int = None
    ):
        if seed is not None:
            torch.manual_seed(seed)

        # Text encoding
        text_emb = self.encode_text(prompt)
        if guidance_scale > 1.0:
            uncond_emb = self.encode_text(negative_prompt)
            text_emb = torch.cat([uncond_emb, text_emb])

        # Initial noise (32× 압축된 크기)
        latent_h, latent_w = height // 32, width // 32
        latents = torch.randn(1, 32, latent_h, latent_w, device="cuda")

        # DDIM denoising
        self.scheduler.set_timesteps(num_inference_steps)

        for t in self.scheduler.timesteps:
            latent_input = torch.cat([latents] * 2) if guidance_scale > 1.0 else latents

            # Predict noise
            noise_pred = self.model(latent_input, t, text_emb)

            # CFG
            if guidance_scale > 1.0:
                uncond_pred, cond_pred = noise_pred.chunk(2)
                noise_pred = uncond_pred + guidance_scale * (cond_pred - uncond_pred)

            # DDIM step
            latents = self.scheduler.step(noise_pred, t, latents).prev_sample

        # Decode with DC-AE
        images = self.dc_ae.decode(latents)
        images = (images / 2 + 0.5).clamp(0, 1)

        return images

    def encode_text(self, text):
        tokens = self.tokenizer(text, return_tensors="pt", padding=True, truncation=True)
        return self.text_encoder(**tokens.to("cuda")).last_hidden_state

future_directions = {
    "adaptive_compression": {
        "idea": "영역별 다른 압축률 적용",
        "benefit": "중요 영역은 높은 해상도 유지"
    },
    "hybrid_attention": {
        "idea": "Linear + Standard Attention 동적 전환",
        "benefit": "효율성과 표현력의 균형"
    },
    "video_extension": {
        "idea": "시간 축 Linear Attention",
        "benefit": "초고속 비디오 생성"
    },
    "distillation": {
        "idea": "더 작은 모델로 지식 증류",
        "benefit": "모바일/엣지 디바이스 배포"
    }
}