use anyhow::{bail, Result};
use std::collections::BTreeMap;

use crate::sync_probe::analyze::report::SyncAnalysisReport;

use super::{median, PulseSegment};

#[path = "correlation_collapse.rs"]
mod collapse;
pub(super) use collapse::collapse_segments_by_phase;

const MARKER_WIDTH_MULTIPLIER: f64 = 1.5;
const PHASE_TOLERANCE_WIDTH_MULTIPLIER: f64 = 2.5;
const STARTUP_PHASE_ANCHOR_TOLERANCE_FRACTION: f64 = 1.0 / 3.0;

#[cfg_attr(not(test), allow(dead_code))]
pub(super) fn correlate_onsets(
    video_onsets_s: &[f64],
    audio_onsets_s: &[f64],
    pulse_period_s: f64,
    max_pair_gap_s: f64,
) -> Result<SyncAnalysisReport> {
    if video_onsets_s.is_empty() {
        bail!("video onset list is empty");
    }
    if audio_onsets_s.is_empty() {
        bail!("audio onset list is empty");
    }
    if max_pair_gap_s <= 0.0 {
        bail!("max pair gap must stay positive");
    }
    if pulse_period_s <= 0.0 {
        bail!("pulse period must stay positive");
    }

    let (video_onsets_s, audio_onsets_s, common_window) =
        trim_onsets_to_common_activity_window(video_onsets_s, audio_onsets_s, max_pair_gap_s);
    let expected_start_skew_ms = (audio_onsets_s[0] - video_onsets_s[0]) * 1000.0;
    let video_pulses = index_onsets_by_spacing(&video_onsets_s, pulse_period_s);
    let audio_pulses = index_onsets_by_spacing(&audio_onsets_s, pulse_period_s);
    let offset_candidates = candidate_index_offsets(&video_pulses, &audio_pulses);
    let mut skews_ms = best_skews_for_index_offsets(
        &video_pulses,
        &audio_pulses,
        &offset_candidates,
        max_pair_gap_s,
        expected_start_skew_ms,
    );

    if skews_ms.is_empty() && video_onsets_s.len() == 1 && audio_onsets_s.len() == 1 {
        let video_phase_s = estimate_phase(&video_onsets_s, pulse_period_s);
        let audio_phase_s = estimate_phase(&audio_onsets_s, pulse_period_s);
        let phase_skew_ms =
            shortest_wrapped_difference(audio_phase_s - video_phase_s, pulse_period_s) * 1000.0;
        if phase_skew_ms.abs() <= max_pair_gap_s * 1000.0 {
            skews_ms.push(phase_skew_ms);
        }
    }

    if skews_ms.is_empty() {
        bail!("no audio/video pulse pairs were close enough to compare");
    }

    Ok(sync_report_from_skews(
        common_window.filter_onsets(&video_onsets_s),
        common_window.filter_onsets(&audio_onsets_s),
        skews_ms,
    ))
}

