use std::collections::BTreeMap;

use crate::sync_probe::analyze::onset_detection::PulseSegment;

use super::shortest_wrapped_difference;

pub(crate) fn collapse_segments_by_phase(
    segments: &[PulseSegment],
    pulse_period_s: f64,
    phase_tolerance_s: f64,
) -> Vec<PulseSegment> {
    if segments.len() <= 1 {
        return segments.to_vec();
    }

    let mut best: Option<(usize, f64, f64, Vec<PulseSegment>)> = None;
    for segment in segments {
        let candidate_phase = segment.start_s.rem_euclid(pulse_period_s);
        let collapsed = collapse_segments_for_reference_phase(
            segments,
            pulse_period_s,
            phase_tolerance_s,
            candidate_phase,
        );
        if collapsed.is_empty() {
            continue;
        }

        let cadence =
            retain_regular_cadence_segments(&collapsed, pulse_period_s, phase_tolerance_s);
        let total_duration_s = cadence
            .iter()
            .map(|segment| segment.duration_s)
            .sum::<f64>();
        let total_error_s = cadence
            .iter()
            .map(|segment| phase_error_for_reference(segment, candidate_phase, pulse_period_s))
            .sum::<f64>();

        match &best {
            Some((best_unique_slots, best_duration_s, best_error_s, _))
                if cadence.len() < *best_unique_slots
                    || (cadence.len() == *best_unique_slots
                        && total_duration_s < *best_duration_s)
                    || (cadence.len() == *best_unique_slots
                        && total_duration_s == *best_duration_s
                        && total_error_s >= *best_error_s) => {}
            _ => best = Some((cadence.len(), total_duration_s, total_error_s, cadence)),
        }
    }

    best.map(|(_, _, _, cadence)| cadence).unwrap_or_else(|| {
        retain_regular_cadence_segments(segments, pulse_period_s, phase_tolerance_s)
    })
}

fn collapse_segments_for_reference_phase(
    segments: &[PulseSegment],
    pulse_period_s: f64,
    phase_tolerance_s: f64,
    reference_phase_s: f64,
) -> Vec<PulseSegment> {
    let mut by_index = BTreeMap::<i64, PulseSegment>::new();
    for segment in segments {
        let pulse_index = ((segment.start_s - reference_phase_s) / pulse_period_s).round() as i64;
        let expected_start_s = reference_phase_s + pulse_index as f64 * pulse_period_s;
        let phase_error_s =
            shortest_wrapped_difference(segment.start_s - expected_start_s, pulse_period_s).abs();
        if phase_error_s > phase_tolerance_s {
            continue;
        }

        match by_index.get_mut(&pulse_index) {
            Some(best) if !prefer_segment(segment, best, expected_start_s) => {}
            Some(best) => *best = *segment,
            None => {
                by_index.insert(pulse_index, *segment);
            }
        }
    }

    by_index.into_values().collect()
}

fn retain_regular_cadence_segments(
    segments: &[PulseSegment],
    pulse_period_s: f64,
    phase_tolerance_s: f64,
) -> Vec<PulseSegment> {
    if segments.len() <= 2 {
        return segments.to_vec();
    }

    let mut best_count = vec![1usize; segments.len()];
    let mut best_duration_s = segments
        .iter()
        .map(|segment| segment.duration_s)
        .collect::<Vec<_>>();
    let mut best_error_s = vec![0.0_f64; segments.len()];
    let mut previous_index = vec![None::<usize>; segments.len()];

    for current in 0..segments.len() {
        for previous in 0..current {
            let Some(cadence_error_s) = cadence_error(
                segments[previous].start_s,
                segments[current].start_s,
                pulse_period_s,
                phase_tolerance_s,
            ) else {
                continue;
            };

            let candidate_count = best_count[previous] + 1;
            let candidate_duration_s = best_duration_s[previous] + segments[current].duration_s;
            let candidate_error_s = best_error_s[previous] + cadence_error_s;
            if candidate_count > best_count[current]
                || (candidate_count == best_count[current]
                    && candidate_duration_s > best_duration_s[current])
                || (candidate_count == best_count[current]
                    && candidate_duration_s == best_duration_s[current]
                    && candidate_error_s < best_error_s[current])
            {
                best_count[current] = candidate_count;
                best_duration_s[current] = candidate_duration_s;
                best_error_s[current] = candidate_error_s;
                previous_index[current] = Some(previous);
            }
        }
    }

    let mut best_terminal = 0usize;
    for index in 1..segments.len() {
        if best_count[index] > best_count[best_terminal]
            || (best_count[index] == best_count[best_terminal]
                && best_duration_s[index] > best_duration_s[best_terminal])
            || (best_count[index] == best_count[best_terminal]
                && best_duration_s[index] == best_duration_s[best_terminal]
                && best_error_s[index] < best_error_s[best_terminal])
        {
            best_terminal = index;
        }
    }

    let mut selected = Vec::new();
    let mut cursor = Some(best_terminal);
    while let Some(index) = cursor {
        selected.push(segments[index]);
        cursor = previous_index[index];
    }
    selected.reverse();
    selected
}

fn cadence_error(
    previous_start_s: f64,
    current_start_s: f64,
    pulse_period_s: f64,
    phase_tolerance_s: f64,
) -> Option<f64> {
    let delta_s = current_start_s - previous_start_s;
    if delta_s <= 0.0 {
        return None;
    }

    let pulse_steps = (delta_s / pulse_period_s).round();
    if pulse_steps < 1.0 {
        return None;
    }

    let error_s = (delta_s - pulse_steps * pulse_period_s).abs();
    (error_s <= phase_tolerance_s).then_some(error_s)
}

