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 <cfloat>

__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 long* __restrict__ obs,
    const float* __restrict__ log_pi,
    const float* __restrict__ log_A,
    const float* __restrict__ log_B,
    float* __restrict__ log_lik,
    int B, int T, int S, int O
) {
    int b = blockIdx.x * blockDim.x + threadIdx.x;
    if (b >= B) return;

    extern __shared__ float s_log_A[];
    // Load log_A into shared? but for simplicity, use registers since S=16 small

    float alpha[16];  // assume S <= 16

    // Initial
    int o0 = obs[b * T];
    for (int j = 0; j < S; j++) {
        alpha[j] = log_pi[j] + log_B[j * O + o0];
    }

    for (int t = 1; t < T; t++) {
        float new_alpha[16];
        int ot = obs[b * T + t];
        for (int j = 0; j < S; j++) {
            float mx = -FLT_MAX;
            for (int i = 0; i < S; i++) {
                float val = alpha[i] + log_A[i * S + j];
                if (val > mx) mx = val;
            }
            float sum = 0.0f;
            for (int i = 0; i < S; i++) {
                float val = alpha[i] + log_A[i * S + j];
                sum += expf(val - mx);
            }
            new_alpha[j] = mx + logf(sum) + log_B[j * O + ot];
        }
        for (int j = 0; j < S; j++) alpha[j] = new_alpha[j];
    }

    float mx = -FLT_MAX;
    for (int i = 0; i < S; i++) if (alpha[i] > mx) mx = alpha[i];
    float sum = 0.0f;
    for (int i = 0; i < S; i++) sum += expf(alpha[i] - mx);
    log_lik[b] = mx + logf(sum);
}

torch::Tensor hmm_forward_cuda(
    torch::Tensor observations,
    torch::Tensor log_init,
    torch::Tensor log_trans,
    torch::Tensor log_emit
) {
    auto B = observations.size(0);
    auto T = observations.size(1);
    auto S = log_init.size(0);
    auto O = log_emit.size(1);

    auto log_pi = log_init - torch::logsumexp(log_init, -1, true);
    auto log_A = log_trans - torch::logsumexp(log_trans, -1, true);
    auto log_B = log_emit - torch::logsumexp(log_emit, -1, true);

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

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

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

    return log_lik;
}
'''

cpp_source = '''
torch::Tensor hmm_forward_cuda(
    torch::Tensor observations,
    torch::Tensor log_init,
    torch::Tensor log_trans,
    torch::Tensor log_emit
);
'''

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))

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