pub(crate) fn correlate_segments(
    video_segments: &[PulseSegment],
    audio_segments: &[PulseSegment],
    pulse_period_s: f64,
    pulse_width_s: f64,
    marker_tick_period: u32,
    max_pair_gap_s: f64,
) -> Result<SyncAnalysisReport> {
    if video_segments.is_empty() {
        bail!("video onset list is empty");
    }
    if audio_segments.is_empty() {
        bail!("audio onset list is empty");
    }
    if pulse_period_s <= 0.0 {
        bail!("pulse period must stay positive");
    }
    if pulse_width_s <= 0.0 {
        bail!("pulse width must stay positive");
    }
    if marker_tick_period == 0 {
        bail!("marker tick period must stay positive");
    }
    if max_pair_gap_s <= 0.0 {
        bail!("max pair gap must stay positive");
    }

    let phase_tolerance_s = segment_phase_tolerance(pulse_period_s, pulse_width_s, max_pair_gap_s);
    let video_segments =
        collapse_segments_by_phase(video_segments, pulse_period_s, phase_tolerance_s);
    let audio_segments =
        collapse_segments_by_phase(audio_segments, pulse_period_s, phase_tolerance_s);

    let video_onsets_s = video_segments
        .iter()
        .map(|segment| segment.start_s)
        .collect::<Vec<_>>();
    let audio_onsets_s = audio_segments
        .iter()
        .map(|segment| segment.start_s)
        .collect::<Vec<_>>();
    if video_onsets_s.is_empty() {
        bail!("video onset list is empty");
    }
    if audio_onsets_s.is_empty() {
        bail!("audio onset list is empty");
    }

    let (video_onsets_s, audio_onsets_s, common_window) =
        trim_onsets_to_common_activity_window(&video_onsets_s, &audio_onsets_s, max_pair_gap_s);
    let expected_start_skew_ms = (audio_onsets_s[0] - video_onsets_s[0]) * 1000.0;
    let video_marker_onsets = marker_onsets(&video_segments, pulse_width_s);
    let audio_marker_onsets = marker_onsets(&audio_segments, pulse_width_s);
    let video_marker_onsets = common_window.filter_onsets(&video_marker_onsets);
    let audio_marker_onsets = common_window.filter_onsets(&audio_marker_onsets);
    let video_indexed = index_onsets_by_spacing(&video_onsets_s, pulse_period_s);
    let audio_indexed = index_onsets_by_spacing(&audio_onsets_s, pulse_period_s);
    let offset_candidates = marker_index_offsets(
        &video_indexed,
        &audio_indexed,
        &video_marker_onsets,
        &audio_marker_onsets,
    );
    let mut skews_ms = best_skews_for_index_offsets(
        &video_indexed,
        &audio_indexed,
        &offset_candidates,
        max_pair_gap_s,
        expected_start_skew_ms,
    );

    if skews_ms.is_empty() && video_onsets_s.len() == 1 && audio_onsets_s.len() == 1 {
        let video_phase_s = estimate_phase(&video_onsets_s, pulse_period_s);
        let audio_phase_s = estimate_phase(&audio_onsets_s, pulse_period_s);
        let phase_skew_ms =
            shortest_wrapped_difference(audio_phase_s - video_phase_s, pulse_period_s) * 1000.0;
        if phase_skew_ms.abs() <= max_pair_gap_s * 1000.0 {
            skews_ms.push(phase_skew_ms);
        }
    }

    if skews_ms.is_empty() {
        bail!("no audio/video pulse pairs were close enough to compare");
    }

    Ok(sync_report_from_skews(
        video_onsets_s,
        audio_onsets_s,
        skews_ms,
    ))
}

#[derive(Clone, Copy)]
struct CommonActivityWindow {
    start_s: f64,
    end_s: f64,
}

impl CommonActivityWindow {
    fn filter_onsets(self, onsets_s: &[f64]) -> &[f64] {
        let start = onsets_s.partition_point(|onset_s| *onset_s < self.start_s);
        let end = onsets_s.partition_point(|onset_s| *onset_s <= self.end_s);
        &onsets_s[start..end]
    }
}

fn trim_onsets_to_common_activity_window<'a>(
    video_onsets_s: &'a [f64],
    audio_onsets_s: &'a [f64],
    max_pair_gap_s: f64,
) -> (&'a [f64], &'a [f64], CommonActivityWindow) {
    let common_window = CommonActivityWindow {
        start_s: (video_onsets_s
            .first()
            .copied()
            .expect("validated video onset list is not empty")
            .max(
                audio_onsets_s
                    .first()
                    .copied()
                    .expect("validated audio onset list is not empty"),
            )
            - max_pair_gap_s)
            .max(0.0),
        end_s: video_onsets_s
            .last()
            .copied()
            .expect("validated video onset list is not empty")
            .min(
                audio_onsets_s
                    .last()
                    .copied()
                    .expect("validated audio onset list is not empty"),
            )
            + max_pair_gap_s,
    };
    let trimmed_video_onsets_s = common_window.filter_onsets(video_onsets_s);
    let trimmed_audio_onsets_s = common_window.filter_onsets(audio_onsets_s);
    if trimmed_video_onsets_s.is_empty() || trimmed_audio_onsets_s.is_empty() {
        return (video_onsets_s, audio_onsets_s, common_window);
    }

    (
        trimmed_video_onsets_s,
        trimmed_audio_onsets_s,
        common_window,
    )
}

