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

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

__global__ void fisherrao_dist_kernel(const float* __restrict__ x, 
                                      const float* __restrict__ y, 
                                      float* __restrict__ out, 
                                      int size) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < size) {
        float eps = 1e-6f;
        float val_x = fabsf(x[idx]) + eps;
        float val_y = fabsf(y[idx]) + eps;
        // Metric: |log(x) - log(y)|
        out[idx] = fabsf(logf(val_x) - logf(val_y));
    }
}

torch::Tensor fisherrao_dist_cuda(torch::Tensor x, torch::Tensor y) {
    auto size = x.numel();
    auto out = torch::empty_like(x);

    const int block_size = 256;
    int grid_size = (size + block_size - 1) / block_size;

    fisherrao_dist_kernel<<<grid_size, block_size>>>(
        x.data_ptr<float>(),
        y.data_ptr<float>(),
        out.data_ptr<float>(),
        size
    );

    return out;
}
"""

cpp_source = "torch::Tensor fisherrao_dist_cuda(torch::Tensor x, torch::Tensor y);"

module = load_inline(
    name="fisherrao_rmsnorm_ext",
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=["fisherrao_dist_cuda"],
    verbose=False,
    with_cuda=True
)


class RMSNorm(nn.Module):
    def __init__(self, dim, eps=1e-6):
        super().__init__()
        self.scale = nn.Parameter(torch.ones(dim))
        self.eps = eps

    def forward(self, x):
        var = x.pow(2).mean(-1, keepdim=True)
        norm_x = x * torch.rsqrt(var + self.eps)
        return norm_x * self.scale


class ModelNew(nn.Module):
    def __init__(self, dim):
        super(ModelNew, self).__init__()
        self.rmsnorm = RMSNorm(dim)
        self.op = module

    def forward(self, x, y):
        dist = self.op.fisherrao_dist_cuda(x.contiguous(), y.contiguous())
        out = self.rmsnorm(dist)
        return out.mean()