import math

import torch
import torch.nn as nn


class Model(nn.Module):
    """
    Scaled low-precision attention with explicit dequantization semantics.

    The inputs are quantized tensors plus per-head scales. This mirrors the data
    path of FP8-style scaled matmul attention even though the reference computes
    everything through explicit dequantization in PyTorch.
    """

    def __init__(self):
        super().__init__()

    def forward(
        self,
        q_q: torch.Tensor,
        k_q: torch.Tensor,
        v_q: torch.Tensor,
        q_scale: torch.Tensor,
        k_scale: torch.Tensor,
        v_scale: torch.Tensor,
    ) -> torch.Tensor:
        q = q_q.float() * q_scale.view(1, 1, -1, 1)
        k = k_q.float() * k_scale.view(1, 1, -1, 1)
        v = v_q.float() * v_scale.view(1, 1, -1, 1)

        batch_size, seq_len, num_heads, head_dim = q.shape
        out = torch.zeros_like(q)
        scale = 1.0 / math.sqrt(head_dim)

        for b in range(batch_size):
            for h in range(num_heads):
                for t in range(seq_len):
                    q_t = q[b, t, h]
                    k_hist = k[b, : t + 1, h]
                    v_hist = v[b, : t + 1, h]
                    scores = (k_hist * q_t.unsqueeze(0)).sum(dim=-1) * scale
                    attn = torch.softmax(scores, dim=0)
                    out[b, t, h] = (attn.unsqueeze(-1) * v_hist).sum(dim=0)

        return out


batch_size = 2
seq_len = 72
num_heads = 4
head_dim = 32


def _quantize_per_head(x: torch.Tensor):
    scale = x.abs().amax(dim=(0, 1, 3)).clamp(min=1e-4) / 127.0
    q = torch.clamp(torch.round(x / scale.view(1, 1, -1, 1)), -127, 127).to(torch.int8)
    return q, scale.to(torch.float32)


def get_inputs():
    p = popcorn_pri
    q = torch.randn(p.jitter_int(batch_size), p.jitter_int(seq_len), p.jitter_int(num_heads), p.jitter_int(head_dim, align=8), dtype=torch.float32)
    k = torch.randn(p.jitter_int(batch_size), p.jitter_int(seq_len), p.jitter_int(num_heads), p.jitter_int(head_dim, align=8), dtype=torch.float32)
    v = torch.randn(p.jitter_int(batch_size), p.jitter_int(seq_len), p.jitter_int(num_heads), p.jitter_int(head_dim, align=8), dtype=torch.float32)
    q_q, q_scale = _quantize_per_head(q)
    k_q, k_scale = _quantize_per_head(k)
    v_q, v_scale = _quantize_per_head(v)
    return [q_q, k_q, v_q, q_scale, k_scale, v_scale]


def get_init_inputs():
    return []

import math
import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline

