timestretch 0.5.0

//! Shared adaptive analysis snapshot for hybrid and dual-plane paths.

use crate::analysis::beat::{default_subdivision_for_preset, detect_beats, snap_to_subdivision};
use crate::analysis::transient::{
    detect_transients_with_options, TransientDetectionOptions, TransientMap,
};
use crate::core::types::StretchParams;
use std::collections::BTreeMap;

/// Maximum FFT size for transient detection (smaller = faster, less frequency resolution).
const TRANSIENT_MAX_FFT: usize = 2048;
/// Maximum hop size for transient detection.
const TRANSIENT_MAX_HOP: usize = 512;
/// Minimum input length (in samples) for beat detection to be worthwhile.
const MIN_SAMPLES_FOR_BEAT_DETECTION: usize = 44_100;
/// Minimum distance (samples) between merged onset/beat positions.
const DEDUP_DISTANCE: usize = 512;
/// Sentinel strength value marking a beat-only segmentation anchor.
pub(crate) const BEAT_ANCHOR_STRENGTH: f32 = f32::NEG_INFINITY;
/// Hard bound for dynamic subdivision-grid generation.
const MAX_SUBDIVISION_GRID_POINTS: usize = 1_000_000;
/// Minimum transient anchors required before enabling live beat-grid merging.
const MIN_LIVE_BEAT_ANCHORS: usize = 2;
/// Minimum anchor strength considered "reliable" for live beat-grid merging.
const MIN_LIVE_BEAT_ANCHOR_STRENGTH: f32 = 0.2;
/// Minimum transient region used for weak onsets in shared segmentation.
const MIN_TRANSIENT_REGION_SECS: f64 = 0.005;
/// Upper bound for ratio-based transient region scaling.
const TRANSIENT_REGION_RATIO_SCALE_MAX: f64 = 1.6;
/// Maximum transient coverage ratio before forcing tonal-only render.
const TONAL_FORCE_MAX_TRANSIENT_COVERAGE: f64 = 0.03;
/// Maximum transient segment count before disabling tonal-force fallback.
const TONAL_FORCE_MAX_TRANSIENT_SEGMENTS: usize = 1;

/// Shared adaptive analysis result for a mono horizon.
#[derive(Debug, Clone)]
pub(crate) struct AdaptiveAnalysisSnapshot {
    pub transient_map: TransientMap,
    pub onsets: Vec<usize>,
    pub strengths: Vec<f32>,
}

/// Shared segmentation primitive used by legacy hybrid and dual-plane policy.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) struct AdaptiveSegment {
    pub start: usize,
    pub end: usize,
    pub is_transient: bool,
}

