bambam 0.3.1

use crate::model::constraint::multimodal::constraint_config::{
    EnergyStateVariable, ModeLegDistanceConstraint, ModeLegEnergyConstraint, ModeLegTimeConstraint,
};
use crate::model::constraint::multimodal::sequence_trie::SubSequenceTrie;
use crate::model::constraint::multimodal::{
    multimodal_frontier_ops as ops, ConstraintConfig, DistanceConstraint, EnergyConstraint,
    TimeConstraint,
};
use bambam_core::model::state::{
    multimodal_state_ops as state_ops, MultimodalMapping, MultimodalStateMapping,
};
use bambam_core::model::{bambam_field, bambam_state};
use routee_compass_core::model::state::StateModelError;
use routee_compass_core::model::traversal::default::fieldname;
use routee_compass_core::model::{
    constraint::ConstraintModelError,
    network::Edge,
    state::{StateModel, StateVariable},
    unit::TimeUnit,
};
use std::collections::{HashMap, HashSet};
use std::num::NonZeroU64;
use uom::si::f64::{Energy, Length, Time};

#[derive(Debug)]
/// types of constraints to limit exponential search expansion in multimodal scenarios.
///
/// only deals with constraints associated with multimodal metadata, since metric-based
/// constraints must be applied _after_ access + traversal metrics have been run.
pub enum Constraint {
    /// Restrict routes to only use allowed transportation modes.
    AllowedModes(HashSet<String>),
    /// Limit the number of times each mode can be used in a route.
    ModeCounts(HashMap<String, usize>),
    /// Require routes to follow one of the specified mode sequences.
    ExactSequences(SubSequenceTrie),
    /// Set distance limit for the trip.
    DistanceConstraint(DistanceConstraint),
    /// Set time limit for the trip.
    TimeConstraint(TimeConstraint),
    /// Set maximum distance limits for each transportation mode.
    ModeDistanceLimit {
        mode_distance_limit: HashMap<String, DistanceConstraint>,
    },
    /// Set maximum time limits for each transportation mode.
    ModeTimeLimit {
        mode_time_limit: HashMap<String, TimeConstraint>,
    },
    /// Set maximum energy limits for each transportation mode.
    ModeEnergyLimit {
        mode_energy_limit: HashMap<String, EnergyConstraint>,
    },
    /// Set distance limits for specific modes on specific trip legs.
    ModeLegDistanceLimit {
        mode_leg_distance_limit: HashMap<String, ModeLegDistanceConstraint>,
    },
    /// Set time limits for specific modes on specific trip legs.
    ModeLegTimeLimit {
        mode_leg_time_limit: HashMap<String, ModeLegTimeConstraint>,
    },
    /// Set energy limits for specific modes on specific trip legs.
    ModeLegEnergyLimit {
        mode_leg_energy_limit: HashMap<String, ModeLegEnergyConstraint>,
    },
}

impl Constraint {
    /// validates an edge for traversal in a multimodal traversal
    pub fn valid_frontier(
        &self,
        edge_mode: &str,
        edge: &Edge,
        state: &[StateVariable],
        state_model: &StateModel,
        mode_to_state: &MultimodalStateMapping,
        max_trip_legs: NonZeroU64,
    ) -> Result<bool, ConstraintModelError> {
        use Constraint as MFC;

        match self {
            MFC::AllowedModes(items) => {
                let result = items.contains(edge_mode);
                Ok(result)
            }

            MFC::ModeCounts(limits) => validate_mode_counts(
                state,
                state_model,
                limits,
                max_trip_legs,
                mode_to_state,
                edge_mode,
            ),

            MFC::ExactSequences(trie) => validate_mode_sequences(
                state,
                state_model,
                trie,
                max_trip_legs,
                mode_to_state,
                edge_mode,
            ),

            MFC::DistanceConstraint(limit) => validate_trip_distance(state, state_model, limit),

            MFC::TimeConstraint(limit) => validate_trip_time(state, state_model, limit),

            MFC::ModeDistanceLimit {
                mode_distance_limit,
            } => validate_mode_distance(
                state,
                state_model,
                mode_distance_limit,
                max_trip_legs,
                mode_to_state,
                edge_mode,
            ),

            MFC::ModeTimeLimit { mode_time_limit } => validate_mode_time(
                state,
                state_model,
                mode_time_limit,
                max_trip_legs,
                mode_to_state,
                edge_mode,
            ),
            MFC::ModeEnergyLimit { mode_energy_limit } => validate_mode_energy(
                state,
                state_model,
                mode_energy_limit,
                max_trip_legs,
                mode_to_state,
                edge_mode,
            ),
            MFC::ModeLegDistanceLimit {
                mode_leg_distance_limit,
            } => validate_mode_leg_distance(
                state,
                state_model,
                mode_leg_distance_limit,
                max_trip_legs,
                mode_to_state,
                edge_mode,
            ),
            MFC::ModeLegTimeLimit {
                mode_leg_time_limit,
            } => validate_mode_leg_time(
                state,
                state_model,
                mode_leg_time_limit,
                max_trip_legs,
                mode_to_state,
                edge_mode,
            ),
            MFC::ModeLegEnergyLimit {
                mode_leg_energy_limit,
            } => validate_mode_leg_energy(
                state,
                state_model,
                mode_leg_energy_limit,
                max_trip_legs,
                mode_to_state,
                edge_mode,
            ),
        }
    }
}

