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 itakura_saito_kernel(
    const float* __restrict__ input,
    const float* __restrict__ target,
    float* __restrict__ output,
    int n,
    float eps
) {
    extern __shared__ float sdata[];
    int tid = threadIdx.x;
    int idx = blockIdx.x * blockDim.x + threadIdx.x;

    float local_sum = 0.0f;
    for (int i = idx; i < n; i += blockDim.x * gridDim.x) {
        float x = input[i];
        float y = target[i];

        float ratio = y / (x + eps);
        float term = ratio - logf(ratio + eps) - 1.0f;

        local_sum += term;
    }

    sdata[tid] = local_sum;
    __syncthreads();

    for (int s = blockDim.x / 2; s > 0; s >>= 1) {
        if (tid < s) {
            sdata[tid] += sdata[tid + s];
        }
        __syncthreads();
    }

    if (tid == 0) {
        atomicAdd(output, sdata[0]);
    }
}

torch::Tensor itakura_saito_cuda(torch::Tensor input, torch::Tensor target, float eps) {
    int n = input.numel();
    auto output = at::zeros({1}, input.options());

    int threads = 256;
    int blocks = 128;
    int shared_mem = threads * sizeof(float);

    itakura_saito_kernel<<<blocks, threads, shared_mem>>>(
        input.data_ptr<float>(),
        target.data_ptr<float>(),
        output.data_ptr<float>(),
        n,
        eps
    );

    return output / n;
}
"""

cpp_source = """
torch::Tensor itakura_saito_cuda(torch::Tensor input, torch::Tensor target, float eps);
"""

itakura_saito_loss = load_inline(
    name="itakura_saito_loss",
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=["itakura_saito_cuda"],
    verbose=False
)


class ModelNew(nn.Module):
    def __init__(self, eps=1e-8):
        super(ModelNew, self).__init__()
        self.eps = eps

    def forward(self, input, target):
        return itakura_saito_loss.itakura_saito_cuda(input, target, self.eps)