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

cpp_source = """
#include <torch/extension.h>

torch::Tensor aconc_cuda_forward(
    const torch::Tensor& x, 
    const torch::Tensor& p1, 
    const torch::Tensor& p2, 
    const torch::Tensor& beta);
"""

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

// Vectorized type for 128-bit access with doubles
// 2 doubles = 16 bytes = 128 bits
struct __align__(16) Double2 {
    double x, y;
};

__device__ __forceinline__ double compute_aconc(double x, double p1, double p2, double beta) {
    double diff = (p1 - p2) * x;
    // use exp() for double, not __expf()
    double sig = 1.0 / (1.0 + exp(-beta * diff));
    return diff * sig + p2 * x;
}

__global__ void aconc_kernel_double(
    const double* __restrict__ input,
    double* __restrict__ output,
    const double* __restrict__ p1,
    const double* __restrict__ p2,
    const double* __restrict__ beta,
    const int num_channels,
    const int spatial_size, // H * W
    const int n_elements) 
{
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int stride = blockDim.x * gridDim.x;

    // Vectorized Loop: Process 2 doubles per iteration (128 bits)
    int vec_loops = n_elements / 2;
    const Double2* vec_input = reinterpret_cast<const Double2*>(input);
    Double2* vec_output = reinterpret_cast<Double2*>(output);

    for (int i = idx; i < vec_loops; i += stride) {
        Double2 x_val = vec_input[i];
        Double2 y_val;
        
        // Scalar index of the first element in the vector
        int scalar_idx = i * 2;
        
        #pragma unroll
        for (int k = 0; k < 2; ++k) {
            int curr_idx = scalar_idx + k;
            // Calculate Channel Index
            int c = (curr_idx / spatial_size) % num_channels;
            
            double p1_val = p1[c];
            double p2_val = p2[c];
            double beta_val = beta[c];
            
            double val = (k == 0) ? x_val.x : x_val.y;
            double res = compute_aconc(val, p1_val, p2_val, beta_val);
            
            if (k == 0) y_val.x = res;
            else y_val.y = res;
        }
        
        vec_output[i] = y_val;
    }

    // Scalar Tail Loop
    int tail_start = vec_loops * 2;
    for (int i = tail_start + idx; i < n_elements; i += stride) {
        int c = (i / spatial_size) % num_channels;
        output[i] = compute_aconc(input[i], p1[c], p2[c], beta[c]);
    }
}

torch::Tensor aconc_cuda_forward(
    const torch::Tensor& x, 
    const torch::Tensor& p1, 
    const torch::Tensor& p2, 
    const torch::Tensor& beta) 
{
    TORCH_CHECK(x.is_cuda(), "Input x must be CUDA");
    TORCH_CHECK(x.scalar_type() == torch::kDouble, "Input must be float64");
    TORCH_CHECK(x.is_contiguous(), "Input x must be contiguous");
    
    int num_channels = x.size(1);
    int spatial_size = 1;
    for (int i = 2; i < x.dim(); ++i) {
        spatial_size *= x.size(i);
    }
    
    int n_elements = x.numel();
    auto output = torch::empty_like(x);

    const int block_size = 256;
    // Grid size calculation adjusted for 2 elements per thread
    int grid_size = (n_elements + block_size * 2 - 1) / (block_size * 2);
    if (grid_size > 65535) grid_size = 65535;

    aconc_kernel_double<<<grid_size, block_size>>>(
        x.data_ptr<double>(),
        output.data_ptr<double>(),
        p1.data_ptr<double>(),
        p2.data_ptr<double>(),
        beta.data_ptr<double>(),
        num_channels,
        spatial_size,
        n_elements
    );

    return output;
}
"""

aconc_op = load_inline(
    name='aconc_op_double',
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=['aconc_cuda_forward'],
    verbose=False,
    extra_cuda_cflags=['-O3'] 
)

class AconCNew(nn.Module):
    def __init__(self, p1, p2, beta):
        super(AconCNew, self).__init__()
        self.p1 = nn.Parameter(p1)
        self.p2 = nn.Parameter(p2)
        self.beta = nn.Parameter(beta)

    def forward(self, x):
        return aconc_op.aconc_cuda_forward(x, self.p1, self.p2, self.beta)

class ModelNew(nn.Module):
    def __init__(self, p1, p2, beta):
        super(ModelNew, self).__init__()
        self.act = AconCNew(p1, p2, beta)

    def forward(self, x):
        return self.act(x.contiguous())