pelorus 0.1.0

use super::{HmmLattice, LogProbability, StateId, ViterbiLattice, ViterbiPath};

/// Viterbi algorithm solver
#[derive(Debug)]
pub struct ViterbiSolver {
    lattice: ViterbiLattice,
}

impl ViterbiSolver {
    /// Create a new Viterbi solver
    pub fn new(lattice: ViterbiLattice) -> Self {
        Self { lattice }
    }

    /// Find the best path through the lattice
    pub fn solve(
        &self,
        emission_probs: &[Vec<f64>],
        transition_probs: &[Vec<Vec<f64>>],
    ) -> Option<ViterbiPath> {
        // Convert to HMM lattice for processing
        let states_per_step = (0..self.lattice.time_steps())
            .map(|t| self.lattice.states_at_time(t).unwrap_or(0))
            .collect();

        let hmm_lattice = HmmLattice::builder()
            .with_structure(self.lattice.time_steps(), states_per_step)
            .with_emission_probs(emission_probs.to_vec())
            .with_transition_probs(transition_probs.to_vec())
            .build()
            .ok()?;

        viterbi_path(&hmm_lattice)
    }
}

/// Find the best path through the HMM lattice using Viterbi algorithm
pub fn viterbi_path(lattice: &HmmLattice) -> Option<ViterbiPath> {
    let time_steps = lattice.time_steps();
    if time_steps == 0 {
        return None;
    }

    // Initialize dynamic programming tables
    // dp[time][state] = max log probability of reaching state at time
    let mut dp: Vec<Vec<LogProbability>> = Vec::new();
    // backpointer[time][state] = previous state that led to this state
    let mut backpointer: Vec<Vec<Option<StateId>>> = Vec::new();

    // Initialize first time step
    let first_states = lattice.states_at_time(0)?;
    let mut first_dp = vec![std::f64::NEG_INFINITY; first_states];
    let first_backpointer = vec![None; first_states];

    // Set initial probabilities (emission probabilities at time 0)
    for state in 0..first_states {
        if let Some(emission_log_prob) = lattice.emission_log_prob(0, state) {
            first_dp[state] = emission_log_prob;
        }
    }

    dp.push(first_dp);
    backpointer.push(first_backpointer);

    // Forward pass: compute maximum probabilities for each time step
    for time in 1..time_steps {
        let current_states = lattice.states_at_time(time)?;
        let prev_states = lattice.states_at_time(time - 1)?;

        let mut current_dp = vec![std::f64::NEG_INFINITY; current_states];
        let mut current_backpointer = vec![None; current_states];

        for to_state in 0..current_states {
            let emission_log_prob = lattice.emission_log_prob(time, to_state)?;

            // Find best previous state
            for from_state in 0..prev_states {
                let prev_prob = dp[time - 1][from_state];
                if prev_prob == std::f64::NEG_INFINITY {
                    continue; // Skip impossible previous states
                }

                let transition_log_prob =
                    lattice.transition_log_prob(time - 1, from_state, to_state)?;
                if transition_log_prob == std::f64::NEG_INFINITY {
                    continue; // Skip impossible transitions
                }

                let total_log_prob =
                    prev_prob + transition_log_prob + emission_log_prob;

                if total_log_prob > current_dp[to_state] {
                    current_dp[to_state] = total_log_prob;
                    current_backpointer[to_state] = Some(from_state);
                }
            }
        }

        dp.push(current_dp);
        backpointer.push(current_backpointer);
    }

    // Find the best final state
    let final_time = time_steps - 1;
    let final_states = lattice.states_at_time(final_time)?;

    let mut best_final_state = None;
    let mut best_final_prob = std::f64::NEG_INFINITY;

    for state in 0..final_states {
        if dp[final_time][state] > best_final_prob {
            best_final_prob = dp[final_time][state];
            best_final_state = Some(state);
        }
    }

    let best_final_state = best_final_state?;

    // Check if we found a valid path
    if best_final_prob == std::f64::NEG_INFINITY {
        return None;
    }

    // Backtrack to reconstruct the path
    let mut path = Vec::new();
    let mut current_state = best_final_state;

    path.push(current_state);

    for time in (1..time_steps).rev() {
        if let Some(prev_state) = backpointer[time][current_state] {
            path.push(prev_state);
            current_state = prev_state;
        } else {
            // No valid backpointer - this shouldn't happen if we found a valid path
            return None;
        }
    }

    // Reverse path since we built it backwards
    path.reverse();

    Some(ViterbiPath::new(path, best_final_prob))
}