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

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


class GeneratorLossCUDAOp(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()

        n = input.numel()

        output, partial_sums = op.generator_loss_forward_cuda(
            input,
            target,
            reduction_id,
            n
        )

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

        return output

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

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

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

        ctx.op.generator_loss_backward_cuda(
            grad_out_scalar,
            input,
            target,
            grad_input,
            grad_target,
            ctx.N
        )

        return grad_input, grad_target, 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>
        #include <vector>

        std::vector<torch::Tensor> generator_loss_forward_cuda(
            torch::Tensor input,
            torch::Tensor target,
            int reduction,
            int64_t n);

        void generator_loss_backward_cuda(
            float grad_out_scalar,
            torch::Tensor input,
            torch::Tensor target,
            torch::Tensor grad_input,
            torch::Tensor grad_target,
            int64_t n);
        """

        cuda_source = """
        #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__ 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;
        }

        __global__ void generator_loss_fwd_kernel(
            const float* __restrict__ input,
            const float* __restrict__ target,
            float* __restrict__ output,
            int64_t n,
            int reduction,
            float* __restrict__ partial_sums) 
        {
            int i = blockIdx.x * blockDim.x + threadIdx.x;
            int stride = blockDim.x * gridDim.x;

            float local_sum = 0.0f;

            for (int idx = i; idx < n; idx += stride) {
                float x = input[idx];
                float t = target[idx];

                float maxi = (x > 0.0f) ? x : 0.0f;
                float loss = maxi - x * t + __logf(1.0f + __expf(-fabsf(x)));

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

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

        __global__ void reduce_kernel(
            const float* __restrict__ partial_sums,
            float* __restrict__ output,
            int n_partials) 
        {
            float local_sum = 0.0f;
            int idx = threadIdx.x;

            for (int i = idx; i < n_partials; i += blockDim.x) {
                local_sum += partial_sums[i];
            }

            local_sum = block_reduce_sum(local_sum);

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

        __global__ void generator_loss_bwd_kernel(
            float grad_out_scalar,
            const float* __restrict__ input,
            const float* __restrict__ target,
            float* __restrict__ grad_input,
            float* __restrict__ grad_target,
            int64_t n)
        {
            int i = blockIdx.x * blockDim.x + threadIdx.x;
            int stride = blockDim.x * gridDim.x;

            for (int idx = i; idx < n; idx += stride) {
                float x = input[idx];
                float t = target[idx];

                float sigma_x = 1.0f / (1.0f + __expf(-x));

                grad_input[idx] = (sigma_x - t) * grad_out_scalar;
                grad_target[idx] = (-x) * grad_out_scalar;
            }
        }

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

            const int block_size = BLOCK_SIZE;
            const int grid_size = std::min(
                (int)((n + block_size - 1) / block_size), 
                MAX_GRID_SIZE
            );

            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);
            }

            generator_loss_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>()
            );

            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 generator_loss_backward_cuda(
            float grad_out_scalar,
            torch::Tensor input,
            torch::Tensor target,
            torch::Tensor grad_input,
            torch::Tensor grad_target,
            int64_t n)
        {
            const int block_size = BLOCK_SIZE;
            const int grid_size = std::min(
                (int)((n + block_size - 1) / block_size), 
                MAX_GRID_SIZE
            );

            generator_loss_bwd_kernel<<<grid_size, block_size>>>(
                grad_out_scalar,
                input.data_ptr<float>(),
                target.data_ptr<float>(),
                grad_input.data_ptr<float>(),
                grad_target.data_ptr<float>(),
                n
            );
        }
        """

        self.op = load_inline(
            name='generator_loss_cuda_v1',
            cpp_sources=cpp_source,
            cuda_sources=cuda_source,
            functions=['generator_loss_forward_cuda', 'generator_loss_backward_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

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