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

# Compile CUDA extension at module load time
hmm_cuda = load_inline(
    name='hmm_forward_cuda',
    cpp_sources='''
extern "C" void hmm_forward_cuda_impl(
    const float* log_pi,
    const float* log_A,
    const float* log_B,
    const long* observations,
    float* log_lik_out,
    int batch_size, int seq_len, int num_states);
''',
    cuda_sources='''
#include <cuda_fp16.h>
#include <cfloat>

__device__ float logsumexp_device(const float* arr, int n) {
    float max_val = -FLT_MAX;
    for (int i = 0; i < n; i++) {
        max_val = max(max_val, arr[i]);
    }
    if (max_val == -FLT_MAX) return -FLT_MAX;
    float sum_exp = 0.0f;
    for (int i = 0; i < n; i++) {
        sum_exp += expf(arr[i] - max_val);
    }
    return max_val + logf(sum_exp);
}

extern "C" __global__ void hmm_forward_kernel(
    const float* __restrict__ log_pi,
    const float* __restrict__ log_A,
    const float* __restrict__ log_B,
    const long* __restrict__ observations,
    float* __restrict__ log_alpha,
    int batch_size, int seq_len, int num_states) {
    
    // Each thread handles one (batch, state) pair
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int total = batch_size * num_states;
    
    if (idx >= total) return;
    
    int b = idx / num_states;
    int j = idx % num_states;
    
    // Time 0: log_alpha[b, j] = log_pi[j] + log_B[j, obs[b, 0]]
    long obs0 = observations[b * seq_len + 0];
    log_alpha[idx] = log_pi[j] + log_B[j * num_states + (int)obs0];
    
    __syncthreads();
    
    // Iterate through time steps
    for (int t = 1; t < seq_len; t++) {
        // Compute logsumexp_i(log_alpha[b, i] + log_A[i, j])
        float buffer[64]; // max num_states = 64
        for (int i = 0; i < num_states; i++) {
            buffer[i] = log_alpha[b * num_states + i] + log_A[i * num_states + j];
        }
        float logsumexp = logsumexp_device(buffer, num_states);
        
        // Add emission
        long obs_t = observations[b * seq_len + t];
        log_alpha[idx] = logsumexp + log_B[j * num_states + (int)obs_t];
        
        __syncthreads();
    }
}

extern "C" __global__ void reduce_logsumexp_kernel(
    const float* __restrict__ log_alpha,
    float* __restrict__ log_lik_out,
    int batch_size, int num_states) {
    
    // Each thread handles one batch element
    int b = blockIdx.x * blockDim.x + threadIdx.x;
    if (b >= batch_size) return;
    
    float max_val = -FLT_MAX;
    for (int j = 0; j < num_states; j++) {
        max_val = max(max_val, log_alpha[b * num_states + j]);
    }
    
    if (max_val == -FLT_MAX) {
        log_lik_out[b] = -FLT_MAX;
        return;
    }
    
    float sum_exp = 0.0f;
    for (int j = 0; j < num_states; j++) {
        sum_exp += expf(log_alpha[b * num_states + j] - max_val);
    }
    log_lik_out[b] = max_val + logf(sum_exp);
}

void hmm_forward_cuda_impl(
    const float* log_pi,
    const float* log_A,
    const float* log_B,
    const long* observations,
    float* log_lik_out,
    int batch_size, int seq_len, int num_states) {
    
    // Allocate temporary buffer for log_alpha
    float* log_alpha;
    cudaMalloc(&log_alpha, batch_size * num_states * sizeof(float));
    
    // Launch forward kernel
    int total = batch_size * num_states;
    int blockSize = 256;
    int numBlocks = (total + blockSize - 1) / blockSize;
    
    hmm_forward_kernel<<<numBlocks, blockSize>>>(
        log_pi, log_A, log_B, observations, log_alpha,
        batch_size, seq_len, num_states);
    
    // Launch reduction kernel
    int reduceBlockSize = 256;
    int reduceNumBlocks = (batch_size + reduceBlockSize - 1) / reduceBlockSize;
    
    reduce_logsumexp_kernel<<<reduceNumBlocks, reduceBlockSize>>>(
        log_alpha, log_lik_out, batch_size, num_states);
    
    cudaFree(log_alpha);
}
''',
    functions=['hmm_forward_cuda_impl'],
    extra_cuda_cflags=['-O3', '-use_fast_math'],
    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 _normalize_log(self, x, dim=-1):
        return x - torch.logsumexp(x, dim=dim, keepdim=True)
    
    def forward(self, observations: torch.Tensor) -> torch.Tensor:
        B, T = observations.shape
        
        # Normalize parameters (done in PyTorch - already optimized)
        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)
        
        # Ensure contiguous and on same device
        log_pi = log_pi.contiguous()
        log_A = log_A.contiguous()
        log_B = log_B.contiguous()
        observations = observations.contiguous()
        
        # Allocate output
        log_lik = torch.empty(B, dtype=torch.float32, device=observations.device)
        
        # Call CUDA kernel
        hmm_cuda.hmm_forward_cuda_impl(
            log_pi.data_ptr(),
            log_A.data_ptr(),
            log_B.data_ptr(),
            observations.data_ptr(),
            log_lik.data_ptr(),
            B, T, self.num_states
        )
        
        return log_lik