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


class ModelNew(nn.Module):
    def __init__(self, epsilon=1.0):
        super().__init__()
        self.epsilon = epsilon
        self._compile_cuda_kernel()

    def _compile_cuda_kernel(self):
        cpp_source = """
        torch::Tensor poly_loss_cuda(torch::Tensor logits, torch::Tensor labels, float epsilon);
        """

        cuda_source = """
        #include <torch/extension.h>
        #include <cuda_runtime.h>
        #include <math.h>

        __global__ void poly_loss_kernel(
            const float* __restrict__ Z,
            const long* __restrict__ Y,
            float* __restrict__ L_out,
            const int batch_size,
            const int num_classes,
            const float epsilon) 
        {
            const int tid = blockIdx.x * blockDim.x + threadIdx.x;
            const int stride = blockDim.x * gridDim.x;

            if (tid >= batch_size) return;

            const float* logit_row = Z + tid * num_classes;
            const int target_class = (int)Y[tid];

            float row_max = -3.402823466e+38F;
            float sum_exp = 0.0f;


            for (int j = 0; j < num_classes; ++j) {
                row_max = fmaxf(row_max, logit_row[j]);
            }


            for (int j = 0; j < num_classes; ++j) {
                sum_exp += expf(logit_row[j] - row_max);
            }

            float log_sum_exp = row_max + logf(sum_exp);


            float logit_t = logit_row[target_class];
            float p_t = expf(logit_t - log_sum_exp);


            float l_ce = log_sum_exp - logit_t;


            L_out[tid] = l_ce + epsilon * (1.0f - p_t);
        }

        torch::Tensor poly_loss_cuda(torch::Tensor logits, torch::Tensor labels, float epsilon) {
            auto Z_c = logits.contiguous();
            auto Y_c = labels.contiguous();

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

            auto output = torch::empty({batch_size}, Z_c.options());

            const int threads = 256;
            const int max_blocks = 65535;
            const int blocks = std::min(batch_size, max_blocks);

            poly_loss_kernel<<<blocks, threads>>>(
                Z_c.data_ptr<float>(),
                Y_c.data_ptr<long>(),
                output.data_ptr<float>(),
                batch_size,
                num_classes,
                epsilon
            );

            return output;
        }
        """

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

    def forward(self, logits, labels):
        loss_elementwise = self.op.poly_loss_cuda(logits, labels, self.epsilon)
        return loss_elementwise.mean()