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

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


class CauchyLossCUDAOp(torch.autograd.Function):

    def forward(ctx, input, target, beta, 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()

        output = op.cauchy_loss_forward_cuda(
            input,
            target,
            beta,
            reduction_id
        )
        return output

    def backward(ctx, grad_output):
        return grad_output, grad_output, None, 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.vector_size = 4  # float4
        self._compile_cuda_kernel()

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

        torch::Tensor cauchy_loss_forward_cuda(
            torch::Tensor input,
            torch::Tensor target,
            float beta,
            int reduction);
        """

        cuda_source = """
        #include <cuda_runtime.h>
        #include <cmath>
        #include <device_functions.h>
        #include <math_functions.h>

        #define BLOCK_SIZE 256

        // Reduction functions (unchanged)
        __inline__ __device__ float warp_reduce_sum(float val) {
            #pragma unroll
            for (int offset = 16; offset > 0; offset /= 2) {
                val += __shfl_down_sync(0xffffffff, val, offset);
            }
            return val;
        }

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

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

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

        // 进一步优化后的 CUDA 内核：简化向量化索引
        __global__ void cauchy_loss_kernel(
            const float* __restrict__ input,
            const float* __restrict__ target,
            float* __restrict__ output,
            int n,
            float beta,
            int reduction
        ) {
            // VEC_SIZE = 4
            const int VEC_SIZE = 4;
            int vec_n = n / VEC_SIZE;
            int rem_start = vec_n * VEC_SIZE;

            // 将索引和步长定义在向量空间
            int vec_idx = blockIdx.x * blockDim.x + threadIdx.x;
            int vec_stride = blockDim.x * gridDim.x;

            float local_sum = 0.0f;

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


            for (int i = vec_idx; i < vec_n; i += vec_stride) {
                float4 in_val = in_ptr[i];
                float4 tgt_val = tgt_ptr[i];
                float4 out_val;
                float losses[4];

                float diff[4];
                diff[0] = fabsf(in_val.x - tgt_val.x);
                diff[1] = fabsf(in_val.y - tgt_val.y);
                diff[2] = fabsf(in_val.z - tgt_val.z);
                diff[3] = fabsf(in_val.w - tgt_val.w);

                #pragma unroll
                for(int k=0; k<4; ++k) {
                    // 使用快速数学函数
                    float ratio = __fdividef(diff[k], beta);
                    float normalized_diff_sq = ratio * ratio;
                    losses[k] = __logf(1.0f + normalized_diff_sq);
                }

                if (reduction == 0) { // none
                    out_val.x = losses[0];
                    out_val.y = losses[1];
                    out_val.z = losses[2];
                    out_val.w = losses[3];
                    out_ptr[i] = out_val;
                } else {
                    local_sum += losses[0] + losses[1] + losses[2] + losses[3];
                }
            }


            int rem_idx = rem_start + threadIdx.x;
            int rem_stride = blockDim.x;

            for (int i = rem_idx; i < n; i += rem_stride) {
                float diff = fabsf(input[i] - target[i]);

                float ratio = __fdividef(diff, beta);
                float normalized_diff_sq = ratio * ratio;
                float loss = __logf(1.0f + normalized_diff_sq);

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


            if (reduction != 0) {
                // 块内归约
                local_sum = block_reduce_sum(local_sum);
                if (threadIdx.x == 0) {
                    // 原子操作加到全局输出
                    atomicAdd(output, local_sum);
                }
            }
        }

        torch::Tensor cauchy_loss_forward_cuda(
            torch::Tensor input,
            torch::Tensor target,
            float beta,
            int reduction) 
        {
            int64_t n = input.numel();
            auto options = input.options();

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

            const int block_size = 256;
            // grid size 现在是基于向量化的元素数量 n/4
            const int grid_size = std::min((int)((n / 4 + block_size - 1) / block_size), 1024);

            // 确保 grid size 至少为 1
            const int final_grid_size = std::max(grid_size, 1); 

            cauchy_loss_kernel<<<final_grid_size, block_size>>>(
                input.data_ptr<float>(),
                target.data_ptr<float>(),
                output.data_ptr<float>(),
                n,
                beta,
                reduction
            );

            if (reduction == 1) { // mean
                output.div_(n);
            }

            return output;
        }
        """

        self.op = load_inline(
            name='cauchy_loss_cuda_opt',
            cpp_sources=cpp_source,
            cuda_sources=cuda_source,
            functions=['cauchy_loss_forward_cuda'],
            extra_cuda_cflags=['-O3', '--use_fast_math'],
            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 not input.is_cuda: input = input.cuda()
        if not target.is_cuda: target = target.cuda()

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

        return CauchyLossCUDAOp.apply(
            input,
            target,
            self.beta,
            self.reduction_id,
            self.op
        )
