pelorus 0.1.0

pub mod break_handler;
pub mod preprocessor;

use crate::probability::{emission_probability, transition_probability};
use crate::road_network::RoadNetwork;
use crate::types::{CandidatePoint, GpsMeasurement, SegmentId};
use crate::viterbi::{HmmLatticeBuilder, ViterbiPath};
use geo::algorithm::euclidean_distance::EuclideanDistance;
use geo::Point;

pub use break_handler::BreakHandler;
pub use preprocessor::TracePreprocessor;

/// Configuration parameters for map matching
#[derive(Debug, Clone)]
pub struct MapMatchingConfig {
    /// Maximum distance from GPS point to consider candidate segments (meters)
    pub max_candidate_distance: f64,
    /// Maximum number of candidate points per GPS measurement
    pub max_candidates_per_point: usize,
    /// GPS measurement noise standard deviation (meters)
    pub gps_sigma: f64,
    /// Minimum time gap to consider a trace break (seconds)
    pub break_time_threshold: f64,
    /// Minimum distance gap to consider a trace break (meters)
    pub break_distance_threshold: f64,
}

impl Default for MapMatchingConfig {
    fn default() -> Self {
        Self {
            max_candidate_distance: 50.0,
            max_candidates_per_point: 10,
            gps_sigma: 5.0,
            break_time_threshold: 120.0,     // 2 minutes
            break_distance_threshold: 200.0, // 200 meters
        }
    }
}

/// A matched GPS point with its corresponding road segment
#[derive(Debug, Clone, PartialEq)]
pub struct MatchedPoint {
    pub gps_measurement: GpsMeasurement,
    pub segment_id: SegmentId,
    pub projected_point: Point<f64>,
    pub emission_probability: f64,
}

/// A sequence of matched points representing a continuous trace
#[derive(Debug, Clone, PartialEq)]
pub struct MatchedTrace {
    pub points: Vec<MatchedPoint>,
    pub total_probability: f64,
}

/// Handles GPS trace break detection and management
#[derive(Debug, Clone, PartialEq)]
pub struct TraceBreak {
    pub start_index: usize,
    pub end_index: usize,
    pub time_gap: f64,
    pub distance_gap: f64,
}

/// Main map matching engine
pub struct MapMatcher {
    config: MapMatchingConfig,
    preprocessor: TracePreprocessor,
    break_handler: BreakHandler,
}

impl MapMatcher {
    /// Create a new map matcher with default configuration
    pub fn new() -> Self {
        let config = MapMatchingConfig::default();
        Self::with_config(config)
    }

    /// Create a new map matcher with custom configuration
    pub fn with_config(config: MapMatchingConfig) -> Self {
        let preprocessor = TracePreprocessor::new(&config);
        let break_handler = BreakHandler::new(&config);

        Self {
            config,
            preprocessor,
            break_handler,
        }
    }

    /// Match a GPS trace to road network segments
    pub fn match_trace(
        &self,
        gps_trace: &[GpsMeasurement],
        road_network: &RoadNetwork,
    ) -> Result<Vec<MatchedTrace>, MapMatchingError> {
        self.match_gps_to_roads(gps_trace, road_network)
    }

