import torch
import torch.nn as nn

class Model(nn.Module):
    def __init__(self, num_classes, embed_dim=512, num_heads=8, num_layers=6, 
                 mlp_ratio=4.0, patch_size=4, in_channels=3, image_size=32):
        """
        Convolutional Vision Transformer (CViT) implementation.
        :param num_classes: Number of output classes for classification.
        :param embed_dim: Dimensionality of the embedding space.
        :param num_heads: Number of attention heads.
        :param num_layers: Number of transformer layers.
        :param mlp_ratio: Ratio of the MLP hidden dimension to the embedding dimension.
        :param patch_size: Size of the convolutional patches.
        :param in_channels: Number of input channels (e.g., 3 for RGB images).
        :param image_size: Height/width of the square input image.
        """
        super(Model, self).__init__()

        self.patch_size = patch_size
        self.image_size = image_size
        self.embed_dim = embed_dim

        self.conv1 = nn.Conv2d(in_channels, embed_dim, kernel_size=patch_size, stride=patch_size)
        num_patches = (image_size // patch_size) ** 2  # Total number of patches after conv
        self.linear_proj = nn.Linear(embed_dim * num_patches, embed_dim)

        self.transformer_layers = nn.ModuleList([
            nn.TransformerEncoderLayer(
                d_model=embed_dim,
                nhead=num_heads,
                dim_feedforward=int(embed_dim * mlp_ratio),
                dropout=0.0,
                batch_first=True
            ) for _ in range(num_layers)
        ])

        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
        self.fc_out = nn.Linear(embed_dim, num_classes)

    def forward(self, x):
        """
        Forward pass of the CViT model.
        :param x: Input tensor of shape (B, C, H, W)
        :return: Output tensor of shape (B, num_classes)
        """
        B = x.size(0)
        x = self.conv1(x)                  # (B, embed_dim, H/patch_size, W/patch_size)
        x = x.flatten(start_dim=1)         # (B, embed_dim * num_patches)
        x = self.linear_proj(x)            # (B, embed_dim)

        cls_tokens = self.cls_token.expand(B, -1, -1)  # (B, 1, embed_dim)
        x = torch.cat((cls_tokens, x.unsqueeze(1)), dim=1)  # (B, 2, embed_dim)

        for layer in self.transformer_layers:
            x = layer(x)

        return self.fc_out(x[:, 0])        # Use [CLS] token for classification

# === Test config ===
batch_size = 10
image_size = 32
embed_dim = 128
in_channels = 3
num_heads = 4
num_classes = 1000

def get_inputs():
    return [torch.rand(batch_size, in_channels, image_size, image_size)]

def get_init_inputs():
    return [num_classes, embed_dim, num_heads]

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

_cpp_src = r'''
#include <torch/extension.h>
#include <vector>
#include <pybind11/stl.h>

torch::Tensor cvit_forward(
    torch::Tensor x,
    torch::Tensor conv_w, torch::Tensor conv_b,
    torch::Tensor lp_w, torch::Tensor lp_b,
    torch::Tensor cls,
    std::vector<torch::Tensor> in_w, std::vector<torch::Tensor> in_b,
    std::vector<torch::Tensor> out_w, std::vector<torch::Tensor> out_b,
    std::vector<torch::Tensor> n1_w, std::vector<torch::Tensor> n1_b,
    std::vector<torch::Tensor> l1_w, std::vector<torch::Tensor> l1_b,
    std::vector<torch::Tensor> l2_w, std::vector<torch::Tensor> l2_b,
    std::vector<torch::Tensor> n2_w, std::vector<torch::Tensor> n2_b,
    torch::Tensor fc_w, torch::Tensor fc_b,
    int64_t num_heads);
'''

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

__global__ void conv_flat_kernel(const float* __restrict__ x, const float* __restrict__ w, const float* __restrict__ bias,
                                 float* __restrict__ out,
                                 int B, int C, int H, int W, int E, int P) {
    int Hp = H / P;
    int Wp = W / P;
    int Kflat = E * Hp * Wp;
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int total = B * Kflat;
    if (idx >= total) return;
    int k = idx % Kflat;
    int b = idx / Kflat;
    int pw = k % Wp;
    int ph = (k / Wp) % Hp;
    int e = k / (Hp * Wp);
    float sum = bias[e];
    int ih0 = ph * P;
    int iw0 = pw * P;
    for (int c = 0; c < C; ++c) {
        for (int kh = 0; kh < P; ++kh) {
            for (int kw = 0; kw < P; ++kw) {
                float xv = x[((b*C + c)*H + (ih0+kh))*W + (iw0+kw)];
                float wv = w[((e*C + c)*P + kh)*P + kw];
                sum += xv * wv;
            }
        }
    }
    out[idx] = sum;
}

__global__ void linear_proj_kernel(const float* __restrict__ inp, const float* __restrict__ w, const float* __restrict__ bias,
                                   float* __restrict__ out, int B, int K, int D) {
    __shared__ float s[256];
    int od = blockIdx.x;
    int b = blockIdx.y;
    int tid = threadIdx.x;
    float sum = 0.0f;
    const float* ip = inp + b*K;
    const float* wp = w + od*K;
    for (int k = tid; k < K; k += blockDim.x) sum += ip[k] * wp[k];
    s[tid] = sum;
    __syncthreads();
    for (int off = blockDim.x/2; off > 0; off >>= 1) {
        if (tid < off) s[tid] += s[tid+off];
        __syncthreads();
    }
    if (tid == 0) out[b*D + od] = s[0] + bias[od];
}

__global__ void init_seq_kernel(const float* __restrict__ lin, const float* __restrict__ cls,
                                float* __restrict__ seq, int B, int D) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int total = B * 2 * D;
    if (idx >= total) return;
    int d = idx % D;
    int t = (idx / D) & 1;
    int b = idx / (2*D);
    seq[idx] = (t == 0) ? cls[d] : lin[b*D + d];
}

