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_src = r'''
#include <torch/extension.h>
torch::Tensor dm_forward(torch::Tensor counts);
void dm_init(int64_t V);
'''

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

#define MAX_N 2304
__constant__ float PRIOR_TABLE[MAX_N + 1];

__device__ __forceinline__ float logfact_int_0_9(float x) {
    int n = __float2int_rz(x);
    if (n <= 1) return 0.0f;
    if (n == 2) return 0.6931471805599453f;
    if (n == 3) return 1.7917594692280550f;
    if (n == 4) return 3.1780538303479458f;
    if (n == 5) return 4.7874917427820458f;
    if (n == 6) return 6.5792512120101012f;
    if (n == 7) return 8.5251613610654143f;
    if (n == 8) return 10.604602902745250f;
    return 12.801827480081469f;
}

__global__ void dm_alpha1_table_kernel(const float* __restrict__ counts,
                                       float* __restrict__ out,
                                       int B, int V) {
    int b = blockIdx.x;
    int tid = threadIdx.x;
    int lane = tid & 31;
    int warp = tid >> 5;
    float nsum = 0.0f;
    float lsum = 0.0f;

    if (tid < V) {
        float c = counts[(size_t)b * V + tid];
        nsum = c;
        lsum = logfact_int_0_9(c);
    }

    #pragma unroll
    for (int off = 16; off > 0; off >>= 1) {
        nsum += __shfl_down_sync(0xffffffff, nsum, off);
        lsum += __shfl_down_sync(0xffffffff, lsum, off);
    }

    __shared__ float sN[8];
    __shared__ float sL[8];
    if (lane == 0) {
        sN[warp] = nsum;
        sL[warp] = lsum;
    }
    __syncthreads();

    if (warp == 0) {
        nsum = (lane < 8) ? sN[lane] : 0.0f;
        lsum = (lane < 8) ? sL[lane] : 0.0f;
        #pragma unroll
        for (int off = 16; off > 0; off >>= 1) {
            nsum += __shfl_down_sync(0xffffffff, nsum, off);
            lsum += __shfl_down_sync(0xffffffff, lsum, off);
        }
        if (lane == 0) {
            out[b] = PRIOR_TABLE[__float2int_rn(nsum)] + lsum;
        }
    }
}

static bool tables_ready = false;

void dm_init(int64_t V64) {
    if (!tables_ready) {
        int V = (int)V64;
        float h[MAX_N + 1];
        float lgV = lgammaf((float)V);
        for (int i = 0; i <= MAX_N; ++i) {
            h[i] = lgV - lgammaf((float)(V + i));
        }
        cudaMemcpyToSymbol(PRIOR_TABLE, h, sizeof(h));
        tables_ready = true;
    }
}

torch::Tensor dm_forward(torch::Tensor counts) {
    const int B = (int)counts.size(0);
    const int V = (int)counts.size(1);
    auto out = torch::empty({B}, counts.options());
    if (!tables_ready) dm_init((int64_t)V);
    dm_alpha1_table_kernel<<<B, 256, 0, at::cuda::getCurrentCUDAStream()>>>(counts.data_ptr<float>(), out.data_ptr<float>(), B, V);
    return out;
}
'''

_ext = load_inline(
    name='dm_ll_ext_final',
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=['dm_forward', 'dm_init'],
    extra_cuda_cflags=['-O3'],
    extra_cflags=['-O3'],
    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))
        _ext.dm_init(vocab_size)

    def forward(self, counts: torch.Tensor) -> torch.Tensor:
        return _ext.dm_forward(counts)