pelorus 0.1.0

use super::{LogProbability, StateId, TimeStep};

/// A Viterbi lattice for dynamic programming
#[derive(Debug, Clone)]
pub struct ViterbiLattice {
    /// Number of time steps
    time_steps: usize,
    /// Number of states per time step
    states_per_step: Vec<usize>,
}

impl ViterbiLattice {
    /// Create a new Viterbi lattice
    pub fn new(time_steps: usize, states_per_step: Vec<usize>) -> Self {
        Self {
            time_steps,
            states_per_step,
        }
    }

    /// Get the number of time steps
    pub fn time_steps(&self) -> usize {
        self.time_steps
    }

    /// Get the number of states at a given time step
    pub fn states_at_time(&self, time: TimeStep) -> Option<usize> {
        self.states_per_step.get(time).copied()
    }
}

/// HMM lattice with emission and transition probabilities
#[derive(Debug, Clone)]
pub struct HmmLattice {
    /// Number of time steps
    time_steps: usize,
    /// Number of states per time step
    states_per_step: Vec<usize>,
    /// Emission probabilities: emission_probs[time][state] = log(P(observation | state))
    emission_probs: Vec<Vec<LogProbability>>,
    /// Transition probabilities: transition_probs[time][from_state][to_state] = log(P(to_state | from_state))
    transition_probs: Vec<Vec<Vec<LogProbability>>>,
}

impl HmmLattice {
    /// Create a builder for constructing an HMM lattice
    pub fn builder() -> HmmLatticeBuilder {
        HmmLatticeBuilder::new()
    }

    /// Get the number of time steps
    pub fn time_steps(&self) -> usize {
        self.time_steps
    }

    /// Get the number of states at a given time step
    pub fn states_at_time(&self, time: TimeStep) -> Option<usize> {
        self.states_per_step.get(time).copied()
    }

    /// Get emission log probability for a state at a time step
    pub fn emission_log_prob(
        &self,
        time: TimeStep,
        state: StateId,
    ) -> Option<LogProbability> {
        self.emission_probs.get(time)?.get(state).copied()
    }

    /// Get transition log probability between states
    pub fn transition_log_prob(
        &self,
        time: TimeStep,
        from_state: StateId,
        to_state: StateId,
    ) -> Option<LogProbability> {
        self.transition_probs
            .get(time)?
            .get(from_state)?
            .get(to_state)
            .copied()
    }
}

/// Builder for HMM lattice
#[derive(Debug)]
pub struct HmmLatticeBuilder {
    time_steps: Option<usize>,
    states_per_step: Option<Vec<usize>>,
    emission_probs: Option<Vec<Vec<f64>>>,
    transition_probs: Option<Vec<Vec<Vec<f64>>>>,
}

impl HmmLatticeBuilder {
    /// Create a new builder
    pub fn new() -> Self {
        Self {
            time_steps: None,
            states_per_step: None,
            emission_probs: None,
            transition_probs: None,
        }
    }

    /// Set the number of time steps and states per step
    pub fn with_structure(
        mut self,
        time_steps: usize,
        states_per_step: Vec<usize>,
    ) -> Self {
        self.time_steps = Some(time_steps);
        self.states_per_step = Some(states_per_step);
        self
    }

    /// Set emission probabilities (will be converted to log probabilities)
    pub fn with_emission_probs(mut self, emission_probs: Vec<Vec<f64>>) -> Self {
        self.emission_probs = Some(emission_probs);
        self
    }

    /// Set transition probabilities (will be converted to log probabilities)
    pub fn with_transition_probs(
        mut self,
        transition_probs: Vec<Vec<Vec<f64>>>,
    ) -> Self {
        self.transition_probs = Some(transition_probs);
        self
    }

    /// Build the HMM lattice
    pub fn build(self) -> Result<HmmLattice, &'static str> {
        let time_steps = self.time_steps.ok_or("Time steps not set")?;
        let states_per_step = self.states_per_step.ok_or("States per step not set")?;
        let emission_probs = self
            .emission_probs
            .ok_or("Emission probabilities not set")?;
        let transition_probs = self
            .transition_probs
            .ok_or("Transition probabilities not set")?;

        // Validate dimensions
        if emission_probs.len() != time_steps {
            return Err("Emission probabilities length doesn't match time steps");
        }
        if transition_probs.len() != time_steps - 1 {
            return Err("Transition probabilities length should be time_steps - 1");
        }

        // Convert to log probabilities
        let log_emission_probs = emission_probs
            .into_iter()
            .map(|time_probs| {
                time_probs
                    .into_iter()
                    .map(|prob| {
                        if prob <= 0.0 {
                            std::f64::NEG_INFINITY
                        } else {
                            prob.ln()
                        }
                    })
                    .collect()
            })
            .collect();

        let log_transition_probs = transition_probs
            .into_iter()
            .map(|time_transitions| {
                time_transitions
                    .into_iter()
                    .map(|from_state_transitions| {
                        from_state_transitions
                            .into_iter()
                            .map(|prob| {
                                if prob <= 0.0 {
                                    std::f64::NEG_INFINITY
                                } else {
                                    prob.ln()
                                }
                            })
                            .collect()
                    })
                    .collect()
            })
            .collect();

        Ok(HmmLattice {
            time_steps,
            states_per_step,
            emission_probs: log_emission_probs,
            transition_probs: log_transition_probs,
        })
    }
}

impl Default for HmmLatticeBuilder {
    fn default() -> Self {
        Self::new()
    }
}