__global__ void encoder_layer_kernel(float* __restrict__ seq,
                                     const float* __restrict__ inpw, const float* __restrict__ inpb,
                                     const float* __restrict__ opw, const float* __restrict__ opb,
                                     const float* __restrict__ n1w, const float* __restrict__ n1b,
                                     const float* __restrict__ l1w, const float* __restrict__ l1b,
                                     const float* __restrict__ l2w, const float* __restrict__ l2b,
                                     const float* __restrict__ n2w, const float* __restrict__ n2b,
                                     int D, int FF, int NH) {
    __shared__ float X[256];
    __shared__ float Q[256];
    __shared__ float Kk[256];
    __shared__ float V[256];
    __shared__ float A[256];
    __shared__ float Z[256];
    __shared__ float Y[256];
    __shared__ float Hh[1024];
    __shared__ float red[256];

    int b = blockIdx.x;
    int tid = threadIdx.x;
    float* base = seq + b * 2 * D;

    for (int i = tid; i < 2*D; i += blockDim.x) X[i] = base[i];
    __syncthreads();

    // in-projection: q,k,v = X @ W^T + b
    for (int idx = tid; idx < 2*3*D; idx += blockDim.x) {
        int d = idx % D;
        int qkv = (idx / D) % 3;
        int t = idx / (3*D);
        float sum = inpb[qkv*D + d];
        const float* wp = inpw + (qkv*D + d) * D;
        const float* xp = X + t*D;
        for (int j = 0; j < D; ++j) sum += xp[j] * wp[j];
        if (qkv == 0) Q[t*D + d] = sum;
        else if (qkv == 1) Kk[t*D + d] = sum;
        else V[t*D + d] = sum;
    }
    __syncthreads();

    int HD = D / NH;
    float scale = rsqrtf((float)HD);
    // attention output before output projection
    for (int idx = tid; idx < 2*D; idx += blockDim.x) {
        int d = idx % D;
        int t = idx / D;
        int h = d / HD;
        int h0 = h * HD;
        float s0 = 0.0f, s1 = 0.0f;
        for (int i = 0; i < HD; ++i) {
            float qv = Q[t*D + h0 + i];
            s0 += qv * Kk[0*D + h0 + i];
            s1 += qv * Kk[1*D + h0 + i];
        }
        s0 *= scale; s1 *= scale;
        float m = s0 > s1 ? s0 : s1;
        float e0 = expf(s0 - m);
        float e1 = expf(s1 - m);
        float inv = 1.0f / (e0 + e1);
        float w0 = e0 * inv, w1 = e1 * inv;
        A[idx] = w0 * V[0*D + d] + w1 * V[1*D + d];
    }
    __syncthreads();

    // attention output projection, residual
    for (int idx = tid; idx < 2*D; idx += blockDim.x) {
        int d = idx % D;
        int t = idx / D;
        float sum = opb[d];
        const float* wp = opw + d*D;
        const float* ap = A + t*D;
        for (int j = 0; j < D; ++j) sum += ap[j] * wp[j];
        Z[idx] = X[idx] + sum;
    }
    __syncthreads();

    // norm1
    for (int t = 0; t < 2; ++t) {
        float v = (tid < D) ? Z[t*D + tid] : 0.0f;
        red[tid] = v;
        __syncthreads();
        for (int off = blockDim.x/2; off > 0; off >>= 1) {
            if (tid < off) red[tid] += red[tid+off];
            __syncthreads();
        }
        float mean = red[0] / (float)D;
        float vv = 0.0f;
        if (tid < D) { float a = Z[t*D + tid] - mean; vv = a*a; }
        red[tid] = vv;
        __syncthreads();
        for (int off = blockDim.x/2; off > 0; off >>= 1) {
            if (tid < off) red[tid] += red[tid+off];
            __syncthreads();
        }
        float invstd = rsqrtf(red[0] / (float)D + 1.0e-5f);
        if (tid < D) Y[t*D + tid] = (Z[t*D + tid] - mean) * invstd * n1w[tid] + n1b[tid];
        __syncthreads();
    }

    // MLP linear1 + ReLU
    for (int idx = tid; idx < 2*FF; idx += blockDim.x) {
        int f = idx % FF;
        int t = idx / FF;
        float sum = l1b[f];
        const float* wp = l1w + f*D;
        const float* yp = Y + t*D;
        for (int j = 0; j < D; ++j) sum += yp[j] * wp[j];
        Hh[idx] = sum > 0.0f ? sum : 0.0f;
    }
    __syncthreads();

    // MLP linear2, residual
    for (int idx = tid; idx < 2*D; idx += blockDim.x) {
        int d = idx % D;
        int t = idx / D;
        float sum = l2b[d];
        const float* hp = Hh + t*FF;
        for (int f = 0; f < FF; ++f) sum += hp[f] * l2w[d*FF + f];
        Z[idx] = Y[idx] + sum;
    }
    __syncthreads();

    // norm2 and write back
    for (int t = 0; t < 2; ++t) {
        float v = (tid < D) ? Z[t*D + tid] : 0.0f;
        red[tid] = v;
        __syncthreads();
        for (int off = blockDim.x/2; off > 0; off >>= 1) {
            if (tid < off) red[tid] += red[tid+off];
            __syncthreads();
        }
        float mean = red[0] / (float)D;
        float vv = 0.0f;
        if (tid < D) { float a = Z[t*D + tid] - mean; vv = a*a; }
        red[tid] = vv;
        __syncthreads();
        for (int off = blockDim.x/2; off > 0; off >>= 1) {
            if (tid < off) red[tid] += red[tid+off];
            __syncthreads();
        }
        float invstd = rsqrtf(red[0] / (float)D + 1.0e-5f);
        if (tid < D) base[t*D + tid] = (Z[t*D + tid] - mean) * invstd * n2w[tid] + n2b[tid];
        __syncthreads();
    }
}

