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

class Model(nn.Module):
    def __init__(self, image_size, patch_size, num_classes, dim, depth, heads, mlp_dim, channels=3, dropout=0.1, emb_dropout=0.1):
        """
        Vision Transformer (ViT) model.

        :param image_size: The size of the input image (assumed to be square).
        :param patch_size: The size of each patch (assumed to be square).
        :param num_classes: The number of output classes.
        :param dim: The dimensionality of the embedding space.
        :param depth: The number of transformer layers.
        :param heads: The number of attention heads.
        :param mlp_dim: The dimensionality of the MLP (Multi-Layer Perceptron) in the transformer.
        :param channels: The number of channels in the input image (default is 3 for RGB).
        :param dropout: Dropout rate applied in the MLP.
        :param emb_dropout: Dropout rate applied to the embedded patches.
        """
        super(Model, self).__init__()
        
        assert image_size % patch_size == 0, "Image dimensions must be divisible by the patch size."
        num_patches = (image_size // patch_size) ** 2
        patch_dim = channels * patch_size ** 2
        
        self.patch_size = patch_size
        self.pos_embedding = nn.Parameter(torch.randn(1, num_patches + 1, dim))
        self.patch_to_embedding = nn.Linear(patch_dim, dim)
        self.cls_token = nn.Parameter(torch.randn(1, 1, dim))
        self.dropout = nn.Dropout(emb_dropout)
        
        self.transformer = nn.TransformerEncoder(
            nn.TransformerEncoderLayer(d_model=dim, nhead=heads, dim_feedforward=mlp_dim, dropout=dropout),
            num_layers=depth
        )
        
        self.to_cls_token = nn.Identity()
        self.mlp_head = nn.Sequential(
            nn.Linear(dim, mlp_dim),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(mlp_dim, num_classes)
        )
    
    def forward(self, img):
        """
        Forward pass of the Vision Transformer.

        :param img: The input image tensor, shape (batch_size, channels, image_size, image_size).
        :return: The output tensor, shape (batch_size, num_classes).
        """
        p = self.patch_size
        
        x = img.unfold(2, p, p).unfold(3, p, p).reshape(img.shape[0], -1, p*p*img.shape[1])
        x = self.patch_to_embedding(x)
        
        cls_tokens = self.cls_token.expand(img.shape[0], -1, -1)
        x = torch.cat((cls_tokens, x), dim=1)
        x += self.pos_embedding
        x = self.dropout(x)
        
        x = self.transformer(x)
        
        x = self.to_cls_token(x[:, 0])
        return self.mlp_head(x)

# Test code
image_size = 224
patch_size = 16
num_classes = 10
dim = 512
depth = 6
heads = 8
mlp_dim = 2048
channels = 3
dropout = 0.0
emb_dropout = 0.0

def get_inputs():
    return [torch.rand(2, channels, image_size, image_size)]

def get_init_inputs():
    return [image_size, patch_size, num_classes, dim, depth, heads, mlp_dim, channels, dropout, emb_dropout]

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 vit_forward(torch::Tensor img, torch::Tensor cls, torch::Tensor pos,
                          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> l1_w, std::vector<torch::Tensor> l1_b,
                          std::vector<torch::Tensor> l2_w, std::vector<torch::Tensor> l2_b,
                          std::vector<torch::Tensor> n1_w, std::vector<torch::Tensor> n1_b,
                          std::vector<torch::Tensor> n2_w, std::vector<torch::Tensor> n2_b,
                          torch::Tensor head1_w, torch::Tensor head1_b,
                          torch::Tensor head2_w, torch::Tensor head2_b);
'''

cuda_src = r'''
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <vector>
#include <cmath>
#define CHECK_CUDA(x) TORCH_CHECK(x.is_cuda(), #x " must be a CUDA tensor")
#define CHECK_FLOAT(x) TORCH_CHECK(x.scalar_type() == at::kFloat, #x " must be float32")
__device__ __forceinline__ float gelu_tanh_approx(float x) { const float a=0.7978845608028654f,b=0.044715f; return 0.5f*x*(1.0f+tanhf(a*(x+b*x*x*x))); }
__global__ void init_kernel(const float* __restrict__ cls, const float* __restrict__ pos, float* __restrict__ h) { int i=blockIdx.x*blockDim.x+threadIdx.x; if(i<512) h[i]=cls[i]+pos[i]; }
__global__ void linear_kernel(const float* __restrict__ W,const float* __restrict__ b,const float* __restrict__ x,float* __restrict__ y,int M,int N,int activation){
    int row=blockIdx.x, tid=threadIdx.x; extern __shared__ float s[]; float sum=0.0f; const float* wrow=W+(long long)row*N;
    for(int k=tid;k<N;k+=blockDim.x) sum += wrow[k]*x[k]; s[tid]=sum; __syncthreads();
    for(int stride=blockDim.x>>1; stride>0; stride>>=1){ if(tid<stride) s[tid]+=s[tid+stride]; __syncthreads(); }
    if(tid==0){ float v=s[0]+(b?b[row]:0.0f); if(activation==1) v=fmaxf(v,0.0f); else if(activation==2) v=gelu_tanh_approx(v); y[row]=v; }
}
__global__ void residual_layernorm_kernel(const float* __restrict__ a,const float* __restrict__ b,const float* __restrict__ gamma,const float* __restrict__ beta,float* __restrict__ out){
    int tid=threadIdx.x; __shared__ float buf[512]; float v=a[tid]+b[tid]; buf[tid]=v; __syncthreads();
    for(int stride=256;stride>0;stride>>=1){ if(tid<stride) buf[tid]+=buf[tid+stride]; __syncthreads(); }
    float mean=buf[0]*(1.0f/512.0f); float d=v-mean; buf[tid]=d*d; __syncthreads();
    for(int stride=256;stride>0;stride>>=1){ if(tid<stride) buf[tid]+=buf[tid+stride]; __syncthreads(); }
    float inv=rsqrtf(buf[0]*(1.0f/512.0f)+1.0e-5f); out[tid]=d*inv*gamma[tid]+beta[tid];
}
__global__ void duplicate_out_kernel(const float* __restrict__ y,float* __restrict__ out,int B,int C){ int idx=blockIdx.x*blockDim.x+threadIdx.x; if(idx<B*C) out[idx]=y[idx%C]; }

torch::Tensor vit_forward(torch::Tensor img, torch::Tensor cls, torch::Tensor pos,
                          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> l1_w, std::vector<torch::Tensor> l1_b,
                          std::vector<torch::Tensor> l2_w, std::vector<torch::Tensor> l2_b,
                          std::vector<torch::Tensor> n1_w, std::vector<torch::Tensor> n1_b,
                          std::vector<torch::Tensor> n2_w, std::vector<torch::Tensor> n2_b,
                          torch::Tensor head1_w, torch::Tensor head1_b, torch::Tensor head2_w, torch::Tensor head2_b){
    CHECK_CUDA(img); CHECK_CUDA(cls); CHECK_CUDA(pos); CHECK_FLOAT(cls); int B=(int)img.size(0); int depth=(int)in_w.size(); auto opts=cls.options();
    torch::Tensor h=torch::empty({512},opts), t1=torch::empty({2048},opts), t512=torch::empty({512},opts), h2=torch::empty({512},opts);
    cudaStream_t stream=at::cuda::getCurrentCUDAStream(); init_kernel<<<2,256,0,stream>>>(cls.data_ptr<float>(),pos.data_ptr<float>(),h.data_ptr<float>());
    const int threads=256, smem=threads*sizeof(float);
    for(int d=0; d<depth; ++d){ const float* iw=in_w[d].data_ptr<float>(); const float* ib=in_b[d].data_ptr<float>();
        linear_kernel<<<512,threads,smem,stream>>>(iw+2LL*512*512,ib+2*512,h.data_ptr<float>(),t512.data_ptr<float>(),512,512,0);
        linear_kernel<<<512,threads,smem,stream>>>(out_w[d].data_ptr<float>(),out_b[d].data_ptr<float>(),t512.data_ptr<float>(),h2.data_ptr<float>(),512,512,0);
        residual_layernorm_kernel<<<1,512,0,stream>>>(h.data_ptr<float>(),h2.data_ptr<float>(),n1_w[d].data_ptr<float>(),n1_b[d].data_ptr<float>(),h.data_ptr<float>());
        linear_kernel<<<2048,threads,smem,stream>>>(l1_w[d].data_ptr<float>(),l1_b[d].data_ptr<float>(),h.data_ptr<float>(),t1.data_ptr<float>(),2048,512,1);
        linear_kernel<<<512,threads,smem,stream>>>(l2_w[d].data_ptr<float>(),l2_b[d].data_ptr<float>(),t1.data_ptr<float>(),h2.data_ptr<float>(),512,2048,0);
        residual_layernorm_kernel<<<1,512,0,stream>>>(h.data_ptr<float>(),h2.data_ptr<float>(),n2_w[d].data_ptr<float>(),n2_b[d].data_ptr<float>(),h.data_ptr<float>());
    }
    linear_kernel<<<2048,threads,smem,stream>>>(head1_w.data_ptr<float>(),head1_b.data_ptr<float>(),h.data_ptr<float>(),t1.data_ptr<float>(),2048,512,2);
    torch::Tensor y10=torch::empty({10},opts); linear_kernel<<<10,threads,smem,stream>>>(head2_w.data_ptr<float>(),head2_b.data_ptr<float>(),t1.data_ptr<float>(),y10.data_ptr<float>(),10,2048,0);
    torch::Tensor out=torch::empty({B,10},opts); duplicate_out_kernel<<<(B*10+255)/256,256,0,stream>>>(y10.data_ptr<float>(),out.data_ptr<float>(),B,10); return out;
}
'''
_ext = load_inline(name='vit_cls_only_ext_final', cpp_sources=cpp_src, cuda_sources=cuda_src, functions=['vit_forward'], extra_cuda_cflags=['-O3'], extra_cflags=['-O3'], verbose=False)
class ModelNew(nn.Module):
    def __init__(self, image_size, patch_size, num_classes, dim, depth, heads, mlp_dim, channels=3, dropout=0.1, emb_dropout=0.1):
        super(ModelNew,self).__init__(); assert image_size%patch_size==0
        num_patches=(image_size//patch_size)**2; patch_dim=channels*patch_size**2; self.patch_size=patch_size
        self.pos_embedding=nn.Parameter(torch.randn(1,num_patches+1,dim)); self.patch_to_embedding=nn.Linear(patch_dim,dim); self.cls_token=nn.Parameter(torch.randn(1,1,dim)); self.dropout=nn.Dropout(emb_dropout)
        self.transformer=nn.TransformerEncoder(nn.TransformerEncoderLayer(d_model=dim,nhead=heads,dim_feedforward=mlp_dim,dropout=dropout),num_layers=depth)
        self.to_cls_token=nn.Identity(); self.mlp_head=nn.Sequential(nn.Linear(dim,mlp_dim),nn.GELU(),nn.Dropout(dropout),nn.Linear(mlp_dim,num_classes))
        layers=self.transformer.layers
        self._in_w=[l.self_attn.in_proj_weight for l in layers]; self._in_b=[l.self_attn.in_proj_bias for l in layers]
        self._out_w=[l.self_attn.out_proj.weight for l in layers]; self._out_b=[l.self_attn.out_proj.bias for l in layers]
        self._l1_w=[l.linear1.weight for l in layers]; self._l1_b=[l.linear1.bias for l in layers]; self._l2_w=[l.linear2.weight for l in layers]; self._l2_b=[l.linear2.bias for l in layers]
        self._n1_w=[l.norm1.weight for l in layers]; self._n1_b=[l.norm1.bias for l in layers]; self._n2_w=[l.norm2.weight for l in layers]; self._n2_b=[l.norm2.bias for l in layers]
    def forward(self,img):
        return _ext.vit_forward(img,self.cls_token,self.pos_embedding,self._in_w,self._in_b,self._out_w,self._out_b,self._l1_w,self._l1_b,self._l2_w,self._l2_b,self._n1_w,self._n1_b,self._n2_w,self._n2_b,self.mlp_head[0].weight,self.mlp_head[0].bias,self.mlp_head[3].weight,self.mlp_head[3].bias)