import torch
import torch.nn as nn


class Model(nn.Module):
    """
    Log-likelihood under the Dirichlet-Multinomial (Polya) distribution.
    Given count vectors n and concentration parameters alpha, computes
    log p(n | alpha) using the log-gamma formulation.  Core computation
    in Bayesian topic models and mixed-membership models.
    """

    def __init__(self, vocab_size):
        super().__init__()
        self.vocab_size = vocab_size
        self.log_alpha = nn.Parameter(torch.zeros(vocab_size))

    def forward(self, counts: torch.Tensor) -> torch.Tensor:
        """
        Args:
            counts: (B, V) – integer count vectors (as float for differentiability)
        Returns:
            log_lik: (B,) – log p(counts | alpha) for each sample
        """
        alpha = torch.exp(self.log_alpha).unsqueeze(0)  # (1, V)
        N = counts.sum(-1)  # (B,)
        alpha_sum = alpha.sum(-1)  # (1,)

        # log p(n | alpha) = log Gamma(alpha_sum) - log Gamma(N + alpha_sum)
        #                   + sum_v [ log Gamma(n_v + alpha_v) - log Gamma(alpha_v) ]
        log_lik = (
            torch.lgamma(alpha_sum)
            - torch.lgamma(N + alpha_sum)
            + (torch.lgamma(counts + alpha) - torch.lgamma(alpha)).sum(-1)
        )
        return log_lik


vocab_size = 256
batch_size = 512


def get_inputs():
    bs = popcorn_pri.jitter_int(batch_size)
    counts = torch.randint(0, 10, (bs, vocab_size)).float()
    return [counts]


def get_init_inputs():
    return [vocab_size]

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

cuda_source = r'''
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <cmath>

__device__ float lgamma_approx(float x) {
    return lgammaf(x);
}

__global__ void dirichlet_multinomial_kernel(
    const float* __restrict__ log_alpha,
    const float* __restrict__ counts,
    float* __restrict__ log_lik,
    int B,
    int V,
    float alpha_sum,
    float log_gamma_alpha_sum
) {
    int b = blockIdx.x * blockDim.x + threadIdx.x;
    if (b >= B) return;

    float sum = 0.0f;
    float N = 0.0f;
    for (int v = 0; v < V; v++) {
        float alpha_v = expf(log_alpha[v]);
        float count = counts[b * V + v];
        N += count;
        sum += lgamma_approx(count + alpha_v) - lgamma_approx(alpha_v);
    }
    float term = log_gamma_alpha_sum - lgammaf(N + alpha_sum);
    log_lik[b] = sum + term;
}

torch::Tensor dirichlet_multinomial_forward(torch::Tensor log_alpha, torch::Tensor counts) {
    int V = log_alpha.size(0);
    int B = counts.size(0);
    
    auto alpha = torch::exp(log_alpha);
    float alpha_sum = alpha.sum().item<float>();
    float log_gamma_alpha_sum = lgammaf(alpha_sum);
    
    auto log_lik = torch::empty({B}, counts.options());
    
    int threads = 256;
    int blocks = (B + threads - 1) / threads;
    
    dirichlet_multinomial_kernel<<<blocks, threads>>>(
        log_alpha.data_ptr<float>(),
        counts.data_ptr<float>(),
        log_lik.data_ptr<float>(),
        B, V, alpha_sum, log_gamma_alpha_sum
    );
    
    return log_lik;
}
'''

cpp_source = r'''
torch::Tensor dirichlet_multinomial_forward(torch::Tensor log_alpha, torch::Tensor counts);
'''

module = load_inline(
    name='dirichlet_multinomial',
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=['dirichlet_multinomial_forward'],
    verbose=True
)

class ModelNew(nn.Module):
    def __init__(self, vocab_size):
        super().__init__()
        self.vocab_size = vocab_size
        self.log_alpha = nn.Parameter(torch.zeros(vocab_size))
    
    def forward(self, counts: torch.Tensor) -> torch.Tensor:
        return module.dirichlet_multinomial_forward(self.log_alpha, counts)