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

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


class FocalLossCUDAOp(torch.autograd.Function):
    epsilon = 1e-8

    def forward(ctx, input, target, alpha, gamma, 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 = input.numel()

        output, partial_sums = op.focal_loss_forward_cuda(
            input,
            target,
            reduction_id,
            n,
            alpha,
            gamma,
            FocalLossCUDAOp.epsilon
        )

        ctx.save_for_backward(input, target)
        ctx.reduction_id = reduction_id
        ctx.alpha = alpha
        ctx.gamma = gamma
        ctx.N = n
        ctx.op = op

        return output

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

        grad_out_scalar = 0.0
        grad_output_n = None

        if ctx.reduction_id != 0:
            grad_out_scalar = grad_output[0]
            if ctx.reduction_id == 1:
                grad_out_scalar = grad_out_scalar / ctx.N
        else:
            grad_output_n = grad_output.contiguous()

        grad_input = torch.empty_like(input)

        ctx.op.focal_loss_backward_cuda(
            grad_out_scalar,
            grad_output_n,
            input,
            target,
            grad_input,
            ctx.reduction_id,
            ctx.N,
            ctx.alpha,
            ctx.gamma,
            FocalLossCUDAOp.epsilon
        )

        return grad_input, None, None, None, None, None


class ModelNew(nn.Module):
    def __init__(self, reduction='mean', alpha=0.25, gamma=2.0):
        super().__init__()
        self.alpha = float(alpha)
        self.gamma = float(gamma)

        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> focal_loss_forward_cuda(
            torch::Tensor input,
            torch::Tensor target,
            int reduction,
            int64_t n,
            float alpha,
            float gamma,
            float epsilon);

        void focal_loss_backward_cuda(
            float grad_out_scalar,
            torch::Tensor grad_output_n,
            torch::Tensor input,
            torch::Tensor target,
            torch::Tensor grad_input,
            int reduction,
            int64_t n,
            float alpha,
            float gamma,
            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 reduce_kernel(
            const float* __restrict__ input,
            float* __restrict__ output,
            int n) 
        {
            double local_sum = 0.0;
            int idx = threadIdx.x;

            for (int i = idx; i < n; i += blockDim.x) {
                local_sum += (double)input[i];
            }

            local_sum = block_reduce_sum_double(local_sum);

            if (threadIdx.x == 0) {
                output[0] = (float)local_sum;
            }
        }

        __global__ void focal_fwd_kernel(
            const float* __restrict__ input,
            const float* __restrict__ target,
            float* __restrict__ output,
            int64_t n,
            int reduction,
            float* __restrict__ partial_sums,
            const double alpha_d,
            const double gamma_d,
            const double eps_d)
        {
            const int64_t n_vec = n / 4;
            const int64_t rem_start = n_vec * 4;

            const int i = blockIdx.x * blockDim.x + threadIdx.x;
            const int stride = blockDim.x * gridDim.x;

            const float4* in_ptr = (const float4*)input;
            const float4* tgt_ptr = (const float4*)target;
            float4* out_ptr = (float4*)output;

            double local_sum = 0.0;

            for (int idx = i; idx < n_vec; idx += stride) {
                const float4 x_vec = in_ptr[idx];
                const float4 t_vec = tgt_ptr[idx];
                float4 loss_vec;

                double loss[4];
                double x_d[4], t_d[4];

                x_d[0] = (double)x_vec.x; t_d[0] = (double)t_vec.x;
                x_d[1] = (double)x_vec.y; t_d[1] = (double)t_vec.y;
                x_d[2] = (double)x_vec.z; t_d[2] = (double)t_vec.z;
                x_d[3] = (double)x_vec.w; t_d[3] = (double)t_vec.w;

                #pragma unroll
                for(int k=0; k<4; ++k) {
                    const double p_d = 1.0 / (1.0 + exp(-x_d[k]));
                    const double p_t = p_d * t_d[k] + (1.0 - p_d) * (1.0 - t_d[k]);
                    const double mod_factor = pow(1.0 - p_t, gamma_d);
                    const double alpha_t = alpha_d * t_d[k] + (1.0 - alpha_d) * (1.0 - t_d[k]);

                    const double max_val = (x_d[k] > 0.0) ? x_d[k] : 0.0;
                    const double stable_bce = max_val - x_d[k] * t_d[k] + log(1.0 + exp(-fabs(x_d[k])));

                    loss[k] = alpha_t * mod_factor * stable_bce;
                }

                if (reduction == 0) {
                    out_ptr[idx] = make_float4(
                        (float)loss[0], (float)loss[1], 
                        (float)loss[2], (float)loss[3]
                    );
                } else {
                    local_sum += loss[0] + loss[1] + loss[2] + loss[3];
                }
            }

            for (int idx = rem_start + i; idx < n; idx += stride) {
                const double x_d = (double)input[idx];
                const double t_d = (double)target[idx];

                const double p_d = 1.0 / (1.0 + exp(-x_d));
                const double p_t = p_d * t_d + (1.0 - p_d) * (1.0 - t_d);
                const double mod_factor = pow(1.0 - p_t, gamma_d);
                const double alpha_t = alpha_d * t_d + (1.0 - alpha_d) * (1.0 - t_d);

                const double max_val = (x_d > 0.0) ? x_d : 0.0;
                const double stable_bce = max_val - x_d * t_d + log(1.0 + exp(-fabs(x_d)));

                const double loss = alpha_t * mod_factor * stable_bce;

                if (reduction == 0) {
                    output[idx] = (float)loss;
                } else {
                    local_sum += loss;
                }
            }

            if (reduction != 0) {
                local_sum = block_reduce_sum_double(local_sum);
                if (threadIdx.x == 0) {
                    partial_sums[blockIdx.x] = (float)local_sum;
                }
            }
        }

        __global__ void focal_bwd_kernel(
            const double grad_out_scalar_d,
            const float* __restrict__ grad_output_n,
            const float* __restrict__ input,
            const float* __restrict__ target,
            float* __restrict__ grad_input,
            const int reduction,
            const int64_t n,
            const double alpha_d,
            const double gamma_d,
            const double eps_d)
        {
            const int64_t n_vec = n / 4;
            const int64_t rem_start = n_vec * 4;

            const int i = blockIdx.x * blockDim.x + threadIdx.x;
            const int stride = blockDim.x * gridDim.x;

            const float4* in_ptr = (const float4*)input;
            const float4* tgt_ptr = (const float4*)target;
            const float4* grad_n_ptr = (const float4*)grad_output_n;
            float4* grad_in_ptr = (float4*)grad_input;

            for (int idx = i; idx < n_vec; idx += stride) {
                const float4 x_vec = in_ptr[idx];
                const float4 t_vec = tgt_ptr[idx];

                double grad_u[4];
                double x_d[4], t_d[4];

                x_d[0] = (double)x_vec.x; t_d[0] = (double)t_vec.x;
                x_d[1] = (double)x_vec.y; t_d[1] = (double)t_vec.y;
                x_d[2] = (double)x_vec.z; t_d[2] = (double)t_vec.z;
                x_d[3] = (double)x_vec.w; t_d[3] = (double)t_vec.w;

                double grad_out_d[4];
                if (reduction == 0) {
                    const float4 grad_n_vec = grad_n_ptr[idx];
                    grad_out_d[0] = (double)grad_n_vec.x;
                    grad_out_d[1] = (double)grad_n_vec.y;
                    grad_out_d[2] = (double)grad_n_vec.z;
                    grad_out_d[3] = (double)grad_n_vec.w;
                } else {
                    grad_out_d[0] = grad_out_scalar_d;
                    grad_out_d[1] = grad_out_scalar_d;
                    grad_out_d[2] = grad_out_scalar_d;
                    grad_out_d[3] = grad_out_scalar_d;
                }

                #pragma unroll
                for(int k=0; k<4; ++k) {
                    const double p_d = 1.0 / (1.0 + exp(-x_d[k]));
                    const double dp_dx_d = p_d * (1.0 - p_d);

                    const double p_t_d = p_d * t_d[k] + (1.0 - p_d) * (1.0 - t_d[k]);
                    const double one_minus_p_t = 1.0 - p_t_d;

                    const double alpha_t_d = alpha_d * t_d[k] + (1.0 - alpha_d) * (1.0 - t_d[k]);

                    const double mod_factor_d = pow(one_minus_p_t, gamma_d);
                    const double mod_factor_g_minus_1_d = pow(one_minus_p_t, gamma_d - 1.0);

                    const double max_val = (x_d[k] > 0.0) ? x_d[k] : 0.0;
                    const double stable_bce_d = max_val - x_d[k] * t_d[k] + log(1.0 + exp(-fabs(x_d[k])));

                    const double term1 = alpha_t_d * mod_factor_d * (p_d - t_d[k]);

                    const double d_p_t = (2.0 * t_d[k] - 1.0) * dp_dx_d;
                    const double d_mod_factor = alpha_t_d * gamma_d * mod_factor_g_minus_1_d * (-d_p_t);

                    const double term2 = stable_bce_d * d_mod_factor;

                    grad_u[k] = (term1 + term2) * grad_out_d[k];
                }

                grad_in_ptr[idx] = make_float4(
                    (float)grad_u[0], (float)grad_u[1],
                    (float)grad_u[2], (float)grad_u[3]
                );
            }

            for (int idx = rem_start + i; idx < n; idx += stride) {
                const double x_d = (double)input[idx];
                const double t_d = (double)target[idx];

                const double grad_out_d = (reduction == 0) ? 
                    (double)grad_output_n[idx] : 
                    grad_out_scalar_d;

                const double p_d = 1.0 / (1.0 + exp(-x_d));
                const double dp_dx_d = p_d * (1.0 - p_d);
                const double p_t_d = p_d * t_d + (1.0 - p_d) * (1.0 - t_d);
                const double one_minus_p_t = 1.0 - p_t_d;
                const double alpha_t_d = alpha_d * t_d + (1.0 - alpha_d) * (1.0 - t_d);
                const double mod_factor_d = pow(one_minus_p_t, gamma_d);
                const double mod_factor_g_minus_1_d = pow(one_minus_p_t, gamma_d - 1.0);
                const double max_val = (x_d > 0.0) ? x_d : 0.0;
                const double stable_bce_d = max_val - x_d * t_d + log(1.0 + exp(-fabs(x_d)));
                const double term1 = alpha_t_d * mod_factor_d * (p_d - t_d);
                const double d_p_t = (2.0 * t_d - 1.0) * dp_dx_d;
                const double d_mod_factor = alpha_t_d * gamma_d * mod_factor_g_minus_1_d * (-d_p_t);
                const double term2 = stable_bce_d * d_mod_factor;

                grad_input[idx] = (float)((term1 + term2) * grad_out_d);
            }
        }

        std::vector<torch::Tensor> focal_loss_forward_cuda(
            torch::Tensor input,
            torch::Tensor target,
            int reduction,
            int64_t n,
            float alpha,
            float gamma,
            float epsilon)
        {
            auto options = input.options();
            torch::Tensor output;
            torch::Tensor partial_sums;

            const int block_size = BLOCK_SIZE;

            // --- FIX IS HERE ---
            // Replaced std::min with ternary operator
            const int A_fwd = (int)((n / 4 + block_size - 1) / block_size);
            const int B_fwd = MAX_GRID_SIZE;
            const int grid_size = (A_fwd < B_fwd ? A_fwd : B_fwd);
            // --- END FIX ---

            if (reduction == 0) {
                output = torch::empty_like(input);
                partial_sums = torch::empty({1}, options);
            } else {
                output = torch::zeros({1}, options);
                partial_sums = torch::zeros({grid_size}, options);
            }

            focal_fwd_kernel<<<grid_size, block_size>>>(
                input.data_ptr<float>(),
                target.data_ptr<float>(),
                output.data_ptr<float>(),
                n,
                reduction,
                partial_sums.data_ptr<float>(),
                (double)alpha,
                (double)gamma,
                (double)epsilon
            );

            if (reduction != 0) {
                reduce_kernel<<<1, BLOCK_SIZE>>>(
                    partial_sums.data_ptr<float>(),
                    output.data_ptr<float>(),
                    grid_size
                );
                if (reduction == 1) {
                    output.div_(n);
                }
            }

            return {output, partial_sums};
        }

        void focal_loss_backward_cuda(
            float grad_out_scalar,
            torch::Tensor grad_output_n,
            torch::Tensor input,
            torch::Tensor target,
            torch::Tensor grad_input,
            int reduction,
            int64_t n,
            float alpha,
            float gamma,
            float epsilon)
        {
            const int block_size = BLOCK_SIZE;

            // --- FIX IS HERE ---
            // Replaced std::min with ternary operator
            const int A_bwd = (int)((n / 4 + block_size - 1) / block_size);
            const int B_bwd = MAX_GRID_SIZE;
            const int grid_size = (A_bwd < B_bwd ? A_bwd : B_bwd);
            // --- END FIX ---

            const float* grad_output_n_ptr = (reduction == 0) ? 
                grad_output_n.data_ptr<float>() : 
                nullptr;

            focal_bwd_kernel<<<grid_size, block_size>>>(
                (double)grad_out_scalar,
                grad_output_n_ptr,
                input.data_ptr<float>(),
                target.data_ptr<float>(),
                grad_input.data_ptr<float>(),
                reduction,
                n,
                (double)alpha,
                (double)gamma,
                (double)epsilon
            );
        }
        """

        self.op = load_inline(
            name='focal_loss_cuda_v2_vectorized_fix',
            cpp_sources=cpp_source,
            cuda_sources=cuda_source,
            functions=['focal_loss_forward_cuda', 'focal_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 FocalLossCUDAOp.apply(
            input,
            target,
            self.alpha,
            self.gamma,
            self.reduction_id,
            self.op
        )