# switchablenorm_cuda.py
import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline

from switchablenorm_torch import N, C, H, W, EPS

class ModelNew(nn.Module):

    def __init__(self, weight, bias, w_in, w_ln, w_bn):
        super().__init__()
        self.register_buffer('weight', weight)
        self.register_buffer('bias', bias)
        self.register_buffer('w_in', w_in)
        self.register_buffer('w_ln', w_ln)
        self.register_buffer('w_bn', w_bn)
        self.eps = EPS
        self.register_buffer('running_mean', torch.zeros(C))
        self.register_buffer('running_var', torch.ones(C))
        self._compile_cuda_kernel()

    def _compile_cuda_kernel(self):
        cpp_source = """
        #include <torch/extension.h>

        torch::Tensor sn_forward_cuda(
            torch::Tensor input, torch::Tensor weight, torch::Tensor bias, 
            torch::Tensor w_in, torch::Tensor w_ln, torch::Tensor w_bn,
            torch::Tensor running_mean, torch::Tensor running_var,
            float eps, int N, int C, int H, int W);
        """
        
        cuda_source = f"""
        #include <cuda_runtime.h>
        #include <cmath>
        #include <torch/extension.h>

        #define BLOCK_SIZE 256
        #define WARP_SIZE 32

        __device__ __forceinline__ float warp_reduce_sum(float val) {{
            for (int offset = WARP_SIZE / 2; offset > 0; offset /= 2)
                val += __shfl_down_sync(0xffffffff, val, offset);
            return val;
        }}

        __global__ void sn_normalize_kernel(
            const float* __restrict__ x,
            const float* __restrict__ ln_mean_ptr, const float* __restrict__ ln_var_ptr,
            const float* __restrict__ in_mean_ptr, const float* __restrict__ in_var_ptr,
            const float* __restrict__ bn_mean_ptr, const float* __restrict__ bn_var_ptr,
            const float* __restrict__ weight_ptr, const float* __restrict__ bias_ptr,
            const float* __restrict__ w_in_ptr, const float* __restrict__ w_ln_ptr, const float* __restrict__ w_bn_ptr,
            float* __restrict__ output,
            int N, int C, int H, int W, float eps
        ) {{
            int n_elements = N * C * H * W;
            int idx = blockIdx.x * blockDim.x + threadIdx.x;

            if (idx >= n_elements) return;

            int nc_idx = idx / (H * W); 
            int sample_idx = nc_idx / C; 
            int channel_idx = nc_idx % C; 

            float w_in = w_in_ptr[channel_idx];
            float w_ln = w_ln_ptr[channel_idx];
            float w_bn = w_bn_ptr[channel_idx];
     
            float w_sum = fabsf(w_in) + fabsf(w_ln) + fabsf(w_bn);
            float w_in_norm = fabsf(w_in) / w_sum;
            float w_ln_norm = fabsf(w_ln) / w_sum;
            float w_bn_norm = fabsf(w_bn) / w_sum;

            float mean_ln = ln_mean_ptr[sample_idx]; 
            float var_ln = ln_var_ptr[sample_idx]; 

            float mean_in = in_mean_ptr[nc_idx]; 
            float var_in = in_var_ptr[nc_idx]; 

            float mean_bn = bn_mean_ptr[channel_idx];
            float var_bn = bn_var_ptr[channel_idx];

            float agg_mean = w_in_norm * mean_in + w_ln_norm * mean_ln + w_bn_norm * mean_bn;

            float m2_in = var_in + mean_in * mean_in;
            float m2_ln = var_ln + mean_ln * mean_ln;
            float m2_bn = var_bn + mean_bn * mean_bn;
            
            float agg_m2 = w_in_norm * m2_in + w_ln_norm * m2_ln + w_bn_norm * m2_bn;
            float agg_var = agg_m2 - agg_mean * agg_mean;

            float val = x[idx];
            float gamma = weight_ptr[channel_idx];
            float beta = bias_ptr[channel_idx];

            float inv_std = rsqrtf(agg_var + eps);
            float normalized = (val - agg_mean) * inv_std;
            
            output[idx] = normalized * gamma + beta;
        }}


        torch::Tensor sn_forward_cuda(
            torch::Tensor input, torch::Tensor weight, torch::Tensor bias, 
            torch::Tensor w_in, torch::Tensor w_ln, torch::Tensor w_bn,
            torch::Tensor running_mean, torch::Tensor running_var,
            float eps, int N, int C, int H, int W) {{
            
  
            torch::Tensor in_mean = input.mean({{2, 3}}, true).squeeze(-1).squeeze(-1).contiguous();
            torch::Tensor in_var = input.var({{2, 3}}, true).squeeze(-1).squeeze(-1).contiguous();

            torch::Tensor ln_mean = input.mean({{1, 2, 3}}, true).squeeze(-1).squeeze(-1).squeeze(-1).contiguous();
            torch::Tensor ln_var = input.var({{1, 2, 3}}, true).squeeze(-1).squeeze(-1).squeeze(-1).contiguous();

            auto output = torch::empty_like(input).contiguous();
            
            int n_elements = input.numel();
            const int blocks = (n_elements + BLOCK_SIZE - 1) / BLOCK_SIZE;

            sn_normalize_kernel<<<blocks, BLOCK_SIZE>>>(
                input.data_ptr<float>(),
                ln_mean.data_ptr<float>(), ln_var.data_ptr<float>(),
                in_mean.data_ptr<float>(), in_var.data_ptr<float>(),
                running_mean.data_ptr<float>(), running_var.data_ptr<float>(),
                weight.data_ptr<float>(), bias.data_ptr<float>(),
                w_in.data_ptr<float>(), w_ln.data_ptr<float>(), w_bn.data_ptr<float>(),
                output.data_ptr<float>(),
                N, C, H, W, eps
            );
            
            return output;
        }}
        """
        
        self.sn_op = load_inline(
            name="sn_fused_op_logic_correct_v3",
            cpp_sources=cpp_source, 
            cuda_sources=cuda_source,
            functions=["sn_forward_cuda"],
            extra_cuda_cflags=["-O3", "--use_fast_math"], 
            verbose=True
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:

        weight = self.weight.squeeze().contiguous()
        bias = self.bias.squeeze().contiguous()
        w_in = self.w_in.contiguous()
        w_ln = self.w_ln.contiguous()
        w_bn = self.w_bn.contiguous()
        running_mean = self.running_mean.contiguous()
        running_var = self.running_var.contiguous()

        N, C, H, W = x.size(0), x.size(1), x.size(2), x.size(3)

        return self.sn_op.sn_forward_cuda(
            x.contiguous(), weight, bias, w_in, w_ln, w_bn, 
            running_mean, running_var, self.eps, N, C, H, W
        )