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

cpp_source = """
#include <torch/extension.h>
#include <string>

torch::Tensor poly1_focal_loss_cuda_forward(
    const torch::Tensor& logits, 
    const torch::Tensor& targets,
    float gamma,
    float alpha,
    float epsilon,
    std::string reduction);
"""

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

#define BLOCK_SIZE 256

struct __align__(16) Float4 {
    float x, y, z, w;
};

__device__ __forceinline__ float compute_loss(
    float logit, float target, float gamma, float alpha, float epsilon) 
{
    // 1. Sigmoid
    float p = 1.0f / (1.0f + __expf(-logit));
    
    // 2. Get Pt and Alpha_t
    // Assume target is 0.0 or 1.0
    float pt, alpha_t;
    if (target > 0.5f) {
        pt = p;
        alpha_t = alpha;
    } else {
        pt = 1.0f - p;
        alpha_t = 1.0f - alpha;
    }
    
    // Clamp for stability
    pt = fmaxf(pt, 1e-8f);
    float one_minus_pt = 1.0f - pt;

    // 3. Compute Fused Loss
    // powf is usually faster than generic pow
    float term1 = __powf(one_minus_pt, gamma);
    
    float ce_part = -alpha_t * __logf(pt);
    float poly_part = epsilon * one_minus_pt;
    
    return term1 * (ce_part + poly_part);
}

__global__ void poly1_focal_loss_kernel(
    float* __restrict__ output,
    const float* __restrict__ logits,
    const float* __restrict__ targets,
    const int n_elements,
    const float gamma,
    const float alpha,
    const float epsilon) 
{
    const int idx = blockIdx.x * blockDim.x + threadIdx.x;
    const int vec_n = n_elements / 4;
    
    // Grid-Stride Loop
    int i = idx;
    const int stride = blockDim.x * gridDim.x;

    // 1. Vectorized Loop
    for (; i < vec_n; i += stride) {
        Float4 l_vec = reinterpret_cast<const Float4*>(logits)[i];
        Float4 t_vec = reinterpret_cast<const Float4*>(targets)[i];
        Float4 out_vec;

        out_vec.x = compute_loss(l_vec.x, t_vec.x, gamma, alpha, epsilon);
        out_vec.y = compute_loss(l_vec.y, t_vec.y, gamma, alpha, epsilon);
        out_vec.z = compute_loss(l_vec.z, t_vec.z, gamma, alpha, epsilon);
        out_vec.w = compute_loss(l_vec.w, t_vec.w, gamma, alpha, epsilon);

        reinterpret_cast<Float4*>(output)[i] = out_vec;
    }

    // 2. Scalar Loop (Tail)
    int start_scalar = vec_n * 4;
    int global_tid = idx; 
    
    if (idx < (n_elements - start_scalar)) {
        int real_idx = start_scalar + idx;
        output[real_idx] = compute_loss(logits[real_idx], targets[real_idx], gamma, alpha, epsilon);
    }
}

torch::Tensor poly1_focal_loss_cuda_forward(
    const torch::Tensor& logits, 
    const torch::Tensor& targets,
    float gamma,
    float alpha,
    float epsilon,
    std::string reduction) 
{
    TORCH_CHECK(logits.is_cuda() && targets.is_cuda(), "Inputs must be CUDA");
    TORCH_CHECK(logits.is_contiguous() && targets.is_contiguous(), "Inputs must be contiguous");
    TORCH_CHECK(logits.numel() == targets.numel(), "Shapes must match");

    const int n = logits.numel();
    auto output = torch::empty_like(logits);

    const int vec_n = n / 4;
    const int grid_size = (vec_n + BLOCK_SIZE - 1) / BLOCK_SIZE;
    int final_grid = (grid_size < 1) ? 1 : grid_size;
    if (final_grid > 65535) final_grid = 65535;

    poly1_focal_loss_kernel<<<final_grid, BLOCK_SIZE>>>(
        output.data_ptr<float>(),
        logits.data_ptr<float>(),
        targets.data_ptr<float>(),
        n,
        gamma,
        alpha,
        epsilon
    );

    if (reduction == "mean") {
        return output.mean();
    } else if (reduction == "sum") {
        return output.sum();
    }

    return output;
}
"""

class ModelNew(nn.Module):
    def __init__(self, gamma=2.0, alpha=0.25, epsilon=1.0, reduction='none'):
        super(ModelNew, self).__init__()
        self.gamma = gamma
        self.alpha = alpha
        self.epsilon = epsilon
        self.reduction = reduction
        
        self.op = load_inline(
            name='poly1_focal_loss_op',
            cpp_sources=cpp_source,
            cuda_sources=cuda_source,
            functions=['poly1_focal_loss_cuda_forward'],
            verbose=False
        )

    def forward(self, logits: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
        return self.op.poly1_focal_loss_cuda_forward(
            logits.contiguous(), 
            targets.contiguous(), 
            self.gamma,
            self.alpha,
            self.epsilon,
            self.reduction
        )