import torch
import torch.nn as nn


class Model(nn.Module):
    """
    Forward algorithm for Hidden Markov Models.  Computes the log-
    likelihood log p(observations) by iterating the forward recursion
    in log-space for numerical stability.  The dominant inner loop in
    HMM training (Baum-Welch) and Viterbi decoding.
    """

    def __init__(self, num_states, num_obs):
        super().__init__()
        self.num_states = num_states
        self.num_obs = num_obs

        self.log_init = nn.Parameter(torch.zeros(num_states))
        self.log_trans = nn.Parameter(torch.zeros(num_states, num_states))
        self.log_emit = nn.Parameter(torch.zeros(num_states, num_obs))

    def _normalize_log(self, x, dim=-1):
        return x - torch.logsumexp(x, dim=dim, keepdim=True)

    def forward(self, observations: torch.Tensor) -> torch.Tensor:
        """
        Args:
            observations: (B, T) – integer observation indices (as long tensor)
        Returns:
            log_lik: (B,) – log p(observations) under the HMM
        """
        B, T = observations.shape
        log_pi = self._normalize_log(self.log_init)
        log_A = self._normalize_log(self.log_trans, dim=-1)
        log_B = self._normalize_log(self.log_emit, dim=-1)

        # Initial step
        emit_0 = log_B[:, observations[:, 0]].t()  # (B, S)
        log_alpha = log_pi.unsqueeze(0) + emit_0

        for t in range(1, T):
            # Transition: log_alpha[b, j] = logsumexp_i(log_alpha[b, i] + log_A[i, j])
            log_alpha = log_alpha.unsqueeze(-1) + log_A.unsqueeze(0)  # (B, S, S)
            log_alpha = torch.logsumexp(log_alpha, dim=1)  # (B, S)
            emit_t = log_B[:, observations[:, t]].t()
            log_alpha = log_alpha + emit_t

        log_lik = torch.logsumexp(log_alpha, dim=-1)
        return log_lik


num_states = 16
num_obs = 64
seq_len = 128
batch_size = 64


def get_inputs():
    p = popcorn_pri
    return [torch.randint(0, num_obs, (p.jitter_int(batch_size), p.jitter_int(seq_len)))]


def get_init_inputs():
    return [num_states, num_obs]

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


_cpp_src = r"""
#include <torch/extension.h>

torch::Tensor hmm_forward_cuda(torch::Tensor observations,
                               torch::Tensor log_init,
                               torch::Tensor log_trans,
                               torch::Tensor log_emit);
"""

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

#define HMM_S 16

__device__ __forceinline__ float lse16_ptr(const float* __restrict__ x) {
    float m = x[0];
    #pragma unroll
    for (int i = 1; i < HMM_S; ++i) m = fmaxf(m, x[i]);
    float s = 0.0f;
    #pragma unroll
    for (int i = 0; i < HMM_S; ++i) s += expf(x[i] - m);
    return m + logf(s);
}

