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


class ModelNew(nn.Module):

    def __init__(self, alpha=0.7, beta=0.3, gamma=0.75, smooth=1.0):
        super().__init__()
        self.alpha = alpha
        self.beta = beta
        self.gamma = gamma
        self.smooth = smooth
        self._compile_cuda_kernel()

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

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


        __device__ __forceinline__ double sigmoid_stable(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);
            }
        }


        __global__ void focal_tversky_kernel(
            const double* __restrict__ logits,
            const double* __restrict__ targets,
            double* __restrict__ tp_out,        // True Positive
            double* __restrict__ fp_out,        // False Positive
            double* __restrict__ fn_out,        // False Negative
            const int batch_size,
            const int feature_dim) 
        {
            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_tp[256];
            __shared__ double s_fp[256];
            __shared__ double s_fn[256];

            double local_tp = 0.0;
            double local_fp = 0.0;
            double local_fn = 0.0;

            const int offset = batch_idx * feature_dim;


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

                // 计算概率
                double prob = sigmoid_stable(logit);
                double one_minus_prob = 1.0 - prob;
                double one_minus_target = 1.0 - target;

                // 累加TP, FP, FN
                local_tp += prob * target;                      // TP = p * y
                local_fp += prob * one_minus_target;            // FP = p * (1-y)
                local_fn += one_minus_prob * target;            // FN = (1-p) * y
            }

            // 写入共享内存
            s_tp[tid] = local_tp;
            s_fp[tid] = local_fp;
            s_fn[tid] = local_fn;
            __syncthreads();


            for (int s = blockDim.x / 2; s > 0; s >>= 1) {
                if (tid < s) {
                    s_tp[tid] += s_tp[tid + s];
                    s_fp[tid] += s_fp[tid + s];
                    s_fn[tid] += s_fn[tid + s];
                }
                __syncthreads();
            }

            if (tid == 0) {
                tp_out[batch_idx] = s_tp[0];
                fp_out[batch_idx] = s_fp[0];
                fn_out[batch_idx] = s_fn[0];
            }
        }


        torch::Tensor focal_tversky_loss_cuda(
            torch::Tensor logits, 
            torch::Tensor targets, 
            double smooth, 
            int batch_size) 
        {
            // 确保内存连续
            auto logits_c = logits.contiguous();
            auto targets_c = targets.contiguous();

            const int feature_dim = logits_c.size(1);

            // 分配输出张量
            auto tp_out = torch::zeros({batch_size}, logits_c.options());
            auto fp_out = torch::zeros({batch_size}, logits_c.options());
            auto fn_out = torch::zeros({batch_size}, logits_c.options());

            // 配置kernel
            const int threads = 256;
            const int blocks = batch_size;

            // 启动kernel
            focal_tversky_kernel<<<blocks, threads>>>(
                logits_c.data_ptr<double>(),
                targets_c.data_ptr<double>(),
                tp_out.data_ptr<double>(),
                fp_out.data_ptr<double>(),
                fn_out.data_ptr<double>(),
                batch_size,
                feature_dim
            );

            // 拼接输出：[tp, fp, fn]
            return torch::cat({tp_out, fp_out, fn_out}, 0);
        }
        """

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

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

        batch_size = logits.size(0)
        feature_dim = logits.size(1)

        comp_result = self.op.focal_tversky_loss_cuda(
            logits,
            targets_f,
            self.smooth,
            batch_size
        )

        tp = comp_result[:batch_size]
        fp = comp_result[batch_size:2 * batch_size]
        fn = comp_result[2 * batch_size:3 * batch_size]

        numerator = tp + self.smooth
        denominator = tp + self.alpha * fn + self.beta * fp + self.smooth
        tversky_index = numerator / denominator

        focal_tversky_loss = torch.pow(1.0 - tversky_index, self.gamma)

        return focal_tversky_loss.mean()