impl TryFrom<&ConstraintConfig> for Constraint {
    type Error = ConstraintModelError;

    fn try_from(value: &ConstraintConfig) -> Result<Self, Self::Error> {
        use ConstraintConfig as MFCC;
        match value {
            MFCC::AllowedModes { values } => {
                let modes = values.iter().cloned().collect::<HashSet<_>>();
                Ok(Self::AllowedModes(modes))
            }
            MFCC::ModeCounts { values } => {
                let counts = values
                    .iter()
                    .map(|(k, v)| {
                        let v_usize: usize = v.get().try_into().map_err(|e| {
                            ConstraintModelError::ConstraintModelError(format!(
                                "while reading mode count limit: {e}"
                            ))
                        })?;
                        Ok((k.clone(), v_usize))
                    })
                    .collect::<Result<HashMap<_, _>, ConstraintModelError>>()?;
                Ok(Self::ModeCounts(counts))
            }
            MFCC::ExactSequences { values } => {
                let mut trie = SubSequenceTrie::new();
                for seq in values.iter() {
                    trie.insert_sequence(seq.clone());
                }
                Ok(Self::ExactSequences(trie))
            }
            MFCC::DistanceConstraint(c) => Ok(Self::DistanceConstraint(c.clone())),
            MFCC::TimeConstraint(c) => Ok(Self::TimeConstraint(c.clone())),
            MFCC::ModeDistanceLimit { values } => Ok(Self::ModeDistanceLimit {
                mode_distance_limit: values.clone(),
            }),
            MFCC::ModeTimeLimit { values } => Ok(Self::ModeTimeLimit {
                mode_time_limit: values.clone(),
            }),
            // MFCC::ModeEnergyLimit { mode_energy_limit } => Ok(Self::ModeEnergyLimit {
            //     mode_energy_limit: mode_energy_limit.clone(),
            // }),
            MFCC::ModeLegDistanceLimit { values } => Ok(Self::ModeLegDistanceLimit {
                mode_leg_distance_limit: values.clone(),
            }),
            MFCC::ModeLegTimeLimit { values } => Ok(Self::ModeLegTimeLimit {
                mode_leg_time_limit: values.clone(),
            }),
            // MFCC::ModeLegEnergyLimit {
            //     mode_leg_energy_limit,
            // } => Ok(Self::ModeLegEnergyLimit {
            //     mode_leg_energy_limit: mode_leg_energy_limit.clone(),
            // }),
        }
    }
}

type ConstraintResult = Result<bool, ConstraintModelError>;

/// runs the constraint model validation logic for mode count constraints
fn validate_mode_counts(
    state: &[StateVariable],
    state_model: &StateModel,
    limits: &HashMap<String, usize>,
    max_trip_legs: NonZeroU64,
    mode_to_state: &MultimodalMapping<String, i64>,
    edge_mode: &str,
) -> ConstraintResult {
    let mut counts = ops::get_mode_counts(state, state_model, max_trip_legs, mode_to_state)?;

    // simulate a mode transition if the incoming edge has a different mode than the trip's active mode
    let active_mode =
        state_ops::get_active_leg_mode(state, state_model, max_trip_legs, mode_to_state).map_err(
            |e| {
                ConstraintModelError::ConstraintModelError(format!(
                    "while applying mode count constraint model, {e}"
                ))
            },
        )?;
    if Some(edge_mode) != active_mode {
        counts
            .entry(edge_mode.to_string())
            .and_modify(|cnt| *cnt += 1)
            .or_insert(1);
    }

    Ok(ops::valid_mode_counts(&counts, limits))
}

