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 = '''
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <cmath>

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

    float sum = 0.0f;
    const float* row = counts + b * V;
    for (int v = 0; v < V; ++v) {
        float c = row[v];
        float a = alpha[v];
        sum += lgammaf(c + a) - lgammaf(a);
    }

    float N = 0.0f;
    for (int v = 0; v < V; ++v) {
        N += row[v];
    }

    log_lik[b] = log_alpha_sum - lgammaf(N + alpha_sum) + sum;
}

torch::Tensor dirichlet_multinomial_forward(
    torch::Tensor counts,
    torch::Tensor log_alpha
) {
    int B = counts.size(0);
    int V = counts.size(1);

    auto alpha = torch::exp(log_alpha);
    float alpha_sum = alpha.sum().item<float>();
    float log_alpha_sum = lgammaf(alpha_sum);

    auto log_lik = torch::empty({B}, counts.options());

    dim3 block(256);
    dim3 grid((B + block.x - 1) / block.x);

    dirichlet_multinomial_kernel<<<grid, block>>>(
        counts.data_ptr<float>(),
        alpha.data_ptr<float>(),
        log_lik.data_ptr<float>(),
        alpha_sum,
        log_alpha_sum,
        B,
        V
    );

    return log_lik;
}
'''

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

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

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(counts, self.log_alpha)