import torch
import torch.nn as nn
import torch.nn.functional as F
import math


class Model(nn.Module):
    """
    Axial attention for 2-D feature maps.  Factorises full 2-D attention
    into sequential row-wise and column-wise attention passes, reducing
    complexity from O(N^2 * M^2) to O(N*M*(N+M)).  Used in bioimage
    analysis (cell segmentation, cryo-EM micrograph processing) and
    axial-attention protein language models.
    """

    def __init__(self, dim, num_heads, height, width, dropout=0.0):
        super().__init__()
        self.height = height
        self.width = width
        self.num_heads = num_heads
        self.head_dim = dim // num_heads
        assert dim % num_heads == 0

        self.row_norm = nn.LayerNorm(dim)
        self.row_qkv = nn.Linear(dim, 3 * dim, bias=False)
        self.row_out = nn.Linear(dim, dim)

        self.col_norm = nn.LayerNorm(dim)
        self.col_qkv = nn.Linear(dim, 3 * dim, bias=False)
        self.col_out = nn.Linear(dim, dim)

        self.dropout = nn.Dropout(dropout)

    def _attn(self, x, qkv_proj, out_proj, norm):
        """Standard multi-head self-attention over last-but-one dim."""
        B_outer, L, D = x.shape
        h, d = self.num_heads, self.head_dim
        x = norm(x)
        qkv = qkv_proj(x).view(B_outer, L, 3, h, d)
        q, k, v = qkv.unbind(dim=2)
        q = q.transpose(1, 2)
        k = k.transpose(1, 2)
        v = v.transpose(1, 2)
        attn = (q @ k.transpose(-1, -2)) / math.sqrt(d)
        attn = F.softmax(attn, dim=-1)
        attn = self.dropout(attn)
        out = (attn @ v).transpose(1, 2).reshape(B_outer, L, D)
        return out_proj(out)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Args:
            x: (B, H, W, D) – 2-D feature map
        Returns:
            (B, H, W, D) – updated feature map
        """
        B, H, W, D = x.shape

        # Row-wise attention: treat each row as a sequence
        x_row = x.reshape(B * H, W, D)
        x = x + self._attn(x_row, self.row_qkv, self.row_out, self.row_norm).reshape(B, H, W, D)

        # Column-wise attention: treat each column as a sequence
        x_col = x.permute(0, 2, 1, 3).reshape(B * W, H, D)
        x = x + self._attn(x_col, self.col_qkv, self.col_out, self.col_norm).reshape(B, W, H, D).permute(0, 2, 1, 3)

        return x


dim = 64
num_heads = 4
height = 32
width = 32
batch_size = 4


def get_inputs():
    return [torch.randn(popcorn_pri.jitter_int(batch_size), height, width, dim)]


def get_init_inputs():
    return [dim, num_heads, height, width]

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

cpp_src = r'''
#include <torch/extension.h>
torch::Tensor axial_forward(
    torch::Tensor x,
    torch::Tensor row_norm_w, torch::Tensor row_norm_b,
    torch::Tensor row_qkv_w, torch::Tensor row_out_w, torch::Tensor row_out_b,
    torch::Tensor col_norm_w, torch::Tensor col_norm_b,
    torch::Tensor col_qkv_w, torch::Tensor col_out_w, torch::Tensor col_out_b);
'''

cuda_src = r'''
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda.h>
#include <cuda_runtime.h>

#define DCONST 64
#define HCONST 4
#define HDIM 16

template<bool IS_COL>
__device__ __forceinline__ int off_calc(int b, int fixed, int token, int channel, int H, int W) {
    if constexpr (!IS_COL) {
        return ((b * H + fixed) * W + token) * DCONST + channel;
    } else {
        return ((b * H + token) * W + fixed) * DCONST + channel;
    }
}

template<bool IS_COL>
__global__ void axial_attn_kernel(
    const float* __restrict__ inp,
    float* __restrict__ out,
    const float* __restrict__ norm_w,
    const float* __restrict__ norm_b,
    const float* __restrict__ qkv_w,
    const float* __restrict__ out_w,
    const float* __restrict__ out_b,
    int B, int H, int W, int L)
{
    extern __shared__ float smem[];
    float* xnorm = smem;
    float* qkv = xnorm + L * DCONST;
    float* ctx = qkv + L * 3 * DCONST;
    int tid = threadIdx.x;
    int seq = blockIdx.x;
    int b, fixed;
    if constexpr (!IS_COL) { b = seq / H; fixed = seq - b * H; }
    else { b = seq / W; fixed = seq - b * W; }

    for (int t = tid; t < L; t += blockDim.x) {
        float mean = 0.f;
        #pragma unroll
        for (int c = 0; c < DCONST; ++c) mean += inp[off_calc<IS_COL>(b, fixed, t, c, H, W)];
        mean *= (1.0f / DCONST);
        float var = 0.f;
        #pragma unroll
        for (int c = 0; c < DCONST; ++c) {
            float v = inp[off_calc<IS_COL>(b, fixed, t, c, H, W)] - mean;
            var += v * v;
        }
        float rstd = rsqrtf(var * (1.0f / DCONST) + 1.0e-5f);
        #pragma unroll
        for (int c = 0; c < DCONST; ++c) {
            float v = (inp[off_calc<IS_COL>(b, fixed, t, c, H, W)] - mean) * rstd;
            xnorm[t * DCONST + c] = v * norm_w[c] + norm_b[c];
        }
    }
    __syncthreads();

    int total_qkv = L * 3 * DCONST;
    for (int idx = tid; idx < total_qkv; idx += blockDim.x) {
        int t = idx / (3 * DCONST);
        int o = idx - t * (3 * DCONST);
        const float* wrow = qkv_w + o * DCONST;
        const float* xrow = xnorm + t * DCONST;
        float sum = 0.f;
        #pragma unroll
        for (int c = 0; c < DCONST; ++c) sum = fmaf(xrow[c], wrow[c], sum);
        qkv[idx] = sum;
    }
    __syncthreads();

    for (int row = tid; row < L * HCONST; row += blockDim.x) {
        int t = row / HCONST;
        int h = row - t * HCONST;
        float scores[32];
        float maxv = -3.402823466e38f;
        const int base_q = t * 3 * DCONST + h * HDIM;
        for (int j = 0; j < L; ++j) {
            const int base_k = j * 3 * DCONST + DCONST + h * HDIM;
            float dot = 0.f;
            #pragma unroll
            for (int d = 0; d < HDIM; ++d) dot = fmaf(qkv[base_q + d], qkv[base_k + d], dot);
            float s = dot * 0.25f;
            scores[j] = s; maxv = fmaxf(maxv, s);
        }
        float denom = 0.f;
        for (int j = 0; j < L; ++j) { float e = __expf(scores[j] - maxv); scores[j] = e; denom += e; }
        float invden = 1.0f / denom;
        #pragma unroll
        for (int d = 0; d < HDIM; ++d) {
            float acc = 0.f;
            for (int j = 0; j < L; ++j) {
                int base_v = j * 3 * DCONST + 2 * DCONST + h * HDIM;
                acc = fmaf(scores[j] * invden, qkv[base_v + d], acc);
            }
            ctx[t * DCONST + h * HDIM + d] = acc;
        }
    }
    __syncthreads();

    int total = L * DCONST;
    for (int idx = tid; idx < total; idx += blockDim.x) {
        int t = idx / DCONST;
        int o = idx - t * DCONST;
        const float* wrow = out_w + o * DCONST;
        const float* crow = ctx + t * DCONST;
        float sum = out_b[o];
        #pragma unroll
        for (int c = 0; c < DCONST; ++c) sum = fmaf(crow[c], wrow[c], sum);
        int off = off_calc<IS_COL>(b, fixed, t, o, H, W);
        out[off] = inp[off] + sum;
    }
}

torch::Tensor axial_forward(
    torch::Tensor x,
    torch::Tensor row_norm_w, torch::Tensor row_norm_b,
    torch::Tensor row_qkv_w, torch::Tensor row_out_w, torch::Tensor row_out_b,
    torch::Tensor col_norm_w, torch::Tensor col_norm_b,
    torch::Tensor col_qkv_w, torch::Tensor col_out_w, torch::Tensor col_out_b)
{
    TORCH_CHECK(x.is_cuda(), "x must be CUDA");
    TORCH_CHECK(x.scalar_type() == at::kFloat, "only float32 supported");
    int B = (int)x.size(0), H = (int)x.size(1), W = (int)x.size(2), D = (int)x.size(3);
    TORCH_CHECK(D == DCONST, "specialized for D=64");
    auto temp = torch::empty_like(x);
    auto y = torch::empty_like(x);
    int threads = 256;
    size_t sh_row = (W * DCONST + W * 3 * DCONST + W * DCONST) * sizeof(float);
    size_t sh_col = (H * DCONST + H * 3 * DCONST + H * DCONST) * sizeof(float);
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    axial_attn_kernel<false><<<B * H, threads, sh_row, stream>>>(x.data_ptr<float>(), temp.data_ptr<float>(), row_norm_w.data_ptr<float>(), row_norm_b.data_ptr<float>(), row_qkv_w.data_ptr<float>(), row_out_w.data_ptr<float>(), row_out_b.data_ptr<float>(), B, H, W, W);
    axial_attn_kernel<true><<<B * W, threads, sh_col, stream>>>(temp.data_ptr<float>(), y.data_ptr<float>(), col_norm_w.data_ptr<float>(), col_norm_b.data_ptr<float>(), col_qkv_w.data_ptr<float>(), col_out_w.data_ptr<float>(), col_out_b.data_ptr<float>(), B, H, W, H);
    return y;
}
'''

_ext = load_inline(name='axial_attn_ext_final', cpp_sources=cpp_src, cuda_sources=cuda_src, functions=['axial_forward'], extra_cuda_cflags=['-O3', '--use_fast_math'], extra_cflags=['-O3'], verbose=False)

class ModelNew(nn.Module):
    def __init__(self, dim, num_heads, height, width, dropout=0.0):
        super().__init__()
        self.height = height
        self.width = width
        self.num_heads = num_heads
        self.head_dim = dim // num_heads
        assert dim % num_heads == 0
        self.row_norm = nn.LayerNorm(dim)
        self.row_qkv = nn.Linear(dim, 3 * dim, bias=False)
        self.row_out = nn.Linear(dim, dim)
        self.col_norm = nn.LayerNorm(dim)
        self.col_qkv = nn.Linear(dim, 3 * dim, bias=False)
        self.col_out = nn.Linear(dim, dim)
        self.dropout = nn.Dropout(dropout)
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return _ext.axial_forward(x.contiguous(), self.row_norm.weight, self.row_norm.bias, self.row_qkv.weight, self.row_out.weight, self.row_out.bias, self.col_norm.weight, self.col_norm.bias, self.col_qkv.weight, self.col_out.weight, self.col_out.bias)