    /// Main map matching function that orchestrates the full pipeline
    pub fn match_gps_to_roads(
        &self,
        gps_trace: &[GpsMeasurement],
        road_network: &RoadNetwork,
    ) -> Result<Vec<MatchedTrace>, MapMatchingError> {
        if gps_trace.is_empty() {
            return Ok(Vec::new());
        }

        // Step 1: Preprocess GPS trace (remove close points, filter outliers)
        let preprocessed_trace = self.preprocessor.preprocess(gps_trace)?;

        if preprocessed_trace.is_empty() {
            return Err(MapMatchingError::InvalidInput(
                "All GPS points were filtered out during preprocessing".to_string(),
            ));
        }

        // Step 2: Find candidates for each point
        let candidate_sets =
            self.generate_candidates(&preprocessed_trace, road_network)?;

        // Check if we have valid candidates
        if candidate_sets.iter().any(|cs| cs.is_empty()) {
            // Some points have no candidates - this might indicate a break
            return self.handle_missing_candidates(
                &preprocessed_trace,
                road_network,
                candidate_sets,
            );
        }

        // Step 3: Build HMM lattice with emission and transition probabilities
        let lattice = self.build_hmm_lattice(&candidate_sets)?;

        // Step 4: Run Viterbi algorithm to find best path
        let best_path =
            crate::viterbi::solver::viterbi_path(&lattice).ok_or_else(|| {
                MapMatchingError::ViterbiSolverError("No valid path found".to_string())
            })?;

        // Step 5: Convert path to matched points
        let matched_points = self.path_to_matched_points(
            &best_path,
            &candidate_sets,
            &preprocessed_trace,
        )?;

        // Step 6: Handle breaks and split traces if necessary
        let breaks = self.break_handler.detect_breaks(&preprocessed_trace)?;

        if breaks.is_empty() {
            // Single continuous trace
            Ok(vec![MatchedTrace {
                points: matched_points,
                total_probability: best_path.probability(),
            }])
        } else {
            // Split into multiple traces at breaks
            let trace_segments =
                self.split_trace_at_breaks(&preprocessed_trace, &breaks);
            let mut matched_traces = Vec::new();

            for segment in trace_segments {
                if !segment.is_empty() {
                    let segment_match = self.match_segment(&segment, road_network)?;
                    matched_traces.push(segment_match);
                }
            }

            Ok(matched_traces)
        }
    }

    /// Handle the case where some GPS points have no candidates
    fn handle_missing_candidates(
        &self,
        gps_trace: &[GpsMeasurement],
        road_network: &RoadNetwork,
        candidate_sets: Vec<Vec<CandidatePoint<'_>>>,
    ) -> Result<Vec<MatchedTrace>, MapMatchingError> {
        // Identify points with no candidates
        let mut points_with_candidates = Vec::new();
        let mut current_segment = Vec::new();

        for (measurement, candidates) in gps_trace.iter().zip(candidate_sets.iter()) {
            if candidates.is_empty() {
                // No candidates for this point - end current segment
                if !current_segment.is_empty() {
                    points_with_candidates.push(current_segment);
                    current_segment = Vec::new();
                }
            } else {
                // Add point to current segment
                current_segment.push(measurement.clone());
            }
        }

        // Add final segment if not empty
        if !current_segment.is_empty() {
            points_with_candidates.push(current_segment);
        }

        // Match each segment that has candidates
        let mut matched_traces = Vec::new();
        for segment in points_with_candidates {
            if segment.len() >= 2 {
                // Need at least 2 points for matching
                let segment_match = self.match_segment(&segment, road_network)?;
                matched_traces.push(segment_match);
            }
        }

        if matched_traces.is_empty() {
            return Err(MapMatchingError::NoValidCandidates(
                "No GPS points could be matched to road network".to_string(),
            ));
        }

        Ok(matched_traces)
    }

    /// Match a single trace segment (no breaks)
    fn match_segment(
        &self,
        gps_trace: &[GpsMeasurement],
        road_network: &RoadNetwork,
    ) -> Result<MatchedTrace, MapMatchingError> {
        if gps_trace.is_empty() {
            return Ok(MatchedTrace {
                points: Vec::new(),
                total_probability: 1.0,
            });
        }

        // Generate candidate points for each GPS measurement
        let candidate_sets = self.generate_candidates(gps_trace, road_network)?;

        // Build HMM lattice
        let lattice = self.build_hmm_lattice(&candidate_sets)?;

        // Solve using Viterbi algorithm
        let best_path =
            crate::viterbi::solver::viterbi_path(&lattice).ok_or_else(|| {
                MapMatchingError::ViterbiSolverError("No valid path found".to_string())
            })?;

        // Convert path to matched points
        let matched_points =
            self.path_to_matched_points(&best_path, &candidate_sets, gps_trace)?;

        Ok(MatchedTrace {
            points: matched_points,
            total_probability: best_path.probability(),
        })
    }

