import torch
from torch.utils.cpp_extension import load_inline

from CrossEntropyLoss_torch import BATCH_SIZE, FEATURE_DIM

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

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

        cuda_source = """
        #include <cuda_runtime.h>
        #include <float.h>
        #include <cmath>

        #define BLOCK_SIZE 256

        __global__ void cel_fused_kernel(
            const float* __restrict__ input,   // [B, C]
            const int64_t* __restrict__ target,// [B]
            float* __restrict__ loss_partial,  // [B]
            int num_classes)
        {
            extern __shared__ float smem[]; // 动态 shared mem: 存 logits & exp
            float* logits = smem;
            float* exp_logits = smem + num_classes;

            int sample_idx = blockIdx.x;
            const float* row = input + sample_idx * num_classes;
            int label = target[sample_idx];

            // 1. load logits to shared memory
            for (int j = threadIdx.x; j < num_classes; j += blockDim.x) {
                logits[j] = row[j];
            }
            __syncthreads();

            // 2. compute row max for numerical stability
            float local_max = -FLT_MAX;
            for (int j = threadIdx.x; j < num_classes; j += blockDim.x)
                local_max = fmaxf(local_max, logits[j]);

            // block reduce max
            __shared__ float row_max;
            if (threadIdx.x == 0) row_max = -FLT_MAX;
            __syncthreads();

            atomicMax((int*)&row_max, __float_as_int(local_max));
            __syncthreads();

            // 3. compute exp(x - max) and sum
            float local_sum = 0.0f;
            for (int j = threadIdx.x; j < num_classes; j += blockDim.x) {
                float e = expf(logits[j] - row_max);
                exp_logits[j] = e;
                local_sum += e;
            }

            // block reduce sum
            __shared__ float row_sum;
            if (threadIdx.x == 0) row_sum = 0.0f;
            __syncthreads();

            atomicAdd(&row_sum, local_sum);
            __syncthreads();

            // 4. compute -log(p_correct)
            float loss_val = 0.0f;
            if (threadIdx.x == 0) {
                float p_correct = exp_logits[label] / row_sum;
                loss_val = -logf(p_correct);
                loss_partial[sample_idx] = loss_val;
            }
        }

        torch::Tensor cel_forward_cuda(torch::Tensor input, torch::Tensor target) {
            TORCH_CHECK(input.is_cuda(), "input must be CUDA tensor");
            TORCH_CHECK(target.is_cuda(), "target must be CUDA tensor");

            input = input.contiguous();
            target = target.contiguous();

            const int batch_size = input.size(0);
            const int num_classes = input.size(1);

            auto loss_buf = torch::empty({batch_size}, input.options());

            const dim3 grid(batch_size);
            const dim3 block(BLOCK_SIZE);
            const size_t shmem_bytes = 2 * num_classes * sizeof(float);

            cel_fused_kernel<<<grid, block, shmem_bytes>>>(
                input.data_ptr<float>(),
                target.data_ptr<int64_t>(),
                loss_buf.data_ptr<float>(),
                num_classes
            );

            return loss_buf.mean();
        }
        """

        self.cel_op = load_inline(
            name="cel_fused_op_v2",
            cpp_sources=cpp_source,
            cuda_sources=cuda_source,
            functions=["cel_forward_cuda"],
            extra_cuda_cflags=["-O3", "--use_fast_math", "-lineinfo"],
            verbose=False
        )

    def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
        return self.cel_op.cel_forward_cuda(input, target)