/// runs the constraint model validation logic for exact sequence constraints
fn validate_mode_sequences(
    state: &[StateVariable],
    state_model: &StateModel,
    trie: &SubSequenceTrie,
    max_trip_legs: NonZeroU64,
    mode_to_state: &MultimodalMapping<String, i64>,
    edge_mode: &str,
) -> ConstraintResult {
    let mut modes = state_ops::get_mode_sequence(state, state_model, max_trip_legs, mode_to_state)
        .map_err(|e| {
            ConstraintModelError::ConstraintModelError(format!(
                "while testing for matching mode sub-sequence, had error: {e}"
            ))
        })?;

    // simulate a mode transition if the incoming edge has a different mode than the trip's active mode
    let active_mode =
        state_ops::get_active_leg_mode(state, state_model, max_trip_legs, mode_to_state).map_err(
            |e| {
                ConstraintModelError::ConstraintModelError(format!(
                    "while testing for matching mode sub-sequence, {e}"
                ))
            },
        )?;
    if Some(edge_mode) != active_mode {
        modes.push(edge_mode.to_string());
    }
    let is_match = trie.contains(&modes);
    Ok(is_match)
}

fn validate_trip_distance(
    state: &[StateVariable],
    state_model: &StateModel,
    limit: &DistanceConstraint,
) -> ConstraintResult {
    let dist: Length = state_model
        .get_distance(state, fieldname::TRIP_DISTANCE)
        .map_err(|e| {
            let msg = format!(
                "while retrieving '{}' from state: {e}",
                fieldname::TRIP_DISTANCE
            );
            ConstraintModelError::ConstraintModelError(msg)
        })?;
    let valid = limit.test(dist, false);
    Ok(valid)
}

fn validate_trip_time(
    state: &[StateVariable],
    state_model: &StateModel,
    limit: &TimeConstraint,
) -> ConstraintResult {
    let time: Time = state_model
        .get_time(state, fieldname::TRIP_TIME)
        .map_err(|e| {
            let msg = format!(
                "while retrieving '{}' from state: {e}",
                fieldname::TRIP_TIME
            );
            ConstraintModelError::ConstraintModelError(msg)
        })?;
    let valid = limit.test(time, false);
    log::debug!(
        "validating trip_time with time {:.2} minutes against limit {limit:?}. valid? {valid}",
        time.get::<uom::si::time::minute>(),
    );
    Ok(valid)
}

/// runs the constraint model validation logic for mode distance constraints
fn validate_mode_distance(
    state: &[StateVariable],
    state_model: &StateModel,
    limits: &HashMap<String, DistanceConstraint>,
    max_trip_legs: NonZeroU64,
    mode_to_state: &MultimodalMapping<String, i64>,
    edge_mode: &str,
) -> ConstraintResult {
    match limits.get(edge_mode) {
        Some(constraint) => {
            let value: Length = get_mode_distance(edge_mode, state, state_model)?;
            let ending_leg =
                check_mode_switch(state, state_model, max_trip_legs, mode_to_state, edge_mode)?;
            let valid = constraint.test(value, ending_leg);
            Ok(valid)
        }
        None => Ok(true),
    }
}

/// runs the constraint model validation logic for mode distance constraints
fn validate_mode_time(
    state: &[StateVariable],
    state_model: &StateModel,
    limits: &HashMap<String, TimeConstraint>,
    max_trip_legs: NonZeroU64,
    mode_to_state: &MultimodalMapping<String, i64>,
    edge_mode: &str,
) -> ConstraintResult {
    match limits.get(edge_mode) {
        Some(constraint) => {
            let value: Time = get_mode_time(edge_mode, state, state_model)?;
            let ending_leg =
                check_mode_switch(state, state_model, max_trip_legs, mode_to_state, edge_mode)?;
            let valid = constraint.test(value, ending_leg);
            Ok(valid)
        }
        None => Ok(true),
    }
}