    /// Generate candidate points for each GPS measurement
    fn generate_candidates<'a>(
        &self,
        gps_trace: &[GpsMeasurement],
        road_network: &'a RoadNetwork,
    ) -> Result<Vec<Vec<CandidatePoint<'a>>>, MapMatchingError> {
        let mut candidate_sets = Vec::new();

        for gps_measurement in gps_trace {
            let candidates =
                self.find_nearby_segments(gps_measurement, road_network)?;
            candidate_sets.push(candidates);
        }

        Ok(candidate_sets)
    }

    /// Find road segments near a GPS measurement
    fn find_nearby_segments<'a>(
        &self,
        gps_measurement: &GpsMeasurement,
        road_network: &'a RoadNetwork,
    ) -> Result<Vec<CandidatePoint<'a>>, MapMatchingError> {
        // Convert max distance from meters to degrees (rough approximation)
        let search_radius_degrees = self.config.max_candidate_distance / 111_000.0;

        // Query spatial index for nearby segments
        let nearby_segments = road_network
            .query_segments_near(gps_measurement.point, search_radius_degrees);

        let mut candidates = Vec::new();

        // Create candidate points for each nearby segment
        for segment in nearby_segments {
            // Create candidate point
            let candidate = CandidatePoint {
                segment_id: segment.id,
                segment,
                gps_point: gps_measurement.point,
            };

            // Check if projection is valid and within distance threshold
            if let Ok(_projected_point) = candidate.projected_point() {
                let distance_meters = candidate.distance_to_gps() * 111_000.0; // Convert to meters

                if distance_meters <= self.config.max_candidate_distance {
                    candidates.push(candidate);
                }
            }
        }

        // Sort candidates by distance (closest first)
        candidates.sort_by(|a, b| {
            let dist_a = a.distance_to_gps();
            let dist_b = b.distance_to_gps();
            dist_a
                .partial_cmp(&dist_b)
                .unwrap_or(std::cmp::Ordering::Equal)
        });

        // Limit to max candidates per point
        candidates.truncate(self.config.max_candidates_per_point);

        Ok(candidates)
    }

    /// Build HMM lattice from candidate sets
    fn build_hmm_lattice<'a>(
        &self,
        candidate_sets: &'a [Vec<CandidatePoint<'a>>],
    ) -> Result<crate::viterbi::HmmLattice, MapMatchingError> {
        if candidate_sets.is_empty() {
            return Err(MapMatchingError::InvalidInput(
                "Empty candidate sets".to_string(),
            ));
        }

        let time_steps = candidate_sets.len();
        let states_per_step: Vec<usize> =
            candidate_sets.iter().map(|cs| cs.len()).collect();

        // Build emission probabilities
        let mut emission_probs = Vec::with_capacity(time_steps);
        for candidates in candidate_sets {
            let mut time_step_probs = Vec::with_capacity(candidates.len());
            for candidate in candidates {
                let emission_prob = match candidate.projected_point() {
                    Ok(projected) => emission_probability(
                        &candidate.gps_point,
                        &projected,
                        self.config.gps_sigma,
                    ),
                    Err(_) => 0.001, // Small probability for projection errors
                };
                time_step_probs.push(emission_prob);
            }
            emission_probs.push(time_step_probs);
        }

        // Build transition probabilities
        let mut transition_probs = Vec::with_capacity(time_steps - 1);
        for time_step in 0..candidate_sets.len() - 1 {
            let current_candidates = &candidate_sets[time_step];
            let next_candidates = &candidate_sets[time_step + 1];

            let mut time_step_transitions =
                Vec::with_capacity(current_candidates.len());
            for current_candidate in current_candidates {
                let mut from_state_transitions =
                    Vec::with_capacity(next_candidates.len());
                for next_candidate in next_candidates {
                    // Calculate transition probability
                    let great_circle_distance = self.calculate_great_circle_distance(
                        current_candidate,
                        next_candidate,
                    )?;

                    // TODO: Calculate actual route distance via road network
                    // For now, use great circle distance as approximation
                    let route_distance = great_circle_distance * 1.2; // Assume 20% detour
                    let beta = 100.0; // Default beta parameter

                    let transition_prob = transition_probability(
                        great_circle_distance,
                        route_distance,
                        beta,
                    );
                    from_state_transitions.push(transition_prob);
                }
                time_step_transitions.push(from_state_transitions);
            }
            transition_probs.push(time_step_transitions);
        }

        // Build the lattice
        let lattice = HmmLatticeBuilder::new()
            .with_structure(time_steps, states_per_step)
            .with_emission_probs(emission_probs)
            .with_transition_probs(transition_probs)
            .build()
            .map_err(|e| MapMatchingError::ViterbiSolverError(e.to_string()))?;

        Ok(lattice)
    }

    /// Convert Viterbi path to matched points
    fn path_to_matched_points<'a>(
        &self,
        path: &ViterbiPath,
        candidate_sets: &'a [Vec<CandidatePoint<'a>>],
        gps_trace: &[GpsMeasurement],
    ) -> Result<Vec<MatchedPoint>, MapMatchingError> {
        let mut matched_points = Vec::new();

        for (time_step, &state_id) in path.states.iter().enumerate() {
            let candidate_idx = state_id % self.config.max_candidates_per_point;
            let candidate = &candidate_sets[time_step][candidate_idx];
            let gps_measurement = &gps_trace[time_step];

            let projected_point = candidate.projected_point()?;
            let emission_prob = emission_probability(
                &candidate.gps_point,
                &projected_point,
                self.config.gps_sigma,
            );

            matched_points.push(MatchedPoint {
                gps_measurement: gps_measurement.clone(),
                segment_id: candidate.segment_id,
                projected_point,
                emission_probability: emission_prob,
            });
        }

        Ok(matched_points)
    }

    /// Split trace into segments at detected breaks
    fn split_trace_at_breaks(
        &self,
        trace: &[GpsMeasurement],
        breaks: &[TraceBreak],
    ) -> Vec<Vec<GpsMeasurement>> {
        if breaks.is_empty() {
            return vec![trace.to_vec()];
        }

        let mut segments = Vec::new();
        let mut start_idx = 0;

        for break_info in breaks {
            if break_info.start_index > start_idx {
                segments.push(trace[start_idx..break_info.start_index].to_vec());
            }
            start_idx = break_info.end_index;
        }

        // Add remaining points after last break
        if start_idx < trace.len() {
            segments.push(trace[start_idx..].to_vec());
        }

        segments
    }

    /// Calculate great circle distance between two candidate points in meters
    fn calculate_great_circle_distance<'a>(
        &self,
        candidate1: &CandidatePoint<'a>,
        candidate2: &CandidatePoint<'a>,
    ) -> Result<f64, MapMatchingError> {
        let projected1 = candidate1.projected_point()?;
        let projected2 = candidate2.projected_point()?;

        let distance_degrees = projected1.euclidean_distance(&projected2);
        let distance_meters = distance_degrees * 111_000.0; // Rough conversion to meters

        Ok(distance_meters)
    }
}

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

/// Errors that can occur during map matching
#[derive(Debug, Clone, PartialEq)]
pub enum MapMatchingError {
    InvalidInput(String),
    NoValidCandidates(String),
    ViterbiSolverError(String),
    ProjectionError(String),
}

impl std::fmt::Display for MapMatchingError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            MapMatchingError::InvalidInput(msg) => write!(f, "Invalid input: {}", msg),
            MapMatchingError::NoValidCandidates(msg) => {
                write!(f, "No valid candidates: {}", msg)
            }
            MapMatchingError::ViterbiSolverError(msg) => {
                write!(f, "Viterbi solver error: {}", msg)
            }
            MapMatchingError::ProjectionError(msg) => {
                write!(f, "Projection error: {}", msg)
            }
        }
    }
}

impl std::error::Error for MapMatchingError {}

impl From<&'static str> for MapMatchingError {
    fn from(msg: &'static str) -> Self {
        MapMatchingError::ProjectionError(msg.to_string())
    }
}