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

class Model(nn.Module):
    def __init__(self, num_classes):
        """
        LeNet-5 architecture implementation in PyTorch.

        :param num_classes: The number of output classes.
        """
        super(Model, self).__init__()
        
        # Convolutional layers
        self.conv1 = nn.Conv2d(in_channels=1, out_channels=6, kernel_size=5, stride=1)
        self.conv2 = nn.Conv2d(in_channels=6, out_channels=16, kernel_size=5, stride=1)
        
        # Fully connected layers
        self.fc1 = nn.Linear(in_features=16*5*5, out_features=120)
        self.fc2 = nn.Linear(in_features=120, out_features=84)
        self.fc3 = nn.Linear(in_features=84, out_features=num_classes)
    
    def forward(self, x):
        """
        Forward pass of the LeNet-5 model.

        :param x: The input tensor, shape (batch_size, 1, 32, 32)
        :return: The output tensor, shape (batch_size, num_classes)
        """
        # First convolutional layer with ReLU activation and max pooling
        x = F.relu(self.conv1(x))
        x = F.max_pool2d(x, kernel_size=2, stride=2)
        
        # Second convolutional layer with ReLU activation and max pooling
        x = F.relu(self.conv2(x))
        x = F.max_pool2d(x, kernel_size=2, stride=2)
        
        # Flatten the output for the fully connected layers
        x = x.view(-1, 16*5*5)
        
        # First fully connected layer with ReLU activation
        x = F.relu(self.fc1(x))
        
        # Second fully connected layer with ReLU activation
        x = F.relu(self.fc2(x))
        
        # Final fully connected layer
        x = self.fc3(x)
        
        return x

# Test code for the LeNet-5 model (larger batch & image)
batch_size = 4096
num_classes = 20

def get_inputs():
    return [torch.rand(batch_size, 1, 32, 32)]

def get_init_inputs():
    return [num_classes]

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

_src_cpp = r"""
#include <torch/extension.h>
torch::Tensor lenet_forward(torch::Tensor x,
                            torch::Tensor c1w, torch::Tensor c1b,
                            torch::Tensor c2w, torch::Tensor c2b,
                            torch::Tensor f1w, torch::Tensor f1b,
                            torch::Tensor f2w, torch::Tensor f2b,
                            torch::Tensor f3w, torch::Tensor f3b);
"""

_src_cuda = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda.h>
#include <cuda_runtime.h>