/// Builds a shared adaptive snapshot for mono content.
pub(crate) fn analyze_adaptive_snapshot_mono(
    input: &[f32],
    params: &StretchParams,
) -> AdaptiveAnalysisSnapshot {
    let confident_pre = params.pre_analysis.as_ref().filter(|artifact| {
        artifact.is_usable(params.sample_rate, params.beat_snap_confidence_threshold)
    });

    // A usable artifact is authoritative for transients (even when it holds
    // none), so online detection is skipped entirely — that is the CPU win
    // and the offline analysis is non-causal, hence at least as accurate.
    let transient_map = if input.is_empty() {
        empty_transient_map(params.hop_size.max(1))
    } else if let Some(artifact) = confident_pre {
        transient_map_from_artifact(artifact, input.len())
    } else {
        detect_transients_with_options(
            input,
            params.sample_rate,
            params.fft_size.min(TRANSIENT_MAX_FFT),
            params.hop_size.min(TRANSIENT_MAX_HOP),
            params.transient_sensitivity,
            TransientDetectionOptions::from_stretch_params(params),
        )
    };

    let mut onsets = transient_map.onsets.clone();
    let mut strengths = if transient_map.strengths.len() == transient_map.onsets.len() {
        transient_map.strengths.clone()
    } else {
        vec![1.0; transient_map.onsets.len()]
    };

    let mut detected_beat_grid = None;

    // Optionally merge live/offline beats into transient anchors.
    if params.beat_aware && input.len() >= MIN_SAMPLES_FOR_BEAT_DETECTION {
        let use_live_beats =
            confident_pre.is_some() || should_use_live_beat_aware_anchors(&strengths);
        if use_live_beats {
            let beats = if let Some(artifact) = confident_pre {
                artifact.beat_positions.as_slice()
            } else {
                let grid = detected_beat_grid
                    .get_or_insert_with(|| detect_beats(input, params.sample_rate));
                grid.beats.as_slice()
            };
            let (merged_onsets, merged_strengths) =
                merge_onsets_and_beats(&onsets, &strengths, beats, input.len());
            onsets = merged_onsets;
            strengths = merged_strengths;
        }
    }

    // Optionally snap transient anchors to beat subdivisions.
    let snap_bpm = params
        .bpm
        .or_else(|| confident_pre.map(|artifact| artifact.bpm))
        .filter(|bpm| bpm.is_finite() && *bpm > 0.0);

    if let Some(bpm) = snap_bpm {
        let tolerance_samples =
            params.sample_rate as f64 * (params.beat_snap_tolerance_ms / 1000.0).max(0.001);
        let subdivision = default_subdivision_for_preset(params.preset);
        let phase_offset = confident_pre
            .map(|artifact| artifact.downbeat_offset_samples)
            .unwrap_or(0);
        let beat_grid = generate_subdivision_grid_with_phase(
            bpm,
            params.sample_rate,
            input.len(),
            subdivision,
            phase_offset,
        );

        let strict_suppression = confident_pre.is_some();
        let had_transients = strengths.iter().copied().any(strength_marks_transient);
        let mut snapped: BTreeMap<usize, f32> = BTreeMap::new();
        for (i, &onset) in onsets.iter().enumerate() {
            let strength = strengths.get(i).copied().unwrap_or(1.0);
            let is_transient = strength_marks_transient(strength);
            let chosen = if is_transient {
                match snap_to_subdivision(onset as f64, &beat_grid, tolerance_samples) {
                    Some(snapped) => snapped.round() as usize,
                    None if strict_suppression => continue,
                    None => onset,
                }
            } else {
                onset
            };

            snapped
                .entry(chosen)
                .and_modify(|existing| *existing = merge_anchor_strength(*existing, strength))
                .or_insert(strength);
        }

        let snapped_has_transients = snapped.values().copied().any(strength_marks_transient);
        if !snapped.is_empty() && (!had_transients || snapped_has_transients) {
            onsets = snapped.keys().copied().collect();
            strengths = snapped.values().copied().collect();
        }
    }

    AdaptiveAnalysisSnapshot {
        transient_map,
        onsets,
        strengths,
    }
}

/// Reconstructs a sparse [`TransientMap`] from a pre-analysis artifact.
///
/// Onsets past `input_len` are dropped (batch inputs are assumed to be the
/// analyzed file from source frame 0). Band flux is written only at onset
/// frames; artifacts without band flux (version 1) leave the map empty there,
/// which downstream degrades to full four-band phase resets.
fn transient_map_from_artifact(
    artifact: &crate::core::preanalysis::PreAnalysisArtifact,
    input_len: usize,
) -> TransientMap {
    let hop = if artifact.analysis_hop_size > 0 {
        artifact.analysis_hop_size
    } else {
        TRANSIENT_MAX_HOP
    };

    let has_band_flux = artifact.onset_band_flux.len() == artifact.transient_onsets.len();
    let mut per_frame_band_flux = if has_band_flux {
        vec![[0.0f32; 4]; input_len / hop + 1]
    } else {
        Vec::new()
    };

    let mut onsets = Vec::with_capacity(artifact.transient_onsets.len());
    let mut onsets_fractional = Vec::with_capacity(artifact.transient_onsets.len());
    let mut strengths = Vec::with_capacity(artifact.transient_onsets.len());
    for (i, &onset) in artifact.transient_onsets.iter().enumerate() {
        if onset >= input_len {
            continue;
        }
        onsets.push(onset);
        onsets_fractional.push(onset as f64);
        strengths.push(artifact.strength_at(i));
        if has_band_flux {
            per_frame_band_flux[onset / hop] = artifact.onset_band_flux[i];
        }
    }

    TransientMap {
        onsets,
        onsets_fractional,
        strengths,
        flux: Vec::new(),
        hop_size: hop,
        per_frame_band_flux,
    }
}

#[inline]
fn empty_transient_map(hop_size: usize) -> TransientMap {
    TransientMap {
        onsets: Vec::new(),
        onsets_fractional: Vec::new(),
        strengths: Vec::new(),
        flux: Vec::new(),
        hop_size: hop_size.max(1),
        per_frame_band_flux: Vec::new(),
    }
}

/// Returns true when an anchor strength encodes a real transient.
#[inline]
pub(crate) fn strength_marks_transient(strength: f32) -> bool {
    strength.is_finite() && strength >= 0.0
}

