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


class ModelNew(nn.Module):
    def __init__(self, reduction='mean'):
        super().__init__()
        self.reduction = reduction
        self._compile_cuda_kernel()

    def _compile_cuda_kernel(self):
        cpp_source = """
        torch::Tensor mse_loss_cuda(torch::Tensor input, torch::Tensor target);
        """

        cuda_source = """
        #include <torch/extension.h>
        #include <cuda_runtime.h>

        __global__ void mse_loss_kernel_vectorized(
            const float* __restrict__ input,
            const float* __restrict__ target,
            float* __restrict__ output,
            const int64_t n)
        {
            const int64_t idx = blockIdx.x * blockDim.x + threadIdx.x;
            const int64_t stride = blockDim.x * gridDim.x;

            const int64_t n_vec = n / 4;
            const float4* in_vec = reinterpret_cast<const float4*>(input);
            const float4* tgt_vec = reinterpret_cast<const float4*>(target);
            float4* out_vec = reinterpret_cast<float4*>(output);

            for (int64_t i = idx; i < n_vec; i += stride) {
                float4 a = in_vec[i];
                float4 b = tgt_vec[i];
                float4 res;
                float d;

                d = __fsub_rn(a.x, b.x); res.x = __fmul_rn(d, d);
                d = __fsub_rn(a.y, b.y); res.y = __fmul_rn(d, d);
                d = __fsub_rn(a.z, b.z); res.z = __fmul_rn(d, d);
                d = __fsub_rn(a.w, b.w); res.w = __fmul_rn(d, d);

                out_vec[i] = res;
            }

            const int64_t tail_offset = n_vec * 4;
            for (int64_t i = tail_offset + idx; i < n; i += stride) {
                float diff = __fsub_rn(input[i], target[i]);
                output[i] = __fmul_rn(diff, diff);
            }
        }

        torch::Tensor mse_loss_cuda(torch::Tensor input, torch::Tensor target) {
            TORCH_CHECK(input.is_cuda(), "Input tensor must be on CUDA");
            TORCH_CHECK(target.is_cuda(), "Target tensor must be on CUDA");
            TORCH_CHECK(input.numel() == target.numel(), "Input and target must have the same number of elements");

            auto input_c = input.contiguous();
            auto target_c = target.contiguous();
            const int64_t n = input_c.numel();

            auto output = torch::empty_like(input_c);

            const int threads = 256;
            const int blocks = min((int64_t)((n + threads * 4 - 1) / (threads * 4)), (int64_t)65535);

            mse_loss_kernel_vectorized<<<blocks, threads>>>(
                input_c.data_ptr<float>(),
                target_c.data_ptr<float>(),
                output.data_ptr<float>(),
                n
            );

            return output;
        }
        """

        self.op = load_inline(
            name="perceptual_loss_opt_op",
            cpp_sources=cpp_source,
            cuda_sources=cuda_source,
            functions=["mse_loss_cuda"],
            extra_cuda_cflags=["-O3", "-fmad=false"],
            verbose=False
        )

    def forward(self, input_features: list[torch.Tensor], target_features: list[torch.Tensor]) -> torch.Tensor:
        if len(input_features) > 0 and not input_features[0].is_cuda:
            input_features = [f.cuda() for f in input_features]
            target_features = [f.cuda() for f in target_features]

        loss_val = 0.0

        for i in range(len(input_features)):
            squared_diff = self.op.mse_loss_cuda(input_features[i], target_features[i])

            if self.reduction == 'mean':
                loss_val += squared_diff.mean()
            elif self.reduction == 'sum':
                loss_val += squared_diff.sum()
            else:
                loss_val += squared_diff.mean()

        return loss_val