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_log_lik(
    const float* __restrict__ counts,
    const float* __restrict__ alpha,
    float* __restrict__ log_lik,
    int B,
    int V,
    float alpha_sum
) {
    int b = blockIdx.x * blockDim.x + threadIdx.x;
    if (b >= B) return;

    double sum = 0.0;
    double N = 0.0;
    for (int v = 0; v < V; ++v) {
        float count = counts[b * V + v];
        float a = alpha[v];
        sum += lgamma(static_cast<double>(count + a)) - lgamma(static_cast<double>(a));
        N += count;
    }
    log_lik[b] = static_cast<float>(lgamma(static_cast<double>(alpha_sum)) - lgamma(static_cast<double>(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);
    auto alpha_sum_t = alpha.sum();
    float alpha_sum = alpha_sum_t.item<float>();
    auto log_lik = torch::empty({B}, counts.options());
    
    dim3 block(256);
    dim3 grid((B + 255) / 256);
    
    dirichlet_multinomial_log_lik<<<grid, block>>>(
        counts.data_ptr<float>(),
        alpha.data_ptr<float>(),
        log_lik.data_ptr<float>(),
        B, V, alpha_sum
    );
    return log_lik;
}
"""

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

module = load_inline(
    name='dirichlet_multinomial',
    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))
        self.module = module

    def forward(self, counts: torch.Tensor) -> torch.Tensor:
        return self.module.dirichlet_multinomial_forward(counts, self.log_alpha)