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

__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;
    if (b >= B) return;

    extern __shared__ float alpha[];
    float* prev_alpha = alpha;
    float* curr_alpha = alpha + S;

    const int64_t* obs = observations + b * T;

    // Initial step
    for (int s = threadIdx.x; s < S; s += blockDim.x) {
        prev_alpha[s] = log_pi[s] + log_B[s * O + obs[0]];
    }
    __syncthreads();

    for (int t = 1; t < T; t++) {
        for (int j = threadIdx.x; j < S; j += blockDim.x) {
            float max_val = -INFINITY;
            for (int i = 0; i < S; i++) {
                float val = prev_alpha[i] + log_A[i * S + j];
                if (val > max_val) max_val = val;
            }
            float sum = 0.0f;
            for (int i = 0; i < S; i++) {
                float val = prev_alpha[i] + log_A[i * S + j];
                sum += expf(val - max_val);
            }
            curr_alpha[j] = max_val + logf(sum) + log_B[j * O + obs[t]];
        }
        __syncthreads();
        float* temp = prev_alpha;
        prev_alpha = curr_alpha;
        curr_alpha = temp;
    }

    // logsumexp over states
    if (threadIdx.x == 0) {
        float max_val = -INFINITY;
        for (int s = 0; s < S; s++) {
            if (prev_alpha[s] > max_val) max_val = prev_alpha[s];
        }
        float sum = 0.0f;
        for (int s = 0; s < S; s++) {
            sum += expf(prev_alpha[s] - max_val);
        }
        log_lik[b] = max_val + logf(sum);
    }
}

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

    // Normalize
    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);

    log_pi = log_pi.contiguous();
    log_A = log_A.contiguous();
    log_B = log_B.contiguous();
    auto obs_long = observations.contiguous().to(torch::kInt64);

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

    int threads = 16;
    int shared_mem = 2 * S * sizeof(float);

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

    return log_lik;
}
'''

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

module = load_inline(
    name='hmm_cuda',
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=['hmm_forward'],
    verbose=False,
    extra_cuda_cflags=['-O3']
)

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 forward(self, observations: torch.Tensor) -> torch.Tensor:
        return self.hmm_cuda.hmm_forward(
            self.log_init, self.log_trans, self.log_emit, observations
        )
```