cuda_source = """
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <cfloat>

__global__ void causal_attention_kernel(
    const float* __restrict__ q,
    const float* __restrict__ k,
    const float* __restrict__ v,
    float* __restrict__ out,
    int batch_size, int seq_len, int num_heads, int head_dim, float scale
) {
    // Each block handles one (batch, head) pair
    int bh = blockIdx.x;
    int b = bh / num_heads;
    int h = bh % num_heads;
    
    int tid = threadIdx.x;
    int block_size = blockDim.x;
    
    // Base offset for this batch and head
    int64_t qk_offset = (int64_t)b * seq_len * num_heads * head_dim + (int64_t)h * head_dim;
    
    // Shared memory for K and V (we'll load them tile by tile)
    extern __shared__ float shared_mem[];
    float* shared_k = shared_mem;                          // [TILE_SIZE][head_dim]
    float* shared_v = shared_mem + blockDim.x * head_dim; // [TILE_SIZE][head_dim]
    float* shared_scores = shared_v + blockDim.x * head_dim; // [TILE_SIZE]
    
    // Process each query position
    for (int t = tid; t < seq_len; t += block_size) {
        // Load query vector for position t
        float q_vec[32]; // assuming head_dim <= 32
        for (int d = 0; d < head_dim; d++) {
            q_vec[d] = q[qk_offset + (int64_t)t * num_heads * head_dim + d];
        }
        
        float max_score = -FLT_MAX;
        float sum_exp = 0.0f;
        float out_vec[32];
        for (int d = 0; d < head_dim; d++) out_vec[d] = 0.0f;
        
        // Process key/value positions 0 to t (causal)
        for (int kv_start = 0; kv_start <= t; kv_start += block_size) {
            int kv_end = min(kv_start + block_size, t + 1);
            int tile_size = kv_end - kv_start;
            
            // Collaborative loading of K and V tiles
            for (int kv_idx = tid; kv_idx < tile_size; kv_idx += block_size) {
                int kv_pos = kv_start + kv_idx;
                for (int d = 0; d < head_dim; d++) {
                    shared_k[kv_idx * head_dim + d] = k[qk_offset + (int64_t)kv_pos * num_heads * head_dim + d];
                    shared_v[kv_idx * head_dim + d] = v[qk_offset + (int64_t)kv_pos * num_heads * head_dim + d];
                }
            }
            __syncthreads();
            
            // Compute scores for this tile
            float local_max = -FLT_MAX;
            for (int kv_idx = 0; kv_idx < tile_size; kv_idx++) {
                float score = 0.0f;
                for (int d = 0; d < head_dim; d++) {
                    score += q_vec[d] * shared_k[kv_idx * head_dim + d];
                }
                score *= scale;
                shared_scores[kv_idx] = score;
                local_max = fmaxf(local_max, score);
            }
            
            // Online softmax update
            float new_max = fmaxf(max_score, local_max);
            float rescale = (max_score == -FLT_MAX) ? 0.0f : expf(max_score - new_max);
            float local_rescale = expf(local_max - new_max);
            
            // Update sum_exp and output accumulator
            sum_exp = sum_exp * rescale;
            for (int d = 0; d < head_dim; d++) {
                out_vec[d] = out_vec[d] * rescale;
            }
            
            for (int kv_idx = 0; kv_idx < tile_size; kv_idx++) {
                float exp_score = expf(shared_scores[kv_idx] - new_max);
                sum_exp += exp_score;
                for (int d = 0; d < head_dim; d++) {
                    out_vec[d] += exp_score * shared_v[kv_idx * head_dim + d];
                }
            }
            
            max_score = new_max;
            __syncthreads();
        }
        
        // Normalize and write output
        float inv_sum = 1.0f / sum_exp;
        for (int d = 0; d < head_dim; d++) {
            out[qk_offset + (int64_t)t * num_heads * head_dim + d] = out_vec[d] * inv_sum;
        }
    }
}

__global__ void dequantize_kernel(
    const int8_t* __restrict__ q_q,
    const int8_t* __restrict__ k_q,
    const int8_t* __restrict__ v_q,
    const float* __restrict__ q_scale,
    const float* __restrict__ k_scale,
    const float* __restrict__ v_scale,
    float* __restrict__ q,
    float* __restrict__ k,
    float* __restrict__ v,
    int batch_size, int seq_len, int num_heads, int head_dim
) {
    int64_t idx = (int64_t)blockIdx.x * blockDim.x + threadIdx.x;
    int64_t total = (int64_t)batch_size * seq_len * num_heads * head_dim;
    
    if (idx < total) {
        int64_t head_idx = (idx / head_dim) % num_heads;
        float qs = q_scale[head_idx];
        float ks = k_scale[head_idx];
        float vs = v_scale[head_idx];
        
        q[idx] = (float)q_q[idx] * qs;
        k[idx] = (float)k_q[idx] * ks;
        v[idx] = (float)v_q[idx] * vs;
    }
}

torch::Tensor causal_attention_cuda(
    torch::Tensor q_q,
    torch::Tensor k_q,
    torch::Tensor v_q,
    torch::Tensor q_scale,
    torch::Tensor k_scale,
    torch::Tensor v_scale
) {
    int batch_size = q_q.size(0);
    int seq_len = q_q.size(1);
    int num_heads = q_q.size(2);
    int head_dim = q_q.size(3);
    
    auto options = torch::TensorOptions().dtype(torch::kFloat32).device(q_q.device());
    auto q = torch::empty({batch_size, seq_len, num_heads, head_dim}, options);
    auto k = torch::empty({batch_size, seq_len, num_heads, head_dim}, options);
    auto v = torch::empty({batch_size, seq_len, num_heads, head_dim}, options);
    auto out = torch::empty({batch_size, seq_len, num_heads, head_dim}, options);
    
    // Dequantize
    int64_t total = (int64_t)batch_size * seq_len * num_heads * head_dim;
    int deq_threads = 256;
    int deq_blocks = (total + deq_threads - 1) / deq_threads;
    
    dequantize_kernel<<<deq_blocks, deq_threads>>>(
        q_q.data_ptr<int8_t>(),
        k_q.data_ptr<int8_t>(),
        v_q.data_ptr<int8_t>(),
        q_scale.data_ptr<float>(),
        k_scale.data_ptr<float>(),
        v_scale.data_ptr<float>(),
        q.data_ptr<float>(),
        k.data_ptr<float>(),
        v.data_ptr<float>(),
        batch_size, seq_len, num_heads, head_dim
    );
    
    // Causal attention
    float scale = 1.0f / sqrtf((float)head_dim);
    int num_bh = batch_size * num_heads;
    int threads_per_block = 128;
    size_t shared_mem_size = 2 * threads_per_block * head_dim * sizeof(float) + threads_per_block * sizeof(float);
    
    causal_attention_kernel<<<num_bh, threads_per_block, shared_mem_size>>>(
        q.data_ptr<float>(),
        k.data_ptr<float>(),
        v.data_ptr<float>(),
        out.data_ptr<float>(),
        batch_size, seq_len, num_heads, head_dim, scale
    );
    
    return out;
}
"""

cpp_source = """
torch::Tensor causal_attention_cuda(
    torch::Tensor q_q,
    torch::Tensor k_q,
    torch::Tensor v_q,
    torch::Tensor q_scale,
    torch::Tensor k_scale,
    torch::Tensor v_scale
);
"""

attention_module = load_inline(
    name="causal_attention",
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=["causal_attention_cuda"],
    extra_cuda_cflags=["-O3"],
    verbose=False
)


class ModelNew(nn.Module):
    def __init__(self):
        super().__init__()
        self.attention_op = attention_module.causal_attention_cuda
    
    def forward(self, q_q, k_q, v_q, q_scale, k_scale, v_scale):
        return self.attention_op(q_q, k_q, v_q, q_scale, k_scale, v_scale)