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

cpp_sources = r'''
#include <torch/extension.h>
torch::Tensor dm_forward_cuda(torch::Tensor counts, torch::Tensor log_alpha);
torch::Tensor dm_forward_unit_cuda(torch::Tensor counts);
torch::Tensor dm_forward(torch::Tensor counts, torch::Tensor log_alpha) { return dm_forward_cuda(counts, log_alpha); }
torch::Tensor dm_forward_unit(torch::Tensor counts) { return dm_forward_unit_cuda(counts); }
'''

cuda_sources = r'''
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <math.h>

__device__ __forceinline__ float warp_sum(float v) {
    unsigned mask = 0xffffffffu;
    v += __shfl_down_sync(mask, v, 16);
    v += __shfl_down_sync(mask, v, 8);
    v += __shfl_down_sync(mask, v, 4);
    v += __shfl_down_sync(mask, v, 2);
    v += __shfl_down_sync(mask, v, 1);
    return v;
}

__device__ __forceinline__ float logfac_int(int c) {
    switch (c) {
        case 0: return 0.0f;
        case 1: return 0.0f;
        case 2: return 0.6931471805599453f;
        case 3: return 1.791759469228055f;
        case 4: return 3.1780538303479458f;
        case 5: return 4.787491742782046f;
        case 6: return 6.579251212010101f;
        case 7: return 8.525161361065415f;
        case 8: return 10.60460290274525f;
        case 9: return 12.801827480081469f;
        case 10: return 15.104412573075516f;
        case 11: return 17.502307845873887f;
        case 12: return 19.987214495661885f;
        case 13: return 22.552163853123425f;
        case 14: return 25.191221182738683f;
        case 15: return 27.899271383840894f;
        case 16: return 30.671860106080675f;
        case 17: return 33.50507345013689f;
        case 18: return 36.39544520803305f;
        case 19: return 39.339884187199495f;
        case 20: return 42.335616460753485f;
        default: return lgammaf((float)c + 1.0f);
    }
}

__global__ void dm_unit_kernel(const float* __restrict__ counts, float* __restrict__ out, int B, int V) {
    int row = blockIdx.x;
    int tid = threadIdx.x;
    int lane = tid & 31;
    int wid = tid >> 5;
    __shared__ float sn[32];
    __shared__ float st[32];
    float n = 0.0f;
    float t = 0.0f;
    if (tid < V) {
        float cf = counts[row * V + tid];
        n = cf;
        t = logfac_int((int)(cf + 0.5f));
    }
    n = warp_sum(n);
    t = warp_sum(t);
    if (lane == 0) { sn[wid] = n; st[wid] = t; }
    __syncthreads();
    if (wid == 0) {
        int nw = (blockDim.x + 31) >> 5;
        n = (lane < nw) ? sn[lane] : 0.0f;
        t = (lane < nw) ? st[lane] : 0.0f;
        n = warp_sum(n);
        t = warp_sum(t);
        if (lane == 0) {
            float asum = (float)V;
            out[row] = lgammaf(asum) - lgammaf(n + asum) + t;
        }
    }
}

__global__ void prepare_kernel(const float* __restrict__ log_alpha, float* __restrict__ alpha, float* __restrict__ alpha_sum, int V) {
    int tid = threadIdx.x;
    int lane = tid & 31;
    int wid = tid >> 5;
    __shared__ float ss[32];
    float val = 0.0f;
    if (tid < V) {
        val = expf(log_alpha[tid]);
        alpha[tid] = val;
    }
    val = warp_sum(val);
    if (lane == 0) ss[wid] = val;
    __syncthreads();
    if (wid == 0) {
        int nw = (blockDim.x + 31) >> 5;
        val = (lane < nw) ? ss[lane] : 0.0f;
        val = warp_sum(val);
        if (lane == 0) alpha_sum[0] = val;
    }
}

__global__ void dm_generic_kernel(const float* __restrict__ counts,
                          const float* __restrict__ alpha,
                          const float* __restrict__ alpha_sum,
                          float* __restrict__ out,
                          int B, int V) {
    int row = blockIdx.x;
    int tid = threadIdx.x;
    int lane = tid & 31;
    int wid = tid >> 5;
    __shared__ float sn[32];
    __shared__ float st[32];
    float n = 0.0f;
    float term = 0.0f;
    if (tid < V) {
        float cf = counts[row * V + tid];
        n = cf;
        int c = (int)(cf + 0.5f);
        float a = alpha[tid];
        for (int k = 0; k < c; ++k) term += logf(a + (float)k);
    }
    n = warp_sum(n);
    term = warp_sum(term);
    if (lane == 0) { sn[wid] = n; st[wid] = term; }
    __syncthreads();
    if (wid == 0) {
        int nw = (blockDim.x + 31) >> 5;
        n = (lane < nw) ? sn[lane] : 0.0f;
        term = (lane < nw) ? st[lane] : 0.0f;
        n = warp_sum(n);
        term = warp_sum(term);
        if (lane == 0) {
            float asum = alpha_sum[0];
            out[row] = lgammaf(asum) - lgammaf(n + asum) + term;
        }
    }
}

torch::Tensor dm_forward_unit_cuda(torch::Tensor counts) {
    const int B = counts.size(0);
    const int V = counts.size(1);
    auto out = torch::empty({B}, counts.options());
    int threads = 1; while (threads < V) threads <<= 1; if (threads < 32) threads = 32; if (threads > 1024) threads = 1024;
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    dm_unit_kernel<<<B, threads, 0, stream>>>(counts.data_ptr<float>(), out.data_ptr<float>(), B, V);
    return out;
}

torch::Tensor dm_forward_cuda(torch::Tensor counts, torch::Tensor log_alpha) {
    const int B = counts.size(0);
    const int V = counts.size(1);
    auto out = torch::empty({B}, counts.options());
    auto alpha = torch::empty({V}, counts.options());
    auto alpha_sum = torch::empty({1}, counts.options());
    int threads = 1; while (threads < V) threads <<= 1; if (threads < 32) threads = 32; if (threads > 1024) threads = 1024;
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    prepare_kernel<<<1, threads, 0, stream>>>(log_alpha.data_ptr<float>(), alpha.data_ptr<float>(), alpha_sum.data_ptr<float>(), V);
    dm_generic_kernel<<<B, threads, 0, stream>>>(counts.data_ptr<float>(), alpha.data_ptr<float>(), alpha_sum.data_ptr<float>(), out.data_ptr<float>(), B, V);
    return out;
}
'''

_ext = load_inline(name='dm_poly_ext_v3', cpp_sources=cpp_sources, cuda_sources=cuda_sources, functions=['dm_forward','dm_forward_unit'], extra_cflags=['-O3'], extra_cuda_cflags=['-O3', '--use_fast_math'], 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:
        counts = counts.contiguous()
        if self.log_alpha._version == 0:
            return _ext.dm_forward_unit(counts)
        return _ext.dm_forward(counts, self.log_alpha)