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

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

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

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

struct __align__(16) Double2 {
    double x, y;
};

__device__ __forceinline__ double compute_meta_aconc(double x, double p1, double p2, double beta) {
    double diff = (p1 - p2) * x;
    double sig = 1.0 / (1.0 + exp(-beta * diff));
    return diff * sig + p2 * x;
}

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

    // Vectorized Loop
    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;
        
        int scalar_idx = i * 2;
        
        #pragma unroll
        for (int k = 0; k < 2; ++k) {
            int curr_idx = scalar_idx + k;
            
            // Calculate Indices for Broadcasting
            // x layout: (N, C, H, W)
            // batch_idx (n) = curr_idx / (C * H * W)
            // rem = curr_idx % (C * H * W)
            // channel_idx (c) = rem / (H * W)
            
            // Optimization using pre-calculated strides
            int n = curr_idx / batch_stride;
            int rem = curr_idx % batch_stride;
            int c = rem / channel_stride;
            
            // Load Parameters
            // p1, p2 are (1, C, 1, 1) -> index by c
            // beta is (N, C, 1, 1)    -> index by n * C + c
            double p1_val = p1[c];
            double p2_val = p2[c];
            double beta_val = beta[n * num_channels + c];
            
            double val = (k == 0) ? x_val.x : x_val.y;
            double res = compute_meta_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;
    }

    int tail_start = vec_loops * 2;
    for (int i = tail_start + idx; i < n_elements; i += stride) {
        int n = i / batch_stride;
        int rem = i % batch_stride;
        int c = rem / channel_stride;
        
        double p1_val = p1[c];
        double p2_val = p2[c];
        // beta flat index: n * C + c
        double beta_val = beta[n * num_channels + c];
        
        output[i] = compute_meta_aconc(input[i], p1_val, p2_val, beta_val);
    }
}

torch::Tensor meta_aconc_cuda_forward(
    const torch::Tensor& x, 
    const torch::Tensor& beta,
    const torch::Tensor& p1, 
    const torch::Tensor& p2) 
{
    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");
    
    // Shapes
    // x: (N, C, H, W)
    int batch_size = x.size(0);
    int num_channels = x.size(1);
    int height = x.size(2);
    int width = x.size(3);
    
    int spatial_size = height * width;
    int channel_stride = spatial_size;
    int batch_stride = num_channels * spatial_size;
    
    int n_elements = x.numel();
    auto output = torch::empty_like(x);

    const int block_size = 256;
    int grid_size = (n_elements + block_size * 2 - 1) / (block_size * 2);
    if (grid_size > 65535) grid_size = 65535;

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

    return output;
}
"""

meta_aconc_op = load_inline(
    name='meta_aconc_op',
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=['meta_aconc_cuda_forward'],
    verbose=False,
    extra_cuda_cflags=['-O3']
)

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

    def forward(self, x, beta):
        return meta_aconc_op.meta_aconc_cuda_forward(x, beta, self.p1, self.p2)

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

    def forward(self, x: torch.Tensor, beta: torch.Tensor) -> torch.Tensor:
        return self.act(x.contiguous(), beta)