__global__ void lenet_fused_kernel(const float* __restrict__ x,
                                   const float* __restrict__ c1w, const float* __restrict__ c1b,
                                   const float* __restrict__ c2w, const float* __restrict__ c2b,
                                   const float* __restrict__ f1w, const float* __restrict__ f1b,
                                   const float* __restrict__ f2w, const float* __restrict__ f2b,
                                   const float* __restrict__ f3w, const float* __restrict__ f3b,
                                   float* __restrict__ out,
                                   int K) {
    int n = blockIdx.x;
    int tid = threadIdx.x;
    __shared__ float p1[6*14*14];
    __shared__ float p2[16*5*5];
    __shared__ float h1[120];
    __shared__ float h2[84];
    __shared__ float sw1[6*25];
    __shared__ float sb1[6];
    __shared__ float sw2[16*6*25];
    __shared__ float sb2[16];

    for (int i = tid; i < 6*25; i += blockDim.x) sw1[i] = c1w[i];
    if (tid < 6) sb1[tid] = c1b[tid];
    for (int i = tid; i < 16*6*25; i += blockDim.x) sw2[i] = c2w[i];
    if (tid < 16) sb2[tid] = c2b[tid];
    __syncthreads();

    const float* xin = x + ((size_t)n)*32*32;

    for (int idx = tid; idx < 6*14*14; idx += blockDim.x) {
        int tmp = idx;
        int px = tmp % 14; tmp /= 14;
        int py = tmp % 14; tmp /= 14;
        int oc = tmp;
        int oy = py * 2;
        int ox = px * 2;
        float b = sb1[oc];
        float s00 = b, s01 = b, s10 = b, s11 = b;
        const float* w = sw1 + oc*25;
        #pragma unroll
        for (int ky = 0; ky < 5; ++ky) {
            int r0 = (oy + ky) * 32 + ox;
            int r1 = r0 + 32;
            #pragma unroll
            for (int kx = 0; kx < 5; ++kx) {
                float ww = w[ky*5 + kx];
                s00 += ww * xin[r0 + kx];
                s01 += ww * xin[r0 + kx + 1];
                s10 += ww * xin[r1 + kx];
                s11 += ww * xin[r1 + kx + 1];
            }
        }
        float m = 0.0f;
        m = fmaxf(m, s00); m = fmaxf(m, s01); m = fmaxf(m, s10); m = fmaxf(m, s11);
        p1[idx] = m;
    }
    __syncthreads();

    for (int idx = tid; idx < 16*5*5; idx += blockDim.x) {
        int tmp = idx;
        int px = tmp % 5; tmp /= 5;
        int py = tmp % 5; tmp /= 5;
        int oc = tmp;
        int oy = py * 2;
        int ox = px * 2;
        float b = sb2[oc];
        float s00 = b, s01 = b, s10 = b, s11 = b;
        const float* woc = sw2 + oc*6*25;
        #pragma unroll
        for (int ic = 0; ic < 6; ++ic) {
            const float* pp = p1 + ic*14*14;
            const float* wwic = woc + ic*25;
            #pragma unroll
            for (int ky = 0; ky < 5; ++ky) {
                int r0 = (oy + ky) * 14 + ox;
                int r1 = r0 + 14;
                #pragma unroll
                for (int kx = 0; kx < 5; ++kx) {
                    float ww = wwic[ky*5 + kx];
                    s00 += ww * pp[r0 + kx];
                    s01 += ww * pp[r0 + kx + 1];
                    s10 += ww * pp[r1 + kx];
                    s11 += ww * pp[r1 + kx + 1];
                }
            }
        }
        float m = 0.0f;
        m = fmaxf(m, s00); m = fmaxf(m, s01); m = fmaxf(m, s10); m = fmaxf(m, s11);
        p2[idx] = m;
    }
    __syncthreads();

    for (int o = tid; o < 120; o += blockDim.x) {
        float s = f1b[o];
        const float* w = f1w + o*400;
        #pragma unroll 4
        for (int i = 0; i < 400; ++i) s += w[i] * p2[i];
        h1[o] = fmaxf(s, 0.0f);
    }
    __syncthreads();

    for (int o = tid; o < 84; o += blockDim.x) {
        float s = f2b[o];
        const float* w = f2w + o*120;
        #pragma unroll 4
        for (int i = 0; i < 120; ++i) s += w[i] * h1[i];
        h2[o] = fmaxf(s, 0.0f);
    }
    __syncthreads();

    for (int o = tid; o < K; o += blockDim.x) {
        float s = f3b[o];
        const float* w = f3w + ((size_t)o)*84;
        #pragma unroll
        for (int i = 0; i < 84; ++i) s += w[i] * h2[i];
        out[((size_t)n)*K + o] = s;
    }
}

torch::Tensor lenet_forward(torch::Tensor x,
                            torch::Tensor c1w, torch::Tensor c1b,
                            torch::Tensor c2w, torch::Tensor c2b,
                            torch::Tensor f1w, torch::Tensor f1b,
                            torch::Tensor f2w, torch::Tensor f2b,
                            torch::Tensor f3w, torch::Tensor f3b) {
    const int N = (int)x.size(0);
    const int K = (int)f3b.size(0);
    auto out = torch::empty({N, K}, x.options());
    lenet_fused_kernel<<<N, 256, 0, at::cuda::getCurrentCUDAStream()>>>(
        x.data_ptr<float>(), c1w.data_ptr<float>(), c1b.data_ptr<float>(),
        c2w.data_ptr<float>(), c2b.data_ptr<float>(),
        f1w.data_ptr<float>(), f1b.data_ptr<float>(),
        f2w.data_ptr<float>(), f2b.data_ptr<float>(),
        f3w.data_ptr<float>(), f3b.data_ptr<float>(), out.data_ptr<float>(), K);
    return out;
}
"""

_ext = load_inline(name="lenet5_fused_ext_v3",
                   cpp_sources=_src_cpp,
                   cuda_sources=_src_cuda,
                   functions=["lenet_forward"],
                   extra_cuda_cflags=["-O3", "--use_fast_math"],
                   extra_cflags=["-O3"],
                   verbose=False)

class ModelNew(nn.Module):
    def __init__(self, num_classes):
        super(ModelNew, self).__init__()
        self.conv1 = nn.Conv2d(1, 6, 5, 1)
        self.conv2 = nn.Conv2d(6, 16, 5, 1)
        self.fc1 = nn.Linear(16*5*5, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, num_classes)
    def forward(self, x):
        return _ext.lenet_forward(x, self.conv1.weight, self.conv1.bias,
                                  self.conv2.weight, self.conv2.bias,
                                  self.fc1.weight, self.fc1.bias,
                                  self.fc2.weight, self.fc2.bias,
                                  self.fc3.weight, self.fc3.bias)