/// runs the constraint model validation logic for mode distance constraints
fn validate_mode_energy(
    state: &[StateVariable],
    state_model: &StateModel,
    limits: &HashMap<String, EnergyConstraint>,
    max_trip_legs: NonZeroU64,
    mode_to_state: &MultimodalMapping<String, i64>,
    edge_mode: &str,
) -> ConstraintResult {
    match limits.get(edge_mode) {
        Some(constraint) => {
            let value: Energy = get_mode_energy(&constraint.variable, state, state_model)?;
            let ending_leg =
                check_mode_switch(state, state_model, max_trip_legs, mode_to_state, edge_mode)?;
            let valid = constraint.test(value, ending_leg);
            Ok(valid)
        }
        None => Ok(true),
    }
}

/// runs the constraint model validation logic for mode distance constraints
fn validate_mode_leg_distance(
    state: &[StateVariable],
    state_model: &StateModel,
    limits: &HashMap<String, ModeLegDistanceConstraint>,
    max_trip_legs: NonZeroU64,
    mode_to_state: &MultimodalMapping<String, i64>,
    edge_mode: &str,
) -> ConstraintResult {
    match limits.get(edge_mode) {
        Some(ModeLegDistanceConstraint { leg, constraint }) => {
            let matches = leg.matches(state, state_model, max_trip_legs)?;
            if !matches {
                return Ok(true);
            }
            let value: Length = match get_active_leg_distance(state, state_model)? {
                Some(v) => v,
                None => return Ok(true), // trip hasn't begun yet
            };
            let ending_leg =
                check_mode_switch(state, state_model, max_trip_legs, mode_to_state, edge_mode)?;
            let valid = constraint.test(value, ending_leg);
            Ok(valid)
        }
        None => Ok(true),
    }
}

/// runs the constraint model validation logic for mode time constraints
fn validate_mode_leg_time(
    state: &[StateVariable],
    state_model: &StateModel,
    limits: &HashMap<String, ModeLegTimeConstraint>,
    max_trip_legs: NonZeroU64,
    mode_to_state: &MultimodalMapping<String, i64>,
    edge_mode: &str,
) -> ConstraintResult {
    match limits.get(edge_mode) {
        Some(ModeLegTimeConstraint { leg, constraint }) => {
            let matches = leg.matches(state, state_model, max_trip_legs)?;
            if !matches {
                return Ok(true);
            }
            let value: Time = match get_active_leg_time(state, state_model)? {
                Some(v) => v,
                None => return Ok(true), // trip hasn't begun yet
            };
            let ending_leg =
                check_mode_switch(state, state_model, max_trip_legs, mode_to_state, edge_mode)?;
            let valid = constraint.test(value, ending_leg);
            Ok(valid)
        }
        None => Ok(true),
    }
}

/// runs the constraint model validation logic for mode time constraints
fn validate_mode_leg_energy(
    state: &[StateVariable],
    state_model: &StateModel,
    limits: &HashMap<String, ModeLegEnergyConstraint>,
    max_trip_legs: NonZeroU64,
    mode_to_state: &MultimodalMapping<String, i64>,
    edge_mode: &str,
) -> ConstraintResult {
    match limits.get(edge_mode) {
        Some(ModeLegEnergyConstraint { leg, constraint }) => {
            let matches = leg.matches(state, state_model, max_trip_legs)?;
            if !matches {
                return Ok(true);
            }
            let value: Energy = match get_active_leg_energy(state, state_model)? {
                Some(v) => v,
                None => return Ok(true), // trip hasn't begun yet
            };
            let mode_switch =
                check_mode_switch(state, state_model, max_trip_legs, mode_to_state, edge_mode)?;
            let valid = constraint.test(value, mode_switch);
            Ok(valid)
        }
        None => Ok(true),
    }
}

/// helper for retrieving distance values from the state vector.
fn get_mode_distance(
    edge_mode: &str,
    state: &[StateVariable],
    state_model: &StateModel,
) -> Result<Length, ConstraintModelError> {
    state_ops::get_mode_distance(state, edge_mode, state_model).map_err(|e| {
        let msg = format!("while retrieving '{edge_mode}' mode distance from state: {e}");
        ConstraintModelError::ConstraintModelError(msg)
    })
}