fn segment_phase_tolerance(pulse_period_s: f64, pulse_width_s: f64, max_pair_gap_s: f64) -> f64 {
    (pulse_width_s * PHASE_TOLERANCE_WIDTH_MULTIPLIER)
        .max(max_pair_gap_s.min(pulse_period_s / 3.0))
        .min(pulse_period_s / 2.5)
}

pub(super) fn estimate_phase(onsets_s: &[f64], pulse_period_s: f64) -> f64 {
    let (sum_sin, sum_cos) =
        onsets_s
            .iter()
            .copied()
            .fold((0.0_f64, 0.0_f64), |(sum_sin, sum_cos), onset| {
                let wrapped = onset.rem_euclid(pulse_period_s);
                let angle = (wrapped / pulse_period_s) * std::f64::consts::TAU;
                (sum_sin + angle.sin(), sum_cos + angle.cos())
            });

    let mean_angle = sum_sin.atan2(sum_cos).rem_euclid(std::f64::consts::TAU);
    (mean_angle / std::f64::consts::TAU) * pulse_period_s
}

pub(super) fn index_onsets_by_spacing(onsets_s: &[f64], pulse_period_s: f64) -> BTreeMap<i64, f64> {
    let mut indexed = BTreeMap::new();
    let Some(first_onset) = onsets_s.first().copied() else {
        return indexed;
    };

    let mut pulse_index = 0_i64;
    let mut previous_onset = first_onset;
    indexed.insert(pulse_index, first_onset);
    for onset in onsets_s.iter().copied().skip(1) {
        let pulse_steps = ((onset - previous_onset) / pulse_period_s).round().max(1.0) as i64;
        pulse_index += pulse_steps;
        indexed.insert(pulse_index, onset);
        previous_onset = onset;
    }
    indexed
}

pub(super) fn candidate_index_offsets(
    video_indexed: &BTreeMap<i64, f64>,
    audio_indexed: &BTreeMap<i64, f64>,
) -> Vec<i64> {
    if video_indexed.is_empty() || audio_indexed.is_empty() {
        return Vec::new();
    }

    let video_min = *video_indexed
        .keys()
        .next()
        .expect("non-empty indexed video onset map has a first key");
    let video_max = *video_indexed
        .keys()
        .next_back()
        .expect("non-empty indexed video onset map has a last key");
    let audio_min = *audio_indexed
        .keys()
        .next()
        .expect("non-empty indexed audio onset map has a first key");
    let audio_max = *audio_indexed
        .keys()
        .next_back()
        .expect("non-empty indexed audio onset map has a last key");

    (audio_min - video_max..=audio_max - video_min).collect()
}

pub(super) fn marker_index_offsets(
    video_indexed: &BTreeMap<i64, f64>,
    audio_indexed: &BTreeMap<i64, f64>,
    video_marker_onsets: &[f64],
    audio_marker_onsets: &[f64],
) -> Vec<i64> {
    let mut offsets = Vec::new();
    if !video_marker_onsets.is_empty() && !audio_marker_onsets.is_empty() {
        let video_markers = pulse_indices_for_onsets(video_indexed, video_marker_onsets);
        let audio_markers = pulse_indices_for_onsets(audio_indexed, audio_marker_onsets);
        for video_marker in &video_markers {
            for audio_marker in &audio_markers {
                offsets.push(audio_marker - video_marker);
            }
        }
    }

    offsets.extend(candidate_index_offsets(video_indexed, audio_indexed));
    offsets.sort_unstable();
    offsets.dedup();
    offsets
}

fn pulse_indices_for_onsets(indexed: &BTreeMap<i64, f64>, marker_onsets: &[f64]) -> Vec<i64> {
    marker_onsets
        .iter()
        .filter_map(|marker_onset| {
            indexed.iter().find_map(|(pulse_index, onset)| {
                ((onset - marker_onset).abs() < 0.000_001).then_some(*pulse_index)
            })
        })
        .collect()
}

