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


class ModelNew(nn.Module):
    def __init__(self, alpha_combo=0.5, gamma_focal=2.0, smooth=1.0):
        super().__init__()
        self.alpha_combo = alpha_combo
        self.gamma_focal = gamma_focal
        self.smooth = smooth
        self._compile_cuda_kernel()

    def _compile_cuda_kernel(self):
        cpp_source = """
        torch::Tensor combo_loss_cuda(torch::Tensor logits, torch::Tensor targets, double gamma, double smooth, int batch_size);
        """

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

        __device__ __forceinline__ double sigmoid_d(double x) {
            if (x >= 0.0) {
                double z = exp(-x);
                return 1.0 / (1.0 + z);
            } else {
                double z = exp(x);
                return z / (1.0 + z);
            }
        }

        __device__ __forceinline__ double log_sigmoid_d(double x) {
            if (x >= 0.0) {
                return -log(1.0 + exp(-x));
            } else {
                return x - log(1.0 + exp(x));
            }
        }

        __global__ void combo_loss_kernel(
            const double* __restrict__ logits,
            const double* __restrict__ targets,
            double* __restrict__ dice_inter,
            double* __restrict__ dice_sum_inputs,
            double* __restrict__ dice_sum_targets,
            double* __restrict__ focal_out,
            const int batch_size,
            const int feature_dim,
            const double gamma) 
        {
            const int batch_idx = blockIdx.x;
            const int tid = threadIdx.x;
            const int stride = blockDim.x;

            if (batch_idx >= batch_size) return;

            __shared__ double s_inter[256];
            __shared__ double s_sum_inputs[256];
            __shared__ double s_sum_targets[256];
            __shared__ double s_focal[256];

            double local_inter = 0.0;
            double local_sum_inputs = 0.0;
            double local_sum_targets = 0.0;
            double local_focal = 0.0;

            const int offset = batch_idx * feature_dim;

            for (int i = tid; i < feature_dim; i += stride) {
                double z = logits[offset + i];
                double y = targets[offset + i];

                double p = sigmoid_d(z);

                local_inter += p * y;
                local_sum_inputs += p;
                local_sum_targets += y;

                double log_p = log_sigmoid_d(z);
                double log_1mp = log_sigmoid_d(-z);
                double bce = -(y * log_p + (1.0 - y) * log_1mp);

                double pt = exp(-bce);
                double focal_weight = pow(1.0 - pt, gamma);
                local_focal += focal_weight * bce;
            }

            s_inter[tid] = local_inter;
            s_sum_inputs[tid] = local_sum_inputs;
            s_sum_targets[tid] = local_sum_targets;
            s_focal[tid] = local_focal;
            __syncthreads();

            for (int s = blockDim.x / 2; s > 0; s >>= 1) {
                if (tid < s) {
                    s_inter[tid] += s_inter[tid + s];
                    s_sum_inputs[tid] += s_sum_inputs[tid + s];
                    s_sum_targets[tid] += s_sum_targets[tid + s];
                    s_focal[tid] += s_focal[tid + s];
                }
                __syncthreads();
            }

            if (tid == 0) {
                dice_inter[batch_idx] = s_inter[0];
                dice_sum_inputs[batch_idx] = s_sum_inputs[0];
                dice_sum_targets[batch_idx] = s_sum_targets[0];
                focal_out[batch_idx] = s_focal[0];
            }
        }

        torch::Tensor combo_loss_cuda(torch::Tensor logits, torch::Tensor targets, double gamma, double smooth, int batch_size) {
            auto Z_c = logits.contiguous();
            auto Y_c = targets.contiguous();

            const int feature_dim = Z_c.size(1);

            auto dice_inter = torch::zeros({batch_size}, Z_c.options());
            auto dice_sum_inputs = torch::zeros({batch_size}, Z_c.options());
            auto dice_sum_targets = torch::zeros({batch_size}, Z_c.options());
            auto focal_out = torch::zeros({batch_size}, Z_c.options());

            const int threads = 256;
            const int blocks = batch_size;

            combo_loss_kernel<<<blocks, threads>>>(
                Z_c.data_ptr<double>(),
                Y_c.data_ptr<double>(),
                dice_inter.data_ptr<double>(),
                dice_sum_inputs.data_ptr<double>(),
                dice_sum_targets.data_ptr<double>(),
                focal_out.data_ptr<double>(),
                batch_size,
                feature_dim,
                gamma
            );

            return torch::cat({dice_inter, dice_sum_inputs, dice_sum_targets, focal_out}, 0);
        }
        """

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

    def forward(self, logits, targets):
        targets_f = targets.to(logits.dtype)

        batch_size = logits.size(0)

        comp_flat = self.op.combo_loss_cuda(logits, targets_f, self.gamma_focal, self.smooth, batch_size)

        dice_inter = comp_flat[:batch_size]
        dice_sum_inputs = comp_flat[batch_size:2 * batch_size]
        dice_sum_targets = comp_flat[2 * batch_size:3 * batch_size]
        focal_out = comp_flat[3 * batch_size:]

        dice = (2.0 * dice_inter + self.smooth) / (dice_sum_inputs + dice_sum_targets + self.smooth)
        dice_loss = 1.0 - dice.mean()

        focal_loss = focal_out.sum() / (batch_size * logits.size(1))

        return self.alpha_combo * dice_loss + (1.0 - self.alpha_combo) * focal_loss