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

N, C, H, W = 32, 1, 64, 64


class DiceLossCUDAOp(torch.autograd.Function):
    epsilon = 1e-6

    def forward(ctx, input, target, reduction_id, op):
        if not input.is_cuda: input = input.cuda()
        if not target.is_cuda: target = target.cuda()

        input = input.contiguous()
        target = target.contiguous()

        N_batch = input.size(0)
        C_class = input.size(1)

        loss_nc, hp_terms = op.dice_loss_forward_cuda(
            input,
            target,
            N_batch,
            C_class,
            DiceLossCUDAOp.epsilon
        )

        ctx.save_for_backward(input, target, hp_terms)
        ctx.reduction_id = reduction_id
        ctx.N_C = N_batch * C_class
        ctx.op = op

        if reduction_id == 1:
            return loss_nc.mean()
        elif reduction_id == 2:
            return loss_nc.sum()
        else:
            return loss_nc

    def backward(ctx, grad_output):
        input, target, hp_terms = ctx.saved_tensors

        grad_out_scalar = grad_output[0]
        if ctx.reduction_id == 1:
            grad_out_scalar = grad_out_scalar / ctx.N_C

        grad_input = torch.empty_like(input)
        grad_target = torch.empty_like(target)

        ctx.op.dice_loss_backward_cuda(
            grad_out_scalar,
            input,
            target,
            hp_terms,
            grad_input,
            grad_target,
            input.numel(),
            input.size(2) * input.size(3),
            DiceLossCUDAOp.epsilon
        )

        return grad_input, grad_target, None, None


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

        self.red_map = {'none': 0, 'mean': 1, 'sum': 2}
        if reduction not in self.red_map:
            raise ValueError("Invalid reduction")
        self.reduction_id = self.red_map[reduction]

        self._compile_cuda_kernel()

    def _compile_cuda_kernel(self):
        cpp_source = """
        #include <torch/extension.h>
        #include <vector>

        std::vector<torch::Tensor> dice_loss_forward_cuda(
            torch::Tensor input,
            torch::Tensor target,
            int N, int C,
            float epsilon);

        void dice_loss_backward_cuda(
            float grad_out_scalar,
            torch::Tensor input,
            torch::Tensor target,
            torch::Tensor hp_terms,
            torch::Tensor grad_input,
            torch::Tensor grad_target,
            int64_t n_total,
            int hw_size,
            float epsilon);
        """

        cuda_source = """
        #include <torch/extension.h>
        #include <vector>
        #include <algorithm>
        #include <cuda_runtime.h>
        #include <cmath>
        #include <device_functions.h>
        #include <math_functions.h>

        #define BLOCK_SIZE 256
        #define MAX_GRID_SIZE 4096

        __inline__ __device__ double warp_reduce_sum_double(double val) {
            #pragma unroll
            for (int offset = 16; offset > 0; offset /= 2) {
                val += __shfl_down_sync(0xffffffff, val, offset);
            }
            return val;
        }

        __inline__ __device__ double block_reduce_sum_double(double val) {
            __shared__ double shared[32];
            int lane = threadIdx.x % 32;
            int wid = threadIdx.x / 32;

            val = warp_reduce_sum_double(val);
            if (lane == 0) shared[wid] = val;
            __syncthreads();

            val = (threadIdx.x < blockDim.x / 32) ? shared[lane] : 0.0;
            if (wid == 0) val = warp_reduce_sum_double(val);
            return val;
        }

        __global__ void dice_fwd_kernel(
            const float* __restrict__ input,
            const float* __restrict__ target,
            float* __restrict__ loss_nc,
            double* __restrict__ hp_terms_ptr,
            int N, int C, int HW,
            float epsilon)
        {
            int nc_idx = blockIdx.x;
            if (nc_idx >= N * C) return;

            const float* input_ptr = input + nc_idx * HW;
            const float* target_ptr = target + nc_idx * HW;

            double local_intersection = 0.0;
            double local_sum_p = 0.0;
            double local_sum_t = 0.0;

            for (int i = threadIdx.x; i < HW; i += blockDim.x) {
                float p_logit = input_ptr[i];
                float t = target_ptr[i];
                double p = (double)(1.0f / (1.0f + expf(-p_logit)));
                double td = (double)t;

                local_intersection += p * td;
                local_sum_p += p;
                local_sum_t += td;
            }

            __shared__ double s_data[3];
            if (threadIdx.x == 0) s_data[0] = 0.0;
            if (threadIdx.x == 1) s_data[1] = 0.0;
            if (threadIdx.x == 2) s_data[2] = 0.0;
            __syncthreads();

            local_intersection = block_reduce_sum_double(local_intersection);
            local_sum_p = block_reduce_sum_double(local_sum_p);
            local_sum_t = block_reduce_sum_double(local_sum_t);

            if (threadIdx.x == 0) {
                s_data[0] = local_intersection;
                s_data[1] = local_sum_p;
                s_data[2] = local_sum_t;
            }
            __syncthreads();

            if (threadIdx.x == 0) {
                double I_d = s_data[0];
                double S_p_d = s_data[1];
                double S_t_d = s_data[2];

                hp_terms_ptr[nc_idx * 3 + 0] = I_d;
                hp_terms_ptr[nc_idx * 3 + 1] = S_p_d;
                hp_terms_ptr[nc_idx * 3 + 2] = S_t_d;

                float num_e = 2.f * (float)I_d + epsilon;
                float den_e = (float)S_p_d + (float)S_t_d + epsilon;

                loss_nc[nc_idx] = 1.0f - num_e / den_e;
            }
        }

        __global__ void dice_bwd_kernel(
            const double grad_out_d,
            const float* __restrict__ input,
            const float* __restrict__ target,
            const double* __restrict__ hp_terms_ptr,
            float* __restrict__ grad_input,
            float* __restrict__ grad_target,
            int64_t n_total,
            int hw_size,
            const double eps_d
        )
        {
            int i = blockIdx.x * blockDim.x + threadIdx.x;
            int stride = blockDim.x * gridDim.x;

            for (int idx = i; idx < n_total; idx += stride) {
                int nc_idx = idx / hw_size;

                const double I = hp_terms_ptr[nc_idx * 3 + 0];
                const double S_p = hp_terms_ptr[nc_idx * 3 + 1];
                const double S_t = hp_terms_ptr[nc_idx * 3 + 2];

                const double S_e = S_p + S_t + eps_d;
                const double TI_e = 2.0 * I + eps_d;

                const float x_i = input[idx];
                const double t_i = (double)target[idx];

                const double p_i = (double)(1.0f / (1.0f + expf(-x_i)));
                const double dp_dx = p_i * (1.0 - p_i);

                const double grad_L_p = (TI_e / (S_e * S_e)) - (2.0 * t_i / S_e);
                grad_input[idx] = (float)(grad_L_p * dp_dx * grad_out_d);

                const double grad_L_t = (TI_e / (S_e * S_e)) - (2.0 * p_i / S_e);
                grad_target[idx] = (float)(grad_L_t * grad_out_d);
            }
        }

        std::vector<torch::Tensor> dice_loss_forward_cuda(
            torch::Tensor input,
            torch::Tensor target,
            int N, int C,
            float epsilon)
        {
            int HW = input.size(2) * input.size(3);
            auto options = input.options();

            torch::Tensor loss_nc = torch::empty({N, C}, options);
            auto hp_options = options.dtype(torch::kFloat64);
            torch::Tensor hp_terms = torch::empty({N, C, 3}, hp_options);

            const int block_size = BLOCK_SIZE;
            const int grid_size = N * C;

            dice_fwd_kernel<<<grid_size, block_size>>>(
                input.data_ptr<float>(),
                target.data_ptr<float>(),
                loss_nc.data_ptr<float>(),
                hp_terms.data_ptr<double>(),
                N, C, HW,
                epsilon
            );

            return {loss_nc, hp_terms};
        }

        void dice_loss_backward_cuda(
            float grad_out_scalar,
            torch::Tensor input,
            torch::Tensor target,
            torch::Tensor hp_terms,
            torch::Tensor grad_input,
            torch::Tensor grad_target,
            int64_t n_total,
            int hw_size,
            float epsilon)
        {
            const int block_size = BLOCK_SIZE;
            const int grid_size = std::min(
                (int)((n_total + block_size - 1) / block_size), 
                MAX_GRID_SIZE
            );

            dice_bwd_kernel<<<grid_size, block_size>>>(
                (double)grad_out_scalar,
                input.data_ptr<float>(),
                target.data_ptr<float>(),
                hp_terms.data_ptr<double>(),
                grad_input.data_ptr<float>(),
                grad_target.data_ptr<float>(),
                n_total,
                hw_size,
                (double)epsilon
            );
        }
        """

        self.op = load_inline(
            name='dice_loss_cuda_v3_full_double',
            cpp_sources=cpp_source,
            cuda_sources=cuda_source,
            functions=['dice_loss_forward_cuda', 'dice_loss_backward_cuda'],
            extra_cuda_cflags=['-O3'],
            verbose=False
        )

    def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
        if isinstance(input, (list, tuple)) and len(input) > 0:
            input = input[0]
            target = target[0] if len(target) > 0 else target

        if target.dtype != input.dtype:
            target = target.to(input.dtype)

        return DiceLossCUDAOp.apply(
            input,
            target,
            self.reduction_id,
            self.op
        )