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


class ModelNew(nn.Module):
    def __init__(self):
        super().__init__()
        self._compile_cuda_kernel()

    def _compile_cuda_kernel(self):
        cpp_source = """
        #include <torch/extension.h>
        torch::Tensor softmax2d_cuda_forward(torch::Tensor input);
        """

        cuda_source = """
        #include <cuda_runtime.h>
        #include <math.h>

        __global__ void softmax2d_kernel(
            const float* __restrict__ input,
            float* __restrict__ output,
            int channels,
            int spatial_stride, // H * W
            int total_spatial   // N * H * W
        ) {
            // Flattened spatial index: [0, N*H*W)
            // Maps to (n, h, w) combined
            int idx = blockIdx.x * blockDim.x + threadIdx.x;

            if (idx >= total_spatial) return;

            // Calculate base pointer offsets
            // Input layout: N, C, H, W
            // We want to iterate over C for a fixed (n, h, w)
            //
            // Decompose idx:
            // n = idx / spatial_stride
            // hw = idx % spatial_stride
            //
            // Address of input[n, c, h, w]:
            // offset = n * (C * HW) + c * HW + hw
            //        = c * HW + (n * C * HW + hw)
            // Let base_offset = n * channels * spatial_stride + (idx % spatial_stride)

            int n = idx / spatial_stride;
            int hw_offset = idx % spatial_stride;

            // Using long long to prevent overflow for large tensors
            long long base_offset = (long long)n * channels * spatial_stride + hw_offset;

            // --- Pass 1: Online Max and Sum Calculation ---
            // Reduces global memory reads from 2N (FindMax + CalcSum) to 1N

            float max_val = -1e37f; // Negative infinity approximation
            float sum_exp = 0.0f;

            for (int c = 0; c < channels; ++c) {
                // Coalesced read: threads i and i+1 read adjacent memory addresses
                // even though the loop stride is large (H*W).
                long long cur_idx = base_offset + (long long)c * spatial_stride;
                float val = input[cur_idx];

                if (val > max_val) {
                    // Update sum with scaling to prevent overflow
                    // sum = sum * exp(old_max - new_max) + exp(new_val - new_max)
                    // Note: exp(new_val - new_max) is exp(0) = 1, but we do standard formula
                    sum_exp = sum_exp * expf(max_val - val) + 1.0f;
                    max_val = val;
                } else {
                    sum_exp += expf(val - max_val);
                }
            }

            // --- Pass 2: Compute Output ---
            // Reciprocal for fast multiplication
            float inv_sum = 1.0f / sum_exp;

            for (int c = 0; c < channels; ++c) {
                long long cur_idx = base_offset + (long long)c * spatial_stride;
                // Re-read input (L2 cache likely hits if C isn't massive)
                float val = input[cur_idx];
                output[cur_idx] = expf(val - max_val) * inv_sum;
            }
        }

        torch::Tensor softmax2d_cuda_forward(torch::Tensor input) {
            auto input_contig = input.contiguous();

            int n = input_contig.size(0);
            int c = input_contig.size(1);
            int h = input_contig.size(2);
            int w = input_contig.size(3);

            auto output = torch::empty_like(input_contig);

            int spatial_stride = h * w;
            int total_spatial = n * h * w;

            int threads = 256;
            int blocks = (total_spatial + threads - 1) / threads;

            softmax2d_kernel<<<blocks, threads>>>(
                input_contig.data_ptr<float>(),
                output.data_ptr<float>(),
                c,
                spatial_stride,
                total_spatial
            );

            return output;
        }
        """

        self.op = load_inline(
            name="softmax2d_opt",
            cpp_sources=cpp_source,
            cuda_sources=cuda_source,
            functions=["softmax2d_cuda_forward"],
            extra_cuda_cflags=["-O3", "--use_fast_math"],
            verbose=False
        )

    def forward(self, x):
        return self.op.softmax2d_cuda_forward(x)