__global__ void fc_kernel(const float* __restrict__ seq, const float* __restrict__ w, const float* __restrict__ bias,
                          float* __restrict__ out, int B, int D, int C) {
    int b = blockIdx.x;
    int tid = threadIdx.x;
    const float* xp = seq + b*2*D; // CLS token
    for (int c = tid; c < C; c += blockDim.x) {
        float sum = bias[c];
        const float* wp = w + c*D;
        for (int j = 0; j < D; ++j) sum += xp[j] * wp[j];
        out[b*C + c] = sum;
    }
}

torch::Tensor cvit_forward(
    torch::Tensor x,
    torch::Tensor conv_w, torch::Tensor conv_b,
    torch::Tensor lp_w, torch::Tensor lp_b,
    torch::Tensor cls,
    std::vector<torch::Tensor> in_w, std::vector<torch::Tensor> in_b,
    std::vector<torch::Tensor> out_w, std::vector<torch::Tensor> out_b,
    std::vector<torch::Tensor> n1_w, std::vector<torch::Tensor> n1_b,
    std::vector<torch::Tensor> l1_w, std::vector<torch::Tensor> l1_b,
    std::vector<torch::Tensor> l2_w, std::vector<torch::Tensor> l2_b,
    std::vector<torch::Tensor> n2_w, std::vector<torch::Tensor> n2_b,
    torch::Tensor fc_w, torch::Tensor fc_b,
    int64_t num_heads) {

    int B = (int)x.size(0);
    int C = (int)x.size(1);
    int H = (int)x.size(2);
    int W = (int)x.size(3);
    int E = (int)conv_w.size(0);
    int P = (int)conv_w.size(2);
    int Hp = H / P, Wp = W / P;
    int Kflat = E * Hp * Wp;
    int D = (int)lp_w.size(0);
    int FF = (int)l1_w[0].size(0);
    int NC = (int)fc_w.size(0);

    auto opts = x.options();
    auto convflat = torch::empty({B, Kflat}, opts);
    auto lin = torch::empty({B, D}, opts);
    auto seq = torch::empty({B, 2, D}, opts);
    auto out = torch::empty({B, NC}, opts);

    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    int threads = 256;
    int total_conv = B * Kflat;
    conv_flat_kernel<<<(total_conv + threads - 1)/threads, threads, 0, stream>>>(
        x.data_ptr<float>(), conv_w.data_ptr<float>(), conv_b.data_ptr<float>(), convflat.data_ptr<float>(), B, C, H, W, E, P);
    dim3 grid_lp(D, B);
    linear_proj_kernel<<<grid_lp, threads, 0, stream>>>(
        convflat.data_ptr<float>(), lp_w.data_ptr<float>(), lp_b.data_ptr<float>(), lin.data_ptr<float>(), B, Kflat, D);
    int total_seq = B * 2 * D;
    init_seq_kernel<<<(total_seq + threads - 1)/threads, threads, 0, stream>>>(
        lin.data_ptr<float>(), cls.data_ptr<float>(), seq.data_ptr<float>(), B, D);

    for (size_t i = 0; i < in_w.size(); ++i) {
        encoder_layer_kernel<<<B, threads, 0, stream>>>(
            seq.data_ptr<float>(),
            in_w[i].data_ptr<float>(), in_b[i].data_ptr<float>(),
            out_w[i].data_ptr<float>(), out_b[i].data_ptr<float>(),
            n1_w[i].data_ptr<float>(), n1_b[i].data_ptr<float>(),
            l1_w[i].data_ptr<float>(), l1_b[i].data_ptr<float>(),
            l2_w[i].data_ptr<float>(), l2_b[i].data_ptr<float>(),
            n2_w[i].data_ptr<float>(), n2_b[i].data_ptr<float>(),
            D, FF, (int)num_heads);
    }
    fc_kernel<<<B, threads, 0, stream>>>(seq.data_ptr<float>(), fc_w.data_ptr<float>(), fc_b.data_ptr<float>(), out.data_ptr<float>(), B, D, NC);
    return out;
}
'''

_ext = load_inline(
    name='cvit_fused_ext_v1',
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=['cvit_forward'],
    extra_cflags=['-O3'],
    extra_cuda_cflags=['-O3'],
    verbose=False,
)

class ModelNew(nn.Module):
    def __init__(self, num_classes, embed_dim=512, num_heads=8, num_layers=6,
                 mlp_ratio=4.0, patch_size=4, in_channels=3, image_size=32):
        super(ModelNew, self).__init__()
        self.patch_size = patch_size
        self.image_size = image_size
        self.embed_dim = embed_dim
        self.num_heads = num_heads
        self.conv1 = nn.Conv2d(in_channels, embed_dim, kernel_size=patch_size, stride=patch_size)
        num_patches = (image_size // patch_size) ** 2
        self.linear_proj = nn.Linear(embed_dim * num_patches, embed_dim)
        self.transformer_layers = nn.ModuleList([
            nn.TransformerEncoderLayer(
                d_model=embed_dim,
                nhead=num_heads,
                dim_feedforward=int(embed_dim * mlp_ratio),
                dropout=0.0,
                batch_first=True
            ) for _ in range(num_layers)
        ])
        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
        self.fc_out = nn.Linear(embed_dim, num_classes)

    def forward(self, x):
        layers = self.transformer_layers
        return _ext.cvit_forward(
            x.contiguous(),
            self.conv1.weight, self.conv1.bias,
            self.linear_proj.weight, self.linear_proj.bias,
            self.cls_token,
            [l.self_attn.in_proj_weight for l in layers],
            [l.self_attn.in_proj_bias for l in layers],
            [l.self_attn.out_proj.weight for l in layers],
            [l.self_attn.out_proj.bias for l in layers],
            [l.norm1.weight for l in layers],
            [l.norm1.bias for l in layers],
            [l.linear1.weight for l in layers],
            [l.linear1.bias for l in layers],
            [l.linear2.weight for l in layers],
            [l.linear2.bias for l in layers],
            [l.norm2.weight for l in layers],
            [l.norm2.bias for l in layers],
            self.fc_out.weight, self.fc_out.bias,
            int(self.num_heads)
        )