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>
torch::Tensor cvit_forward_cuda(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> layer_tensors, 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_patch_kernel(const float* __restrict__ x, const float* __restrict__ w, const float* __restrict__ bias, float* __restrict__ y, int B, int C, int H, int W, int E, int K, int OH, int OW) {
    int idx=blockIdx.x*blockDim.x+threadIdx.x; int total=B*E*OH*OW; if(idx>=total) return;
    int p=idx%(OH*OW), oc=(idx/(OH*OW))%E, b=idx/(E*OH*OW); int py=p/OW, px=p-py*OW; float sum=bias[oc];
    for(int c=0;c<C;++c) for(int dy=0;dy<K;++dy) for(int dx=0;dx<K;++dx) sum=fmaf(x[((b*C+c)*H+(py*K+dy))*W+(px*K+dx)], w[((oc*C+c)*K+dy)*K+dx], sum);
    y[idx]=sum;
}
__global__ void linear_kernel(const float* __restrict__ inp,const float* __restrict__ w,const float* __restrict__ bias,float* __restrict__ out,int rows,int K,int N,int do_relu){
    extern __shared__ float sdata[]; int row=blockIdx.x,n=blockIdx.y,tid=threadIdx.x; float sum=0.f; const float* inrow=inp+row*K; const float* wrow=w+n*K;
    for(int k=tid;k<K;k+=blockDim.x) sum=fmaf(inrow[k],wrow[k],sum); sdata[tid]=sum; __syncthreads();
    for(int s=blockDim.x/2;s>0;s>>=1){if(tid<s)sdata[tid]+=sdata[tid+s]; __syncthreads();}
    if(tid==0){float v=sdata[0]+bias[n]; if(do_relu&&v<0.f)v=0.f; out[row*N+n]=v;}
}
__global__ void linear_strided_kernel(const float* __restrict__ inp,const float* __restrict__ w,const float* __restrict__ bias,float* __restrict__ out,int rows,int K,int N,int stride){
    extern __shared__ float sdata[]; int row=blockIdx.x,n=blockIdx.y,tid=threadIdx.x; float sum=0.f; const float* inrow=inp+row*stride; const float* wrow=w+n*K;
    for(int k=tid;k<K;k+=blockDim.x) sum=fmaf(inrow[k],wrow[k],sum); sdata[tid]=sum; __syncthreads();
    for(int s=blockDim.x/2;s>0;s>>=1){if(tid<s)sdata[tid]+=sdata[tid+s]; __syncthreads();}
    if(tid==0) out[row*N+n]=sdata[0]+bias[n];
}
__global__ void make_seq_kernel(const float* __restrict__ z,const float* __restrict__ cls,float* __restrict__ seq,int B,int E){int idx=blockIdx.x*blockDim.x+threadIdx.x,total=B*2*E; if(idx>=total)return; int e=idx%E,t=(idx/E)%2,b=idx/(2*E); seq[(b*2+t)*E+e]=(t==0)?cls[e]:z[b*E+e];}
__global__ void attn2_kernel(const float* __restrict__ qkv,float* __restrict__ ctx,int B,int E,int Hh){
    extern __shared__ float sh[]; int b=blockIdx.x,h=blockIdx.y,tid=threadIdx.x,D=E/Hh; int base0=(b*2)*3*E+h*D, base1=(b*2+1)*3*E+h*D; float p00=0,p01=0,p10=0,p11=0;
    for(int d=tid;d<D;d+=blockDim.x){float q0=qkv[base0+d],q1=qkv[base1+d],k0=qkv[base0+E+d],k1=qkv[base1+E+d]; p00=fmaf(q0,k0,p00); p01=fmaf(q0,k1,p01); p10=fmaf(q1,k0,p10); p11=fmaf(q1,k1,p11);} 
    sh[tid]=p00; __syncthreads(); for(int s=blockDim.x/2;s>0;s>>=1){if(tid<s)sh[tid]+=sh[tid+s]; __syncthreads();} p00=sh[0];
    sh[tid]=p01; __syncthreads(); for(int s=blockDim.x/2;s>0;s>>=1){if(tid<s)sh[tid]+=sh[tid+s]; __syncthreads();} p01=sh[0];
    sh[tid]=p10; __syncthreads(); for(int s=blockDim.x/2;s>0;s>>=1){if(tid<s)sh[tid]+=sh[tid+s]; __syncthreads();} p10=sh[0];
    sh[tid]=p11; __syncthreads(); for(int s=blockDim.x/2;s>0;s>>=1){if(tid<s)sh[tid]+=sh[tid+s]; __syncthreads();} p11=sh[0];
    float scale=rsqrtf((float)D); float a=p00*scale,c=p01*scale,m=fmaxf(a,c),e0=expf(a-m),e1=expf(c-m),inv=1.f/(e0+e1),a00=e0*inv,a01=e1*inv; a=p10*scale;c=p11*scale;m=fmaxf(a,c);e0=expf(a-m);e1=expf(c-m);inv=1.f/(e0+e1);float a10=e0*inv,a11=e1*inv;
    for(int d=tid;d<D;d+=blockDim.x){float v0=qkv[base0+2*E+d],v1=qkv[base1+2*E+d]; ctx[(b*2)*E+h*D+d]=a00*v0+a01*v1; ctx[(b*2+1)*E+h*D+d]=a10*v0+a11*v1;}
}
__global__ void add_ln_kernel(const float* __restrict__ x,const float* __restrict__ add,const float* __restrict__ gamma,const float* __restrict__ beta,float* __restrict__ out,int rows,int E){
    extern __shared__ float sh[]; int row=blockIdx.x,tid=threadIdx.x; float sum=0.f; for(int e=tid;e<E;e+=blockDim.x) sum+=x[row*E+e]+add[row*E+e]; sh[tid]=sum; __syncthreads(); for(int s=blockDim.x/2;s>0;s>>=1){if(tid<s)sh[tid]+=sh[tid+s]; __syncthreads();} float mean=sh[0]/(float)E; float ss=0.f; for(int e=tid;e<E;e+=blockDim.x){float v=x[row*E+e]+add[row*E+e]-mean; ss+=v*v;} sh[tid]=ss; __syncthreads(); for(int s=blockDim.x/2;s>0;s>>=1){if(tid<s)sh[tid]+=sh[tid+s]; __syncthreads();} float rstd=rsqrtf(sh[0]/(float)E+1e-5f); for(int e=tid;e<E;e+=blockDim.x){float v=x[row*E+e]+add[row*E+e]; out[row*E+e]=(v-mean)*rstd*gamma[e]+beta[e];}
}

torch::Tensor cvit_forward_cuda(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> layer_tensors, torch::Tensor fc_w, torch::Tensor fc_b, int64_t num_heads){
    int B=x.size(0),C=x.size(1),IH=x.size(2),IW=x.size(3),E=conv_w.size(0),Kp=conv_w.size(2),OH=IH/Kp,OW=IW/Kp,P=OH*OW; auto opts=x.options(); cudaStream_t stream=at::cuda::getCurrentCUDAStream(); int threads=256;
    auto conv=torch::empty({B,E,P},opts); auto z=torch::empty({B,E},opts); auto seq=torch::empty({B*2,E},opts);
    conv_patch_kernel<<<(B*E*P+threads-1)/threads,threads,0,stream>>>(x.data_ptr<float>(),conv_w.data_ptr<float>(),conv_b.data_ptr<float>(),conv.data_ptr<float>(),B,C,IH,IW,E,Kp,OH,OW);
    linear_kernel<<<dim3(B,E),threads,threads*sizeof(float),stream>>>(conv.data_ptr<float>(),lp_w.data_ptr<float>(),lp_b.data_ptr<float>(),z.data_ptr<float>(),B,E*P,E,0);
    make_seq_kernel<<<(B*2*E+threads-1)/threads,threads,0,stream>>>(z.data_ptr<float>(),cls.data_ptr<float>(),seq.data_ptr<float>(),B,E);
    int rows=B*2; auto qkv=torch::empty({rows,3*E},opts),ctx=torch::empty({rows,E},opts),proj=torch::empty({rows,E},opts),ff1=torch::empty({rows,(int)layer_tensors[9].size(0)},opts),ff2=torch::empty({rows,E},opts); int L=layer_tensors.size()/12;
    for(int l=0;l<L;++l){auto inpw=layer_tensors[l*12],inpb=layer_tensors[l*12+1],outw=layer_tensors[l*12+2],outb=layer_tensors[l*12+3],n1w=layer_tensors[l*12+4],n1b=layer_tensors[l*12+5],n2w=layer_tensors[l*12+6],n2b=layer_tensors[l*12+7],l1w=layer_tensors[l*12+8],l1b=layer_tensors[l*12+9],l2w=layer_tensors[l*12+10],l2b=layer_tensors[l*12+11];
      linear_kernel<<<dim3(rows,3*E),threads,threads*sizeof(float),stream>>>(seq.data_ptr<float>(),inpw.data_ptr<float>(),inpb.data_ptr<float>(),qkv.data_ptr<float>(),rows,E,3*E,0);
      attn2_kernel<<<dim3(B,(int)num_heads),64,64*sizeof(float),stream>>>(qkv.data_ptr<float>(),ctx.data_ptr<float>(),B,E,(int)num_heads);
      linear_kernel<<<dim3(rows,E),threads,threads*sizeof(float),stream>>>(ctx.data_ptr<float>(),outw.data_ptr<float>(),outb.data_ptr<float>(),proj.data_ptr<float>(),rows,E,E,0);
      add_ln_kernel<<<rows,threads,threads*sizeof(float),stream>>>(seq.data_ptr<float>(),proj.data_ptr<float>(),n1w.data_ptr<float>(),n1b.data_ptr<float>(),seq.data_ptr<float>(),rows,E);
      int F=l1b.size(0); linear_kernel<<<dim3(rows,F),threads,threads*sizeof(float),stream>>>(seq.data_ptr<float>(),l1w.data_ptr<float>(),l1b.data_ptr<float>(),ff1.data_ptr<float>(),rows,E,F,1);
      linear_kernel<<<dim3(rows,E),threads,threads*sizeof(float),stream>>>(ff1.data_ptr<float>(),l2w.data_ptr<float>(),l2b.data_ptr<float>(),ff2.data_ptr<float>(),rows,F,E,0);
      add_ln_kernel<<<rows,threads,threads*sizeof(float),stream>>>(seq.data_ptr<float>(),ff2.data_ptr<float>(),n2w.data_ptr<float>(),n2b.data_ptr<float>(),seq.data_ptr<float>(),rows,E);
    }
    int classes=fc_b.size(0); auto out=torch::empty({B,classes},opts); linear_strided_kernel<<<dim3(B,classes),threads,threads*sizeof(float),stream>>>(seq.data_ptr<float>(),fc_w.data_ptr<float>(),fc_b.data_ptr<float>(),out.data_ptr<float>(),B,E,classes,2*E); return out;
}
'''
_ext = load_inline(name='cvit_ext_final', cpp_sources=cpp_src, cuda_sources=cuda_src, functions=['cvit_forward_cuda'], extra_cuda_cflags=['-O3', '--use_fast_math'], 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)
        self._lt=[]
        for layer in self.transformer_layers:
            self._lt.extend([layer.self_attn.in_proj_weight,layer.self_attn.in_proj_bias,layer.self_attn.out_proj.weight,layer.self_attn.out_proj.bias,layer.norm1.weight,layer.norm1.bias,layer.norm2.weight,layer.norm2.bias,layer.linear1.weight,layer.linear1.bias,layer.linear2.weight,layer.linear2.bias])
    def forward(self,x):
        return _ext.cvit_forward_cuda(x,self.conv1.weight,self.conv1.bias,self.linear_proj.weight,self.linear_proj.bias,self.cls_token,self._lt,self.fc_out.weight,self.fc_out.bias,self.num_heads)