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

cuda_source = '''
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <math.h>

__device__ __forceinline__ float logsumexp2(float a, float b) {
    float max_val = fmaxf(a, b);
    return max_val + logf(expf(a - max_val) + expf(b - max_val));
}

__global__ void hmm_forward_kernel(
    const float* __restrict__ log_pi,
    const float* __restrict__ log_A,
    const float* __restrict__ log_B,
    const int64_t* __restrict__ observations,
    float* __restrict__ log_lik,
    int B, int T, int S, int O
) {
    int b = blockIdx.x * blockDim.x + threadIdx.x;
    if (b >= B) return;

    float alpha_reg[16];

    // Initial
    int64_t obs0 = observations[b * T];
    for (int s = 0; s < S; s++) {
        alpha_reg[s] = log_pi[s] + log_B[s * O + obs0];
    }

    for (int t = 1; t < T; t++) {
        int64_t obs = observations[b * T + t];
        float new_alpha[16];
        for (int j = 0; j < S; j++) {
            float sum = -INFINITY;
            for (int i = 0; i < S; i++) {
                float val = alpha_reg[i] + log_A[i * S + j];
                sum = logsumexp2(sum, val);
            }
            new_alpha[j] = sum + log_B[j * O + obs];
        }
        for (int s = 0; s < S; s++) alpha_reg[s] = new_alpha[s];
    }

    float lik = -INFINITY;
    for (int s = 0; s < S; s++) {
        lik = logsumexp2(lik, alpha_reg[s]);
    }
    log_lik[b] = lik;
}

torch::Tensor hmm_forward_cuda(
    torch::Tensor log_pi,
    torch::Tensor log_A,
    torch::Tensor log_B,
    torch::Tensor observations
) {
    observations = observations.contiguous();

    int B = observations.size(0);
    int T = observations.size(1);
    int S = log_pi.size(0);
    int O = log_B.size(1);

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

    int threads = 64;
    int blocks = (B + threads - 1) / threads;

    hmm_forward_kernel<<<blocks, threads>>>(
        log_pi.data_ptr<float>(),
        log_A.data_ptr<float>(),
        log_B.data_ptr<float>(),
        observations.data_ptr<int64_t>(),
        log_lik.data_ptr<float>(),
        B, T, S, O
    );

    return log_lik;
}
'''

cpp_source = '''
torch::Tensor hmm_forward_cuda(torch::Tensor log_pi, torch::Tensor log_A, torch::Tensor log_B, torch::Tensor observations);
'''

module = load_inline(
    name='hmm_cuda',
    cpp_sources=[cpp_source],
    cuda_sources=[cuda_source],
    functions=['hmm_forward_cuda'],
    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))
        self.hmm_cuda = module

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

    def forward(self, observations: torch.Tensor) -> torch.Tensor:
        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)
        return self.hmm_cuda.hmm_forward_cuda(
            log_pi, log_A, log_B, observations
        )