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

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

torch::Tensor scalenorm_cuda_forward(
    const torch::Tensor& input,
    const torch::Tensor& weight,
    const torch::Tensor& bias,
    double eps);
"""

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

#define BLOCK_SIZE 256
#define ILP 4

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

template<typename T>
__device__ __forceinline__ T warp_reduce_sum(T val) {
    #pragma unroll
    for (int offset = 16; offset > 0; offset /= 2) {
        val += __shfl_down_sync(0xffffffff, val, offset);
    }
    return val;
}

__device__ __forceinline__ double block_reduce_sum(double val) {
    static __shared__ double shared[32];
    int lane = threadIdx.x % 32;
    int wid = threadIdx.x / 32;

    val = warp_reduce_sum(val);
    if (lane == 0) shared[wid] = val;
    __syncthreads();

    val = (threadIdx.x < blockDim.x / 32) ? shared[lane] : 0.0;
    if (wid == 0) val = warp_reduce_sum(val);
    
    return val;
}

__global__ void scalenorm_kernel(
    double* __restrict__ output,
    const double* __restrict__ input,
    const double* __restrict__ weight,
    const double* __restrict__ bias, // Can be nullptr
    const double eps,
    const int cols)
{
    int row_idx = blockIdx.x;
    int tid = threadIdx.x;

    int offset = row_idx * cols;
    const double* row_in = input + offset;
    double* row_out = output + offset;

    // Register Cache
    Double2 cache[ILP];

    // --- Pass 1: Load Input and Compute Sum of Squares (L2 Norm part) ---
    double local_sum_sq = 0.0;
    int idx = tid * 2;

    #pragma unroll
    for (int k = 0; k < ILP; ++k) {
        cache[k].x = 0.0; cache[k].y = 0.0;
        
        if (idx < cols) {
            if (idx + 1 < cols) {
                // Vectorized load assuming alignment
                Double2 v = *(reinterpret_cast<const Double2*>(&row_in[idx]));
                cache[k] = v;
            } else {
                cache[k].x = row_in[idx];
            }
            // Accumulate sum squares
            local_sum_sq += cache[k].x * cache[k].x + cache[k].y * cache[k].y;
        }
        idx += BLOCK_SIZE * 2;
    }

    // Global Reduction
    double sum_sq_total = block_reduce_sum(local_sum_sq);
    
    __shared__ double s_inv_norm;
    if (tid == 0) {
        // Source: norm = sqrt(sum(x^2))
        // scale = 1 / (norm + eps)
        double norm = sqrt(sum_sq_total);
        s_inv_norm = 1.0 / (norm + eps);
    }
    __syncthreads();
    
    double inv_norm = s_inv_norm;

    // --- Pass 2: Apply Scale and Shift (bias) ---
    idx = tid * 2;
    #pragma unroll
    for (int k = 0; k < ILP; ++k) {
        if (idx < cols) {
            Double2 w_vec = {1.0, 1.0};
            Double2 b_vec = {0.0, 0.0};
            
            // Load weight/bias (Broadcasted across batch rows)
            if (idx + 1 < cols) {
                w_vec = *(reinterpret_cast<const Double2*>(&weight[idx]));
                if (bias != nullptr) {
                    b_vec = *(reinterpret_cast<const Double2*>(&bias[idx]));
                }
            } else {
                w_vec.x = weight[idx];
                if (bias != nullptr) b_vec.x = bias[idx];
            }

            Double2 out_val;
            // output = weight * input * inv_norm + bias
            out_val.x = cache[k].x * inv_norm * w_vec.x + b_vec.x;
            out_val.y = cache[k].y * inv_norm * w_vec.y + b_vec.y;

            if (idx + 1 < cols) {
                *(reinterpret_cast<Double2*>(&row_out[idx])) = out_val;
            } else {
                row_out[idx] = out_val.x;
            }
        }
        idx += BLOCK_SIZE * 2;
    }
}

torch::Tensor scalenorm_cuda_forward(
    const torch::Tensor& input,
    const torch::Tensor& weight,
    const torch::Tensor& bias,
    double eps) 
{
    TORCH_CHECK(input.is_cuda(), "Input must be CUDA");
    TORCH_CHECK(weight.is_cuda(), "Weight must be CUDA");
    TORCH_CHECK(input.scalar_type() == torch::kDouble, "Input must be float64");
    TORCH_CHECK(input.is_contiguous(), "Input must be contiguous");
    
    int batch_size = input.size(0);
    int dim = input.size(1);
    
    auto output = torch::empty_like(input);
    
    dim3 grid(batch_size);
    dim3 block(BLOCK_SIZE);
    
    const double* bias_ptr = nullptr;
    if (bias.defined()) {
        TORCH_CHECK(bias.is_cuda(), "Bias must be CUDA");
        bias_ptr = bias.data_ptr<double>();
    }

    scalenorm_kernel<<<grid, block>>>(
        output.data_ptr<double>(),
        input.data_ptr<double>(),
        weight.data_ptr<double>(),
        bias_ptr,
        eps,
        dim
    );
    
    return output;
}
"""

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

class ScaleNormNew(nn.Module):
    def __init__(self, normalized_shape, eps, bias):
        super(ScaleNormNew, self).__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.empty(normalized_shape, dtype=torch.float64))
        if bias:
            self.bias = nn.Parameter(torch.empty(normalized_shape, dtype=torch.float64))
        else:
            self.register_parameter('bias', None)
        
        # Reset parameters
        nn.init.ones_(self.weight)
        if self.bias is not None:
            nn.init.zeros_(self.bias)

    def forward(self, input):
        # Handle optional bias tensor
        bias_t = self.bias if self.bias is not None else torch.Tensor()
        return scalenorm_op.scalenorm_cuda_forward(input, self.weight, bias_t, self.eps)

class ModelNew(nn.Module):
    def __init__(self, normalized_shape, eps, bias):
        super(ModelNew, self).__init__()
        self.norm = ScaleNormNew(normalized_shape, eps, bias)
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.norm(x)