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

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


class WelschLossCUDAOp(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.welsch_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._compile_cuda_kernel()

    def _compile_cuda_kernel(self):
        cpp_source = """
        #include <torch/extension.h>
        torch::Tensor welsch_loss_forward_cuda(torch::Tensor input, torch::Tensor target, float beta, int reduction);
        """

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

        __inline__ __device__ float warp_reduce(float v) {
            #pragma unroll
            for (int d = 16; d > 0; d >>= 1) v += __shfl_down_sync(0xffffffff, v, d);
            return v;
        }

        __inline__ __device__ float block_reduce(float v) {
            __shared__ float s[32];
            int l = threadIdx.x & 31;
            int w = threadIdx.x >> 5;
            v = warp_reduce(v);
            if (l == 0) s[w] = v;
            __syncthreads();
            v = (threadIdx.x < (blockDim.x >> 5)) ? s[l] : 0.0f;
            if (w == 0) v = warp_reduce(v);
            return v;
        }

        __global__ void welsch_kernel(const float* __restrict__ in, const float* __restrict__ tgt, float* __restrict__ out, int n, float k, int red) {
            int tid = blockIdx.x * blockDim.x + threadIdx.x;
            int stride = blockDim.x * gridDim.x;
            int n4 = n >> 2;
            float acc = 0.0f;

            const float4* in4 = (const float4*)in;
            const float4* tgt4 = (const float4*)tgt;
            float4* out4 = (float4*)out;

            for (int i = tid; i < n4; i += stride) {
                float4 iv = __ldg(&in4[i]);
                float4 tv = __ldg(&tgt4[i]);

                float dx = iv.x - tv.x;
                float dy = iv.y - tv.y;
                float dz = iv.z - tv.z;
                float dw = iv.w - tv.w;

                float lx = 1.0f - __expf(__fmul_rn(__fmul_rn(dx, dx), k));
                float ly = 1.0f - __expf(__fmul_rn(__fmul_rn(dy, dy), k));
                float lz = 1.0f - __expf(__fmul_rn(__fmul_rn(dz, dz), k));
                float lw = 1.0f - __expf(__fmul_rn(__fmul_rn(dw, dw), k));

                if (red == 0) {
                    out4[i] = make_float4(lx, ly, lz, lw);
                } else {
                    acc += lx + ly + lz + lw;
                }
            }

            int rem = n4 << 2;
            for (int i = rem + threadIdx.x; i < n; i += blockDim.x) {
                float d = in[i] - tgt[i];
                float l = 1.0f - __expf(__fmul_rn(__fmul_rn(d, d), k));
                if (red == 0) out[i] = l;
                else acc += l;
            }

            if (red != 0) {
                acc = block_reduce(acc);
                if (threadIdx.x == 0) out[blockIdx.x] = acc;
            }
        }

        torch::Tensor welsch_loss_forward_cuda(torch::Tensor input, torch::Tensor target, float beta, int reduction) {
            int64_t n = input.numel();
            auto opts = input.options();
            float k = -1.0f / (beta * beta);

            const int bs = 256;
            const int gs = std::min((int)((n + (bs << 2) - 1) / (bs << 2)), 1024);
            const int fgs = std::max(gs, 1);

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

            welsch_kernel<<<fgs, bs>>>(input.data_ptr<float>(), target.data_ptr<float>(), output.data_ptr<float>(), n, k, reduction);

            if (reduction != 0) {
                float s = output.sum().item<float>();
                if (reduction == 1) s /= n;
                return torch::tensor({s}, opts);
            }
            return output;
        }
        """

        self.op = load_inline(
            name='welsch_loss_cuda_v2',
            cpp_sources=cpp_source,
            cuda_sources=cuda_source,
            functions=['welsch_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)): input = input[0]
        if isinstance(target, (list, tuple)): target = target[0]
        if not input.is_cuda: input = input.cuda()
        if not target.is_cuda: target = target.cuda()
        input = input.contiguous()
        target = target.contiguous()
        return WelschLossCUDAOp.apply(input, target, self.beta, self.reduction_id, self.op)