fn best_skews_for_index_offsets(
    video_indexed: &BTreeMap<i64, f64>,
    audio_indexed: &BTreeMap<i64, f64>,
    offset_candidates: &[i64],
    max_pair_gap_s: f64,
    expected_start_skew_ms: f64,
) -> Vec<f64> {
    let max_pair_gap_ms = max_pair_gap_s * 1000.0;
    let startup_phase_anchor_tolerance_ms =
        max_pair_gap_ms * STARTUP_PHASE_ANCHOR_TOLERANCE_FRACTION;
    let mut best: Option<(bool, usize, f64, f64, Vec<f64>)> = None;

    for offset in offset_candidates.iter().copied() {
        let skews_ms = video_indexed
            .iter()
            .filter_map(|(pulse_index, video_time)| {
                audio_indexed
                    .get(&(pulse_index + offset))
                    .map(|audio_time| (audio_time - video_time) * 1000.0)
            })
            .filter(|skew_ms| skew_ms.abs() <= max_pair_gap_ms)
            .collect::<Vec<_>>();
        if skews_ms.is_empty() {
            continue;
        }

        let mean_abs_skew_ms =
            skews_ms.iter().map(|skew_ms| skew_ms.abs()).sum::<f64>() / skews_ms.len() as f64;
        let startup_phase_anchor_error_ms = (skews_ms[0] - expected_start_skew_ms).abs();
        let startup_phase_anchor_consistent =
            startup_phase_anchor_error_ms <= startup_phase_anchor_tolerance_ms;
        match &best {
            Some((
                best_anchor_consistent,
                best_count,
                best_anchor_error_ms,
                best_mean_abs_skew_ms,
                _,
            )) if startup_phase_anchor_consistent < *best_anchor_consistent
                || (startup_phase_anchor_consistent == *best_anchor_consistent
                    && (skews_ms.len() < *best_count
                        || (skews_ms.len() == *best_count
                            && (startup_phase_anchor_error_ms > *best_anchor_error_ms
                                || (startup_phase_anchor_error_ms == *best_anchor_error_ms
                                    && mean_abs_skew_ms >= *best_mean_abs_skew_ms))))) => {}
            _ => {
                best = Some((
                    startup_phase_anchor_consistent,
                    skews_ms.len(),
                    startup_phase_anchor_error_ms,
                    mean_abs_skew_ms,
                    skews_ms,
                ))
            }
        }
    }

    best.map(|(_, _, _, _, skews)| skews).unwrap_or_default()
}

pub(super) fn marker_onsets(segments: &[PulseSegment], pulse_width_s: f64) -> Vec<f64> {
    let threshold = pulse_width_s * MARKER_WIDTH_MULTIPLIER;
    segments
        .iter()
        .filter(|segment| segment.duration_s >= threshold)
        .map(|segment| segment.start_s)
        .collect()
}

pub(super) fn shortest_wrapped_difference(delta_s: f64, pulse_period_s: f64) -> f64 {
    let half_period = pulse_period_s / 2.0;
    ((delta_s + half_period).rem_euclid(pulse_period_s)) - half_period
}

fn sync_report_from_skews(
    video_onsets_s: &[f64],
    audio_onsets_s: &[f64],
    skews_ms: Vec<f64>,
) -> SyncAnalysisReport {
    let mut sorted_skews = skews_ms.clone();
    sorted_skews.sort_by(|left, right| left.total_cmp(right));
    let first_skew_ms = *skews_ms.first().expect("paired skew list is not empty");
    let last_skew_ms = *skews_ms.last().expect("paired skew list is not empty");
    let mean_skew_ms = skews_ms.iter().sum::<f64>() / skews_ms.len() as f64;
    let median_skew_ms = median(sorted_skews);
    let max_abs_skew_ms = skews_ms
        .iter()
        .copied()
        .map(f64::abs)
        .fold(0.0_f64, f64::max);

    SyncAnalysisReport {
        video_event_count: video_onsets_s.len(),
        audio_event_count: audio_onsets_s.len(),
        paired_event_count: skews_ms.len(),
        first_skew_ms,
        last_skew_ms,
        mean_skew_ms,
        median_skew_ms,
        max_abs_skew_ms,
        drift_ms: last_skew_ms - first_skew_ms,
        skews_ms,
        video_onsets_s: video_onsets_s.to_vec(),
        audio_onsets_s: audio_onsets_s.to_vec(),
    }
}