fn phase_error_for_reference(
    segment: &PulseSegment,
    reference_phase_s: f64,
    pulse_period_s: f64,
) -> f64 {
    let pulse_index = ((segment.start_s - reference_phase_s) / pulse_period_s).round() as i64;
    let expected_start_s = reference_phase_s + pulse_index as f64 * pulse_period_s;
    shortest_wrapped_difference(segment.start_s - expected_start_s, pulse_period_s).abs()
}

fn prefer_segment(
    candidate: &PulseSegment,
    incumbent: &PulseSegment,
    expected_start_s: f64,
) -> bool {
    if candidate.duration_s > incumbent.duration_s {
        return true;
    }
    if candidate.duration_s < incumbent.duration_s {
        return false;
    }

    let candidate_error_s = (candidate.start_s - expected_start_s).abs();
    let incumbent_error_s = (incumbent.start_s - expected_start_s).abs();
    candidate_error_s < incumbent_error_s
}

#[cfg(test)]
mod tests {
    use super::*;

    fn seg(start_s: f64, duration_s: f64) -> PulseSegment {
        PulseSegment {
            start_s,
            end_s: start_s + duration_s,
            duration_s,
        }
    }

    #[test]
    fn collapse_segments_by_phase_falls_back_when_phase_tolerance_excludes_everything() {
        let segments = [seg(4.0, 0.12), seg(5.0, 0.12)];
        let collapsed = collapse_segments_by_phase(&segments, 1.0, -0.01);
        assert_eq!(collapsed, segments);
    }

    #[test]
    fn collapse_segments_by_phase_replaces_shorter_slot_winner_with_longer_one() {
        let collapsed = collapse_segments_by_phase(
            &[seg(4.02, 0.04), seg(4.05, 0.12), seg(5.02, 0.12)],
            1.0,
            0.32,
        );
        assert_eq!(collapsed.len(), 2);
        assert!((collapsed[0].start_s - 4.05).abs() < 0.001);
    }

    #[test]
    fn collapse_segments_by_phase_keeps_longer_incumbent_and_breaks_equal_duration_ties_by_error() {
        let keep_longer = collapse_segments_by_phase(
            &[seg(4.02, 0.12), seg(4.05, 0.04), seg(5.02, 0.12)],
            1.0,
            0.32,
        );
        assert_eq!(keep_longer.len(), 2);
        assert!((keep_longer[0].start_s - 4.02).abs() < 0.001);

        let prefer_closer = collapse_segments_by_phase(
            &[seg(4.08, 0.12), seg(4.02, 0.12), seg(5.02, 0.12)],
            1.0,
            0.32,
        );
        assert_eq!(prefer_closer.len(), 2);
        assert!((prefer_closer[0].start_s - 4.02).abs() < 0.001);
    }

    #[test]
    fn retain_regular_cadence_segments_prefers_longer_predecessor_chain_when_counts_tie() {
        let selected = retain_regular_cadence_segments(
            &[
                seg(0.00, 0.10),
                seg(0.49, 0.05),
                seg(0.51, 0.20),
                seg(1.50, 0.10),
            ],
            1.0,
            0.1,
        );
        assert_eq!(selected.len(), 2);
        assert!((selected[0].start_s - 0.51).abs() < 0.001);
        assert!((selected[1].start_s - 1.50).abs() < 0.001);
    }

    #[test]
    fn retain_regular_cadence_segments_prefers_best_terminal_by_duration_then_error() {
        let longer_terminal = retain_regular_cadence_segments(
            &[
                seg(0.00, 0.10),
                seg(1.00, 0.10),
                seg(0.49, 0.05),
                seg(1.50, 0.20),
            ],
            1.0,
            0.1,
        );
        assert_eq!(longer_terminal, vec![seg(0.49, 0.05), seg(1.50, 0.20)]);

        let lower_error_terminal = retain_regular_cadence_segments(
            &[
                seg(0.00, 0.10),
                seg(1.05, 0.10),
                seg(0.49, 0.10),
                seg(1.50, 0.10),
            ],
            1.0,
            0.1,
        );
        assert_eq!(lower_error_terminal, vec![seg(0.49, 0.10), seg(1.50, 0.10)]);
    }

    #[test]
    fn cadence_error_rejects_non_forward_and_sub_pulse_gaps() {
        assert_eq!(cadence_error(1.0, 1.0, 1.0, 0.1), None);
        assert_eq!(cadence_error(1.0, 0.9, 1.0, 0.1), None);
        assert_eq!(cadence_error(1.0, 1.49, 1.0, 0.1), None);
        let multi_step = cadence_error(1.0, 3.02, 1.0, 0.05).expect("multi-step cadence");
        assert!((multi_step - 0.02).abs() < 0.000_001);
        assert_eq!(cadence_error(1.0, 2.2, 1.0, 0.05), None);
    }

    #[test]
    fn prefer_segment_prefers_longer_and_then_closer_candidates() {
        assert!(prefer_segment(&seg(4.0, 0.20), &seg(4.0, 0.10), 4.0));
        assert!(!prefer_segment(&seg(4.0, 0.05), &seg(4.0, 0.10), 4.0));
        assert!(prefer_segment(&seg(4.02, 0.10), &seg(4.08, 0.10), 4.0));
        assert!(!prefer_segment(&seg(4.08, 0.10), &seg(4.02, 0.10), 4.0));
    }
}