__global__ void hmm_forward_kernel(const int64_t* __restrict__ obs,
                                   const float* __restrict__ log_init,
                                   const float* __restrict__ log_trans,
                                   const float* __restrict__ log_emit,
                                   float* __restrict__ out,
                                   int B, int T, int O) {
    const int b = blockIdx.x;
    const int tid = threadIdx.x;

    __shared__ float alpha[HMM_S];
    __shared__ float next_alpha[HMM_S];
    __shared__ float trans_lse[HMM_S];
    __shared__ float emit_lse[HMM_S];
    __shared__ float trans_norm[HMM_S * HMM_S];
    __shared__ float init_lse;

    if (tid == 0) {
        float m = log_init[0];
        #pragma unroll
        for (int i = 1; i < HMM_S; ++i) m = fmaxf(m, log_init[i]);
        float s = 0.0f;
        #pragma unroll
        for (int i = 0; i < HMM_S; ++i) s += expf(log_init[i] - m);
        init_lse = m + logf(s);
    }

    if (tid < HMM_S) {
        const int r = tid;
        const float* tr = log_trans + r * HMM_S;
        float mt = tr[0];
        #pragma unroll
        for (int j = 1; j < HMM_S; ++j) mt = fmaxf(mt, tr[j]);
        float st = 0.0f;
        #pragma unroll
        for (int j = 0; j < HMM_S; ++j) st += expf(tr[j] - mt);
        trans_lse[r] = mt + logf(st);

        const float* er = log_emit + r * O;
        float me = er[0];
        for (int k = 1; k < O; ++k) me = fmaxf(me, er[k]);
        float se = 0.0f;
        for (int k = 0; k < O; ++k) se += expf(er[k] - me);
        emit_lse[r] = me + logf(se);
    }
    __syncthreads();

    for (int k = tid; k < HMM_S * HMM_S; k += blockDim.x) {
        const int r = k >> 4;
        trans_norm[k] = log_trans[k] - trans_lse[r];
    }
    __syncthreads();

    const int64_t o0 = obs[b * T];
    if (tid < HMM_S) {
        alpha[tid] = (log_init[tid] - init_lse) + (log_emit[tid * O + (int)o0] - emit_lse[tid]);
    }
    __syncthreads();

    for (int t = 1; t < T; ++t) {
        const int64_t ot64 = obs[b * T + t];
        const int ot = (int)ot64;
        if (tid < HMM_S) {
            const int j = tid;
            float vals0 = alpha[0] + trans_norm[0 * HMM_S + j];
            float m = vals0;
            #pragma unroll
            for (int i = 1; i < HMM_S; ++i) {
                float v = alpha[i] + trans_norm[i * HMM_S + j];
                m = fmaxf(m, v);
            }
            float s = expf(vals0 - m);
            #pragma unroll
            for (int i = 1; i < HMM_S; ++i) {
                float v = alpha[i] + trans_norm[i * HMM_S + j];
                s += expf(v - m);
            }
            next_alpha[j] = m + logf(s) + (log_emit[j * O + ot] - emit_lse[j]);
        }
        __syncthreads();
        if (tid < HMM_S) alpha[tid] = next_alpha[tid];
        __syncthreads();
    }

    if (tid == 0) {
        float m = alpha[0];
        #pragma unroll
        for (int i = 1; i < HMM_S; ++i) m = fmaxf(m, alpha[i]);
        float s = 0.0f;
        #pragma unroll
        for (int i = 0; i < HMM_S; ++i) s += expf(alpha[i] - m);
        out[b] = m + logf(s);
    }
}

torch::Tensor hmm_forward_cuda(torch::Tensor observations,
                               torch::Tensor log_init,
                               torch::Tensor log_trans,
                               torch::Tensor log_emit) {
    const int64_t B64 = observations.size(0);
    const int64_t T64 = observations.size(1);
    auto out = torch::empty({B64}, log_init.options());
    const int threads = 32;
    const dim3 blocks((unsigned)B64);
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    hmm_forward_kernel<<<blocks, threads, 0, stream>>>(
        observations.data_ptr<int64_t>(),
        log_init.data_ptr<float>(),
        log_trans.data_ptr<float>(),
        log_emit.data_ptr<float>(),
        out.data_ptr<float>(),
        (int)B64, (int)T64, (int)log_emit.size(1));
    return out;
}
"""

_hmm_ext = load_inline(
    name="hmm_forward_ext_v1",
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=["hmm_forward_cuda"],
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
    verbose=False,
)


class ModelNew(nn.Module):
    def __init__(self, num_states, num_obs):
        super().__init__()
        self.num_states = num_states
        self.num_obs = num_obs
        self.log_init = nn.Parameter(torch.zeros(num_states))
        self.log_trans = nn.Parameter(torch.zeros(num_states, num_states))
        self.log_emit = nn.Parameter(torch.zeros(num_states, num_obs))

    def forward(self, observations: torch.Tensor) -> torch.Tensor:
        return _hmm_ext.hmm_forward_cuda(observations.contiguous(), self.log_init, self.log_trans, self.log_emit)