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

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

torch::Tensor pairnorm_cuda_forward(const torch::Tensor& input, const std::string& mode, float scale);
"""

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

#define BLOCK_SIZE 256

__global__ void stats_kernel(
    const float* __restrict__ input,
    double* __restrict__ col_sums,
    double* __restrict__ col_sq_sums,
    int N,
    int D)
{
    int d = blockIdx.x * blockDim.x + threadIdx.x;
    if (d >= D) return;

    double local_sum = 0.0;
    double local_sq_sum = 0.0;

    for (int n = blockIdx.y; n < N; n += gridDim.y) {
        double dval = static_cast<double>(input[n * D + d]);
        local_sum += dval;
        local_sq_sum += dval * dval;
    }

    atomicAdd(&col_sums[d], local_sum);
    atomicAdd(&col_sq_sums[d], local_sq_sum);
}

struct __align__(16) Float4 {
    float x, y, z, w;
};

__global__ void apply_pn_kernel(
    const float* __restrict__ input,
    float* __restrict__ output,
    const float* __restrict__ col_means, 
    float global_factor, 
    int num_elements,
    int D)
{
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int stride = blockDim.x * gridDim.x;
    int vec_loops = num_elements / 4;
    
    const Float4* vec_input = reinterpret_cast<const Float4*>(input);
    Float4* vec_output = reinterpret_cast<Float4*>(output);
    
    for (int i = idx; i < vec_loops; i += stride) {
        Float4 val = vec_input[i];
        Float4 res;
        
        int flat_idx = i * 4;
        #pragma unroll
        for(int k=0; k<4; ++k) {
            int d = (flat_idx + k) % D;
            float m = col_means[d];
            float v = (k==0) ? val.x : (k==1 ? val.y : (k==2 ? val.z : val.w));
            float out_v = (v - m) * global_factor;
            if(k==0) res.x = out_v; else if(k==1) res.y = out_v; else if(k==2) res.z = out_v; else res.w = out_v;
        }
        vec_output[i] = res;
    }
    
    for (int i = vec_loops * 4 + idx; i < num_elements; i += stride) {
        int d = i % D;
        output[i] = (input[i] - col_means[d]) * global_factor;
    }
}

template<int MODE> // 1: PN-SI, 2: PN-SCS
__global__ void apply_row_wise_kernel(
    const float* __restrict__ input,
    float* __restrict__ output,
    const float* __restrict__ col_means,
    float scale,
    int N,
    int D)
{
    // Each block handles one row 'n'
    int n = blockIdx.x;
    if (n >= N) return;

    // Shared memory for block reduction
    __shared__ double sdata[BLOCK_SIZE];

    int tid = threadIdx.x;
    const float* row_in = input + n * D;
    float* row_out = output + n * D;

    double local_sq_sum = 0.0;

    // Step 1: Calculate local Sum Sq
    for (int d = tid; d < D; d += blockDim.x) {
        float val = row_in[d];
        float term;
        if (MODE == 1) { // PN-SI: norm of (x - mean)
            float mean = col_means[d];
            term = val - mean;
        } else { // PN-SCS: norm of (x)
            term = val;
        }
        local_sq_sum += (double)(term * term);
    }
    sdata[tid] = local_sq_sum;
    __syncthreads();

    // Step 2: Block Reduction to find Row Sum Sq
    // Assumes BLOCK_SIZE >= D or D is looped. 
    // Standard tree reduction
    for (unsigned int s = blockDim.x / 2; s > 0; s >>= 1) {
        if (tid < s) {
            sdata[tid] += sdata[tid + s];
        }
        __syncthreads();
    }

    // Step 3: Compute Row Norm factor
    // Using 1e-6 epsilon as per reference
    float row_norm = sqrtf(1e-6f + (float)sdata[0]);
    float inv_norm = scale / row_norm;

    // Step 4: Write Output
    for (int d = tid; d < D; d += blockDim.x) {
        float val = row_in[d];
        float mean = col_means[d];
        
        if (MODE == 1) { // PN-SI: scale * (x - mean) / norm
            row_out[d] = (val - mean) * inv_norm;
        } else { // PN-SCS: scale * x / norm - mean
            row_out[d] = (val * inv_norm) - mean;
        }
    }
}


torch::Tensor pairnorm_cuda_forward(const torch::Tensor& input, const std::string& mode, float scale) {
    if (mode == "None") return input;

    TORCH_CHECK(input.is_cuda(), "Input must be CUDA");
    TORCH_CHECK(input.is_contiguous(), "Input must be contiguous");
    int N = input.size(0);
    int D = input.size(1);

    // --- Common Pass 1: Column Stats ---
    auto options = input.options();
    auto double_options = input.options().dtype(torch::kFloat64);
    
    auto col_sums = torch::zeros({D}, double_options);
    auto col_sq_sums = torch::zeros({D}, double_options);

    dim3 block_stats(256);
    dim3 grid_stats((D + 255) / 256, 64);
    stats_kernel<<<grid_stats, block_stats>>>(
        input.data_ptr<float>(), col_sums.data_ptr<double>(), col_sq_sums.data_ptr<double>(), N, D
    );

    // Prepare Col Means on GPU
    auto col_sums_cpu = col_sums.cpu();
    auto col_means = torch::empty({D}, options);
    float* mean_ptr_cpu = (float*)malloc(D * sizeof(float)); // Temporary host buffer
    const double* sum_ptr = col_sums_cpu.data_ptr<double>();
    
    double inv_n = 1.0 / N;
    for(int i=0; i<D; ++i) mean_ptr_cpu[i] = (float)(sum_ptr[i] * inv_n);
    
    // Copy computed means back to device
    cudaMemcpy(col_means.data_ptr<float>(), mean_ptr_cpu, D * sizeof(float), cudaMemcpyHostToDevice);
    free(mean_ptr_cpu);

    auto output = torch::empty_like(input);

    // --- Pass 2: Dispatch based on Mode ---
    if (mode == "PN") {
        // Calculate PN global scale factor on CPU
        auto col_sq_sums_cpu = col_sq_sums.cpu();
        const double* sq_sum_ptr = col_sq_sums_cpu.data_ptr<double>();
        
        double total_centered_sq_sum = 0.0;
        for(int d=0; d<D; ++d) {
            double s = sum_ptr[d];
            double ss = sq_sum_ptr[d];
            // sum((x-u)^2) = ss - s*s/N
            total_centered_sq_sum += (ss - (s * s * inv_n));
        }
        
        double mean_row_norm = total_centered_sq_sum * inv_n;
        float rownorm_mean = sqrt(1e-6 + mean_row_norm);
        float global_factor = scale / rownorm_mean;

        int num_elements = N * D;
        int block = 256;
        int grid = (num_elements / 4 + 255) / 256;
        if(grid > 65535) grid = 65535;
        
        apply_pn_kernel<<<grid, block>>>(
            input.data_ptr<float>(), output.data_ptr<float>(), 
            col_means.data_ptr<float>(), global_factor, num_elements, D
        );

    } else if (mode == "PN-SI") {
        dim3 block(256); // Assuming D <= 256 or handled by loop
        dim3 grid(N);    // One block per row
        
        // Template instantiation for MODE=1
        apply_row_wise_kernel<1><<<grid, block>>>(
            input.data_ptr<float>(), output.data_ptr<float>(),
            col_means.data_ptr<float>(), scale, N, D
        );
        
    } else if (mode == "PN-SCS") {
        dim3 block(256);
        dim3 grid(N);
        
        // Template instantiation for MODE=2
        apply_row_wise_kernel<2><<<grid, block>>>(
            input.data_ptr<float>(), output.data_ptr<float>(),
            col_means.data_ptr<float>(), scale, N, D
        );
    }

    return output;
}
"""

pairnorm_op = load_inline(
    name='pairnorm_full_op',
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=['pairnorm_cuda_forward'],
    verbose=False,
    extra_cuda_cflags=['-O3']
)

class PairNormNew(nn.Module):
    def __init__(self, mode, scale=1.0):
        super(PairNormNew, self).__init__()
        self.mode = mode
        self.scale = scale

    def forward(self, x):
        return pairnorm_op.pairnorm_cuda_forward(x, self.mode, self.scale)

class ModelNew(nn.Module):
    def __init__(self, mode, scale=1.0):
        super(ModelNew, self).__init__()
        self.norm = PairNormNew(mode=mode, scale=scale)

    def forward(self, x):
        return self.norm(x)