/// Merges two anchor strengths at the same position.
///
/// Transients always win over beat-only anchors. If both are transients,
/// the stronger one is kept.
#[inline]
fn merge_anchor_strength(existing: f32, candidate: f32) -> f32 {
    let existing_is_transient = strength_marks_transient(existing);
    let candidate_is_transient = strength_marks_transient(candidate);

    match (existing_is_transient, candidate_is_transient) {
        (false, true) => candidate,
        (true, false) => existing,
        (true, true) => existing.max(candidate),
        (false, false) => existing,
    }
}

#[inline]
pub(crate) fn should_use_live_beat_aware_anchors(strengths: &[f32]) -> bool {
    strengths
        .iter()
        .copied()
        .filter(|&s| strength_marks_transient(s) && s >= MIN_LIVE_BEAT_ANCHOR_STRENGTH)
        .count()
        >= MIN_LIVE_BEAT_ANCHORS
}

/// Merges transient onsets with beat grid positions, deduplicating nearby entries.
///
/// Beat positions that fall within `DEDUP_DISTANCE` samples of an existing
/// transient onset are dropped to avoid creating overly short segments.
///
/// The returned strengths are aligned with returned onset positions:
/// - finite `>= 0.0`: transient anchor with that strength
/// - non-finite (`BEAT_ANCHOR_STRENGTH`): beat-only anchor
pub(crate) fn merge_onsets_and_beats(
    onsets: &[usize],
    strengths: &[f32],
    beats: &[usize],
    input_len: usize,
) -> (Vec<usize>, Vec<f32>) {
    let mut merged_positions: Vec<usize> = Vec::with_capacity(onsets.len() + beats.len());
    let mut merged_strengths: Vec<f32> = Vec::with_capacity(onsets.len() + beats.len());

    for (i, &onset) in onsets.iter().enumerate() {
        if onset >= input_len {
            continue;
        }
        merged_positions.push(onset);
        merged_strengths.push(strengths.get(i).copied().unwrap_or(1.0));
    }

    for &beat in beats {
        if beat >= input_len {
            continue;
        }
        let too_close = merged_positions
            .iter()
            .any(|&pos| pos.abs_diff(beat) < DEDUP_DISTANCE);
        if !too_close {
            merged_positions.push(beat);
            merged_strengths.push(BEAT_ANCHOR_STRENGTH);
        }
    }

    let mut pairs: Vec<(usize, f32)> = merged_positions.into_iter().zip(merged_strengths).collect();
    pairs.sort_unstable_by_key(|(pos, _)| *pos);

    let mut out_onsets = Vec::with_capacity(pairs.len());
    let mut out_strengths = Vec::with_capacity(pairs.len());
    for (pos, strength) in pairs {
        if let Some(last_pos) = out_onsets.last().copied() {
            if last_pos == pos {
                if let Some(last_strength) = out_strengths.last_mut() {
                    *last_strength = merge_anchor_strength(*last_strength, strength);
                }
                continue;
            }
        }
        out_onsets.push(pos);
        out_strengths.push(strength);
    }

    (out_onsets, out_strengths)
}

/// Builds shared transient/tonal segment ranges from adaptive anchors.
pub(crate) fn build_adaptive_segments(
    input_len: usize,
    onsets: &[usize],
    strengths: &[f32],
    params: &StretchParams,
    global_ratio: f64,
) -> Vec<AdaptiveSegment> {
    if onsets.is_empty() {
        return vec![AdaptiveSegment {
            start: 0,
            end: input_len,
            is_transient: false,
        }];
    }

    let transient_ratio_scale = global_ratio.clamp(1.0, TRANSIENT_REGION_RATIO_SCALE_MAX);
    let min_transient_size =
        (params.sample_rate as f64 * MIN_TRANSIENT_REGION_SECS).round() as usize;
    let max_transient_size =
        ((params.sample_rate as f64 * params.transient_region_secs * transient_ratio_scale).round()
            as usize)
            .max(min_transient_size);

    let mut segments = Vec::new();
    let mut pos = 0usize;

    for (i, &onset) in onsets.iter().enumerate() {
        if onset < pos {
            continue;
        }

        let tonal_end = onset.min(input_len);
        if tonal_end > pos {
            segments.push(AdaptiveSegment {
                start: pos,
                end: tonal_end,
                is_transient: false,
            });
        }

        let strength_raw = strengths.get(i).copied().unwrap_or(1.0);
        if !strength_marks_transient(strength_raw) {
            // Beat-only anchor: tonal split without transient region.
            pos = tonal_end;
            continue;
        }

        let strength = strength_raw.clamp(0.0, 1.0);
        let scale = 0.3 + 0.7 * strength as f64;
        let transient_size = min_transient_size
            + ((max_transient_size - min_transient_size) as f64 * scale) as usize;
        let trans_end = (onset + transient_size).min(input_len);
        if trans_end > onset {
            segments.push(AdaptiveSegment {
                start: onset,
                end: trans_end,
                is_transient: true,
            });
        }
        pos = trans_end;
    }

    if pos < input_len {
        segments.push(AdaptiveSegment {
            start: pos,
            end: input_len,
            is_transient: false,
        });
    }

    segments
}

