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

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

__global__ void ssim_softplus_kernel(const float* __restrict__ x, 
                                     const float* __restrict__ y, 
                                     float* __restrict__ out, 
                                     int dim) {
    int bid = blockIdx.x;
    int tid = threadIdx.x;

    const float* row_x = x + bid * dim;
    const float* row_y = y + bid * dim;

    float sum_x = 0.0f;
    float sum_y = 0.0f;
    float sum_sq_x = 0.0f;
    float sum_sq_y = 0.0f;
    float sum_xy = 0.0f;

    for (int i = tid; i < dim; i += blockDim.x) {
        float val_x = row_x[i];
        float val_y = row_y[i];

        sum_x += val_x;
        sum_y += val_y;
        sum_sq_x += val_x * val_x;
        sum_sq_y += val_y * val_y;
        sum_xy += val_x * val_y;
    }

    extern __shared__ float s_mem[];
    float* s_sum_x = s_mem;
    float* s_sum_y = s_mem + 256;
    float* s_sum_sq_x = s_mem + 512;
    float* s_sum_sq_y = s_mem + 768;
    float* s_sum_xy = s_mem + 1024;

    s_sum_x[tid] = sum_x;
    s_sum_y[tid] = sum_y;
    s_sum_sq_x[tid] = sum_sq_x;
    s_sum_sq_y[tid] = sum_sq_y;
    s_sum_xy[tid] = sum_xy;
    __syncthreads();

    for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) {
        if (tid < stride) {
            s_sum_x[tid] += s_sum_x[tid + stride];
            s_sum_y[tid] += s_sum_y[tid + stride];
            s_sum_sq_x[tid] += s_sum_sq_x[tid + stride];
            s_sum_sq_y[tid] += s_sum_sq_y[tid + stride];
            s_sum_xy[tid] += s_sum_xy[tid + stride];
        }
        __syncthreads();
    }

    if (tid == 0) {
        float N = (float)dim;
        float mu_x = s_sum_x[0] / N;
        float mu_y = s_sum_y[0] / N;
        float var_x = (s_sum_sq_x[0] / N) - (mu_x * mu_x);
        float var_y = (s_sum_sq_y[0] / N) - (mu_y * mu_y);
        float cov_xy = (s_sum_xy[0] / N) - (mu_x * mu_y);

        float C1 = 0.0001f;
        float C2 = 0.0009f;

        float num = (2.0f * mu_x * mu_y + C1) * (2.0f * cov_xy + C2);
        float den = (mu_x * mu_x + mu_y * mu_y + C1) * (var_x + var_y + C2);
        float ssim = num / den;

        out[bid] = log1pf(expf(ssim));
    }
}

torch::Tensor ssim_softplus_cuda(torch::Tensor x, torch::Tensor y) {
    int batch_size = x.size(0);
    int dim = x.size(1);
    auto out = torch::empty({batch_size}, x.options());

    int shared_size = 5 * 256 * sizeof(float);

    ssim_softplus_kernel<<<batch_size, 256, shared_size>>>(
        x.data_ptr<float>(),
        y.data_ptr<float>(),
        out.data_ptr<float>(),
        dim
    );

    return out;
}
"""

cpp_source = "torch::Tensor ssim_softplus_cuda(torch::Tensor x, torch::Tensor y);"

module = load_inline(
    name="structural_similarity_softplus_ext",
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=["ssim_softplus_cuda"],
    verbose=False,
    with_cuda=True
)


class ModelNew(nn.Module):
    def __init__(self):
        super(ModelNew, self).__init__()
        self.op = module

    def forward(self, x, y):
        res = self.op.ssim_softplus_cuda(x.contiguous(), y.contiguous())
        return res.mean()