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


class ModelNew(nn.Module):
    def __init__(self, reduction='mean'):
        super().__init__()
        self.reduction = reduction
        self._compile_cuda_kernel()

    def _compile_cuda_kernel(self):
        cpp_source = """
        torch::Tensor adversarial_loss_cuda(torch::Tensor input, torch::Tensor target);
        """

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

        __global__ void adversarial_loss_kernel(
            const float* __restrict__ input,
            const float* __restrict__ target,
            float* __restrict__ output,
            const int n_elements)
        {
            const int tid = blockIdx.x * blockDim.x + threadIdx.x;
            const int stride = blockDim.x * gridDim.x;

            for (int i = tid; i < n_elements; i += stride) {
                float x = input[i];
                float y = target[i];

                // Binary Cross Entropy with Logits Stability Formula:
                // max(x, 0) - x * y + log(1 + exp(-abs(x)))

                float max_val = fmaxf(x, 0.0f);
                float abs_val = fabsf(x);
                float log_term = log1pf(expf(-abs_val));

                output[i] = max_val - x * y + log_term;
            }
        }

        torch::Tensor adversarial_loss_cuda(torch::Tensor input, torch::Tensor target) {
            auto input_c = input.contiguous();
            auto target_c = target.contiguous();
            const int n_elements = input_c.numel();

            auto output = torch::empty_like(input_c);

            const int threads = 256;
            const int blocks = min((n_elements + threads - 1) / threads, 65535);

            adversarial_loss_kernel<<<blocks, threads>>>(
                input_c.data_ptr<float>(),
                target_c.data_ptr<float>(),
                output.data_ptr<float>(),
                n_elements
            );

            return output;
        }
        """

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

    def forward(self, input, target):
        loss_elementwise = self.op.adversarial_loss_cuda(input, target)

        if self.reduction == 'mean':
            return loss_elementwise.mean()
        elif self.reduction == 'sum':
            return loss_elementwise.sum()
        else:
            return loss_elementwise