/// helper for retrieving time values from the state vector.
fn get_mode_time(
    edge_mode: &str,
    state: &[StateVariable],
    state_model: &StateModel,
) -> Result<Time, ConstraintModelError> {
    state_ops::get_mode_time(state, edge_mode, state_model).map_err(|e| {
        let msg = format!("while retrieving '{edge_mode}' mode time from state: {e}");
        ConstraintModelError::ConstraintModelError(msg)
    })
}

/// helper for retrieving all energy values from the state vector based on the constraint's
/// expected energy fieldnames.
fn get_mode_energy(
    variable: &EnergyStateVariable,
    state: &[StateVariable],
    state_model: &StateModel,
) -> Result<Energy, ConstraintModelError> {
    todo!("add a state_ops::get_mode_energy + update multimodal traversal model accounting for energy")
}

/// helper for retrieving energy values from the state vector.
fn get_energy(
    fieldname: &str,
    state: &[StateVariable],
    state_model: &StateModel,
) -> Result<Energy, ConstraintModelError> {
    state_model.get_energy(state, fieldname).map_err(|e| {
        let msg = format!("while retrieving {} from state: {e}", fieldname);
        ConstraintModelError::ConstraintModelError(msg)
    })
}

/// inspect the active trip mode and this edge's travel mode. if they do not match,
/// then we are ending a leg by making this transition.
fn check_mode_switch(
    state: &[StateVariable],
    state_model: &StateModel,
    max_trip_legs: NonZeroU64,
    mode_to_state: &MultimodalMapping<String, i64>,
    edge_mode: &str,
) -> Result<bool, ConstraintModelError> {
    let active_leg =
        state_ops::get_active_leg_mode(state, state_model, max_trip_legs, mode_to_state).map_err(
            |e| {
                ConstraintModelError::ConstraintModelError(format!(
                    "while checking for mode switch in metric-based constraint, {e}"
                ))
            },
        )?;
    match active_leg {
        None => Ok(false),
        Some(leg) => Ok(leg != edge_mode),
    }
}

fn get_active_leg_distance(
    state: &[StateVariable],
    state_model: &StateModel,
) -> Result<Option<Length>, ConstraintModelError> {
    let leg_idx_opt = state_ops::get_active_leg_idx(state, state_model).map_err(|e| {
        let msg = format!("while validating mode leg distance, {e}");
        ConstraintModelError::ConstraintModelError(msg)
    })?;
    let leg_idx = match leg_idx_opt {
        None => return Ok(None), // we haven't started the trip yet
        Some(idx) => idx,
    };
    let value: Length = state_ops::get_leg_distance(state, leg_idx, state_model).map_err(|e| {
        let msg = format!("while validating mode leg distance, {e}");
        ConstraintModelError::ConstraintModelError(msg)
    })?;
    Ok(Some(value))
}

fn get_active_leg_time(
    state: &[StateVariable],
    state_model: &StateModel,
) -> Result<Option<Time>, ConstraintModelError> {
    let leg_idx_opt = state_ops::get_active_leg_idx(state, state_model).map_err(|e| {
        let msg = format!("while validating mode leg time, {e}");
        ConstraintModelError::ConstraintModelError(msg)
    })?;
    let leg_idx = match leg_idx_opt {
        None => return Ok(None), // we haven't started the trip yet
        Some(idx) => idx,
    };
    let value: Time = state_ops::get_leg_time(state, leg_idx, state_model).map_err(|e| {
        let msg = format!("while validating mode leg time, {e}");
        ConstraintModelError::ConstraintModelError(msg)
    })?;
    Ok(Some(value))
}

fn get_active_leg_energy(
    state: &[StateVariable],
    state_model: &StateModel,
) -> Result<Option<Energy>, ConstraintModelError> {
    let leg_idx_opt = state_ops::get_active_leg_idx(state, state_model).map_err(|e| {
        let msg = format!("while validating mode leg energy, {e}");
        ConstraintModelError::ConstraintModelError(msg)
    })?;
    let leg_idx = match leg_idx_opt {
        None => return Ok(None), // we haven't started the trip yet
        Some(idx) => idx,
    };
    let value: Energy = state_ops::get_leg_energy(state, leg_idx, state_model).map_err(|e| {
        let msg = format!("while validating mode leg energy, {e}");
        ConstraintModelError::ConstraintModelError(msg)
    })?;
    Ok(Some(value))
}