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

_cpp_src = r"""
#include <torch/extension.h>

at::Tensor lenet_forward_cuda(at::Tensor x,
                              at::Tensor c1w, at::Tensor c1b,
                              at::Tensor c2w, at::Tensor c2b,
                              at::Tensor f1w, at::Tensor f1b,
                              at::Tensor f2w, at::Tensor f2b,
                              at::Tensor f3w, at::Tensor f3b);
"""

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

__global__ void conv1_pool_kernel(const float* __restrict__ x,
                                  const float* __restrict__ w,
                                  const float* __restrict__ b,
                                  float* __restrict__ y,
                                  int N) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int total = N * 6 * 14 * 14;
    if (idx >= total) return;

    int t = idx;
    int ox = t % 14; t /= 14;
    int oy = t % 14; t /= 14;
    int oc = t % 6;  t /= 6;
    int n = t;

    const float* xp = x + n * 32 * 32;
    const float* wp = w + oc * 25;
    float best = -FLT_MAX;

    #pragma unroll
    for (int py = 0; py < 2; ++py) {
        #pragma unroll
        for (int px = 0; px < 2; ++px) {
            int iy = oy * 2 + py;
            int ix = ox * 2 + px;
            float s = b[oc];
            #pragma unroll
            for (int ky = 0; ky < 5; ++ky) {
                #pragma unroll
                for (int kx = 0; kx < 5; ++kx) {
                    s = fmaf(xp[(iy + ky) * 32 + (ix + kx)], wp[ky * 5 + kx], s);
                }
            }
            best = fmaxf(best, s);
        }
    }
    y[idx] = fmaxf(best, 0.0f);
}

__global__ void conv2_pool_kernel(const float* __restrict__ x,
                                  const float* __restrict__ w,
                                  const float* __restrict__ b,
                                  float* __restrict__ y,
                                  int N) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int total = N * 16 * 5 * 5;
    if (idx >= total) return;

    int t = idx;
    int ox = t % 5; t /= 5;
    int oy = t % 5; t /= 5;
    int oc = t % 16; t /= 16;
    int n = t;

    const float* xp = x + n * (6 * 14 * 14);
    const float* wp0 = w + oc * (6 * 25);
    float best = -FLT_MAX;

    #pragma unroll
    for (int py = 0; py < 2; ++py) {
        #pragma unroll
        for (int px = 0; px < 2; ++px) {
            int iy = oy * 2 + py;
            int ix = ox * 2 + px;
            float s = b[oc];
            #pragma unroll
            for (int ic = 0; ic < 6; ++ic) {
                const float* xpc = xp + ic * 14 * 14;
                const float* wpc = wp0 + ic * 25;
                #pragma unroll
                for (int ky = 0; ky < 5; ++ky) {
                    #pragma unroll
                    for (int kx = 0; kx < 5; ++kx) {
                        s = fmaf(xpc[(iy + ky) * 14 + (ix + kx)], wpc[ky * 5 + kx], s);
                    }
                }
            }
            best = fmaxf(best, s);
        }
    }
    y[idx] = fmaxf(best, 0.0f);
}

__global__ void fc1_kernel(const float* __restrict__ x,
                           const float* __restrict__ w,
                           const float* __restrict__ b,
                           float* __restrict__ y,
                           int N) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int total = N * 120;
    if (idx >= total) return;
    int o = idx % 120;
    int n = idx / 120;
    const float* xp = x + n * 400;
    const float* wp = w + o * 400;
    float s = b[o];
    #pragma unroll 4
    for (int k = 0; k < 400; ++k) {
        s = fmaf(xp[k], wp[k], s);
    }
    y[idx] = fmaxf(s, 0.0f);
}

__global__ void fc2_kernel(const float* __restrict__ x,
                           const float* __restrict__ w,
                           const float* __restrict__ b,
                           float* __restrict__ y,
                           int N) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int total = N * 84;
    if (idx >= total) return;
    int o = idx % 84;
    int n = idx / 84;
    const float* xp = x + n * 120;
    const float* wp = w + o * 120;
    float s = b[o];
    #pragma unroll 4
    for (int k = 0; k < 120; ++k) {
        s = fmaf(xp[k], wp[k], s);
    }
    y[idx] = fmaxf(s, 0.0f);
}

__global__ void fc3_kernel(const float* __restrict__ x,
                           const float* __restrict__ w,
                           const float* __restrict__ b,
                           float* __restrict__ y,
                           int N, int C) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int total = N * C;
    if (idx >= total) return;
    int o = idx % C;
    int n = idx / C;
    const float* xp = x + n * 84;
    const float* wp = w + o * 84;
    float s = b[o];
    #pragma unroll 4
    for (int k = 0; k < 84; ++k) {
        s = fmaf(xp[k], wp[k], s);
    }
    y[idx] = s;
}

at::Tensor lenet_forward_cuda(at::Tensor x,
                              at::Tensor c1w, at::Tensor c1b,
                              at::Tensor c2w, at::Tensor c2b,
                              at::Tensor f1w, at::Tensor f1b,
                              at::Tensor f2w, at::Tensor f2b,
                              at::Tensor f3w, at::Tensor f3b) {
    const int N = (int)x.size(0);
    const int C = (int)f3b.size(0);
    auto opts = x.options();
    at::Tensor y1 = at::empty({N, 6, 14, 14}, opts);
    at::Tensor y2 = at::empty({N, 16, 5, 5}, opts);
    at::Tensor y3 = at::empty({N, 120}, opts);
    at::Tensor y4 = at::empty({N, 84}, opts);
    at::Tensor out = at::empty({N, C}, opts);

    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    const int block = 256;
    int total1 = N * 6 * 14 * 14;
    conv1_pool_kernel<<<(total1 + block - 1) / block, block, 0, stream>>>(
        x.data_ptr<float>(), c1w.data_ptr<float>(), c1b.data_ptr<float>(), y1.data_ptr<float>(), N);
    int total2 = N * 16 * 5 * 5;
    conv2_pool_kernel<<<(total2 + block - 1) / block, block, 0, stream>>>(
        y1.data_ptr<float>(), c2w.data_ptr<float>(), c2b.data_ptr<float>(), y2.data_ptr<float>(), N);
    int total3 = N * 120;
    fc1_kernel<<<(total3 + block - 1) / block, block, 0, stream>>>(
        y2.data_ptr<float>(), f1w.data_ptr<float>(), f1b.data_ptr<float>(), y3.data_ptr<float>(), N);
    int total4 = N * 84;
    fc2_kernel<<<(total4 + block - 1) / block, block, 0, stream>>>(
        y3.data_ptr<float>(), f2w.data_ptr<float>(), f2b.data_ptr<float>(), y4.data_ptr<float>(), N);
    int total5 = N * C;
    fc3_kernel<<<(total5 + block - 1) / block, block, 0, stream>>>(
        y4.data_ptr<float>(), f3w.data_ptr<float>(), f3b.data_ptr<float>(), out.data_ptr<float>(), N, C);
    return out;
}
"""

_ext = load_inline(
    name="lenet5_fused_forward_ext",
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=["lenet_forward_cuda"],
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3", "--use_fast_math"],
    verbose=False,
)

class ModelNew(nn.Module):
    def __init__(self, num_classes):
        super(ModelNew, self).__init__()
        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)
        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):
        return _ext.lenet_forward_cuda(
            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,
        )