/// Returns true when sparse transient coverage should force tonal-only render.
pub(crate) fn should_force_tonal_render(segments: &[AdaptiveSegment], input_len: usize) -> bool {
    if input_len == 0 || segments.is_empty() {
        return false;
    }

    let transient_count = segments.iter().filter(|s| s.is_transient).count();
    if transient_count == 0 || transient_count > TONAL_FORCE_MAX_TRANSIENT_SEGMENTS {
        return false;
    }

    let transient_samples: usize = segments
        .iter()
        .filter(|s| s.is_transient)
        .map(|s| s.end.saturating_sub(s.start))
        .sum();
    let coverage = transient_samples as f64 / input_len as f64;
    coverage <= TONAL_FORCE_MAX_TRANSIENT_COVERAGE
}

/// Generates a subdivision grid with a phase/downbeat offset.
pub(crate) fn generate_subdivision_grid_with_phase(
    bpm: f64,
    sample_rate: u32,
    total_samples: usize,
    subdivision: u32,
    phase_offset_samples: usize,
) -> Vec<f64> {
    if bpm <= 0.0 || subdivision == 0 || total_samples == 0 {
        return Vec::new();
    }

    let beat_interval_samples = 60.0 * sample_rate as f64 / bpm;
    let sub_interval = beat_interval_samples / subdivision as f64;
    if sub_interval <= 0.0 {
        return Vec::new();
    }

    let phase = (phase_offset_samples as f64).rem_euclid(sub_interval);
    let estimated_count = (total_samples as f64 / sub_interval).ceil() as usize + 1;
    let max_points = estimated_count.min(MAX_SUBDIVISION_GRID_POINTS);
    let mut grid = Vec::with_capacity(max_points);
    let mut pos = phase;
    for _ in 0..max_points {
        if pos >= total_samples as f64 {
            break;
        }
        grid.push(pos);
        pos += sub_interval;
    }
    grid
}

#[cfg(test)]
mod tests {
    use super::{
        build_adaptive_segments, should_force_tonal_render, AdaptiveSegment, BEAT_ANCHOR_STRENGTH,
    };
    use crate::core::types::StretchParams;

    #[test]
    fn build_adaptive_segments_creates_tonal_and_transient_regions() {
        let params = StretchParams::new(1.0)
            .with_sample_rate(44_100)
            .with_transient_region_secs(0.020);
        let segments = build_adaptive_segments(44_100, &[10_000], &[1.0], &params, 1.0);
        assert!(segments.len() >= 2, "expected tonal + transient regions");
        assert!(!segments[0].is_transient);
        assert!(segments.iter().any(|s| s.is_transient));
    }

    #[test]
    fn build_adaptive_segments_treats_beat_only_anchor_as_tonal_split() {
        let params = StretchParams::new(1.0).with_sample_rate(44_100);
        let segments = build_adaptive_segments(
            44_100,
            &[12_000, 20_000],
            &[BEAT_ANCHOR_STRENGTH, 1.0],
            &params,
            1.0,
        );
        assert!(segments.iter().any(|s| !s.is_transient && s.end == 12_000));
        assert!(segments.iter().any(|s| s.is_transient && s.start == 20_000));
    }

    #[test]
    fn should_force_tonal_render_for_single_tiny_transient() {
        let segments = vec![
            AdaptiveSegment {
                start: 0,
                end: 20_000,
                is_transient: false,
            },
            AdaptiveSegment {
                start: 20_000,
                end: 20_300,
                is_transient: true,
            },
            AdaptiveSegment {
                start: 20_300,
                end: 44_100,
                is_transient: false,
            },
        ];
        assert!(should_force_tonal_render(&segments, 44_100));
    }
}