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

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


class ModelNew(nn.Module):

    def __init__(self, reduction='mean', eps=1e-3):
        super().__init__()
        self.reduction = reduction
        self.eps = float(eps)

        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.block_size = 256
        self._compile_cuda_kernel()

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

        torch::Tensor charbonnier_forward_cuda(
            torch::Tensor input,
            torch::Tensor target,
            float eps,
            int reduction);
        """

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

        #define BLOCK_SIZE 256
        #define WARP_SIZE 32

        #define ILP 4


        __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]; // max 32 warps
            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;
        }


        __global__ void __launch_bounds__(BLOCK_SIZE) charbonnier_kernel_ilp(
            const float* __restrict__ input,
            const float* __restrict__ target,
            float* __restrict__ output,
            int n,
            float eps_sq,
            int reduction
        ) {

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


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

            int n_vec = n / 4;
            float local_sum = 0.0f;


            int i = idx;

            for (; i + stride * (ILP - 1) < n_vec; i += stride * ILP) {
                float4 r_in[ILP];
                float4 r_tgt[ILP];
                float4 r_res[ILP]; // 仅用于 reduction=none


                #pragma unroll
                for (int k = 0; k < ILP; ++k) {
                    r_in[k] = in_ptr[i + k * stride];
                    r_tgt[k] = tgt_ptr[i + k * stride];
                }


                #pragma unroll
                for (int k = 0; k < ILP; ++k) {
                    float4 res;

                    float d0 = r_in[k].x - r_tgt[k].x;
                    float d1 = r_in[k].y - r_tgt[k].y;
                    float d2 = r_in[k].z - r_tgt[k].z;
                    float d3 = r_in[k].w - r_tgt[k].w;

                    // sqrtf(x^2 + eps^2)
                    res.x = sqrtf(d0 * d0 + eps_sq);
                    res.y = sqrtf(d1 * d1 + eps_sq);
                    res.z = sqrtf(d2 * d2 + eps_sq);
                    res.w = sqrtf(d3 * d3 + eps_sq);

                    if (reduction == 0) { // none
                        r_res[k] = res;
                    } else {
                        local_sum += res.x + res.y + res.z + res.w;
                    }
                }


                if (reduction == 0) {
                    #pragma unroll
                    for (int k = 0; k < ILP; ++k) {
                        out_ptr[i + k * stride] = r_res[k];
                    }
                }
            }


            for (; i < n_vec; i += stride) {
                float4 in_val = in_ptr[i];
                float4 tgt_val = tgt_ptr[i];

                float d0 = in_val.x - tgt_val.x;
                float d1 = in_val.y - tgt_val.y;
                float d2 = in_val.z - tgt_val.z;
                float d3 = in_val.w - tgt_val.w;

                float l0 = sqrtf(d0 * d0 + eps_sq);
                float l1 = sqrtf(d1 * d1 + eps_sq);
                float l2 = sqrtf(d2 * d2 + eps_sq);
                float l3 = sqrtf(d3 * d3 + eps_sq);

                if (reduction == 0) {
                    out_ptr[i] = make_float4(l0, l1, l2, l3);
                } else {
                    local_sum += l0 + l1 + l2 + l3;
                }
            }


            int scalar_start = n_vec * 4;

            int scalar_idx = scalar_start + (blockIdx.x * blockDim.x + threadIdx.x);
            int scalar_stride = blockDim.x * gridDim.x;

            for (int k = scalar_idx; k < n; k += scalar_stride) {
                float d = input[k] - target[k];
                float l = sqrtf(d * d + eps_sq);

                if (reduction == 0) {
                    output[k] = l;
                } else {
                    local_sum += l;
                }
            }


            if (reduction != 0) {
                local_sum = block_reduce_sum(local_sum);
                if (threadIdx.x == 0) {
                    atomicAdd(output, local_sum);
                }
            }
        }

        torch::Tensor charbonnier_forward_cuda(
            torch::Tensor input,
            torch::Tensor target,
            float eps,
            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 max_grid = 2048;

            int n_vec = (n + 3) / 4;
            const int grid_size = std::min((int)((n_vec + BLOCK_SIZE - 1) / BLOCK_SIZE), max_grid);

            float eps_sq = eps * eps;

            charbonnier_kernel_ilp<<<grid_size, BLOCK_SIZE>>>(
                input.data_ptr<float>(),
                target.data_ptr<float>(),
                output.data_ptr<float>(),
                n,
                eps_sq,
                reduction
            );

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

            return output;
        }
        """

        self.op = load_inline(
            name='charbonnier_cuda_opt_v2',
            cpp_sources=cpp_source,
            cuda_sources=cuda_source,
            functions=['charbonnier_forward_cuda'],
            extra_cuda_cflags=['-O3', '--use_fast_math'],
            verbose=False
        )

    def forward(self, input, target):
        if not input.is_cuda: input = input.cuda()
        if not target.is_cuda: target = target.cuda()

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

        return self.op.charbonnier_forward_cuda(
            input,
            target,
            self.eps,
            self.reduction_id
        )