//! Chop engine: slice a sample into individual one-shots by transient, by equal
//! divisions, or by a BPM grid.
//!
//! Transient detection reuses the same spectral-flux onset measure the analysis
//! pipeline uses for `onset_strength` (see [`crate::analysis::spectral`]) — here
//! it is peak-picked to recover onset *positions* rather than a single aggregate.
//! BPM-grid slicing reuses stratum-dsp tempo detection via [`detect_bpm`].

use realfft::RealFftPlanner;

use crate::error::CoreError;

/// A half-open slice of a sample, in frame units: `[start_frame, end_frame)`.
#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub struct Slice {
    pub start_frame: usize,
    pub end_frame: usize,
}

impl Slice {
    /// Number of frames in this slice.
    pub fn len_frames(&self) -> usize {
        self.end_frame.saturating_sub(self.start_frame)
    }
}

/// How to chop a sample into slices.
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
pub enum ChopMethod {
    /// Slice at detected transients. `sensitivity` in `0.0..=1.0` — higher finds
    /// more (quieter) onsets, lower finds only the strongest hits.
    Transient { sensitivity: f32 },
    /// Slice into `n` equal divisions (e.g. 2, 4, 8, 16, 32).
    EqualDivisions(usize),
    /// Slice on a tempo grid. `subdivisions_per_beat` of 1 = one slice per beat,
    /// 2 = eighth notes, 4 = sixteenths.
    BpmGrid { bpm: f64, subdivisions_per_beat: u32 },
}

/// Compute slice boundaries for `samples` (interleaved, `channels`-wide) using
/// the given method. The returned slices always tile the whole sample with no
/// gaps or overlaps, the first starting at frame 0 and the last ending at the
/// final frame. Returns an empty vec for empty / zero-channel input.
pub fn compute_slices(
    samples: &[f32],
    channels: u16,
    sample_rate: u32,
    method: &ChopMethod,
) -> Result<Vec<Slice>, CoreError> {
    let ch = channels as usize;
    if ch == 0 {
        return Err(CoreError::Internal("chop: channels must be > 0".to_string()));
    }
    let total_frames = samples.len() / ch;
    if total_frames == 0 {
        return Ok(Vec::new());
    }

    let boundaries = match method {
        ChopMethod::Transient { sensitivity } => {
            let mono = mixdown_mono(samples, ch);
            let mut starts = detect_onsets(&mono, sample_rate, *sensitivity);
            // Force the first boundary to 0 so any pre-attack lead-in is folded
            // into the first slice instead of being dropped.
            if starts.first() != Some(&0) {
                starts.insert(0, 0);
            }
            starts
        }
        ChopMethod::EqualDivisions(n) => {
            let n = (*n).max(1);
            (0..n).map(|i| i * total_frames / n).collect()
        }
        ChopMethod::BpmGrid {
            bpm,
            subdivisions_per_beat,
        } => {
            if *bpm <= 0.0 {
                return Err(CoreError::Internal(
                    "chop: bpm must be positive for BPM-grid chop".to_string(),
                ));
            }
            let subdiv = (*subdivisions_per_beat).max(1) as f64;
            let frames_per_slice = (sample_rate as f64 * 60.0 / (bpm * subdiv)).round() as usize;
            if frames_per_slice == 0 {
                return Err(CoreError::Internal("chop: BPM grid step rounds to 0 frames".to_string()));
            }
            (0..total_frames).step_by(frames_per_slice).collect()
        }
    };

    Ok(boundaries_to_slices(boundaries, total_frames))
}

/// Convert sorted, deduplicated start boundaries into half-open slices covering
/// `[0, total_frames)`. Boundaries at or past the end are dropped.
fn boundaries_to_slices(mut starts: Vec<usize>, total_frames: usize) -> Vec<Slice> {
    starts.retain(|&s| s < total_frames);
    starts.sort_unstable();
    starts.dedup();
    if starts.is_empty() {
        starts.push(0);
    }
    let mut slices = Vec::with_capacity(starts.len());
    for i in 0..starts.len() {
        let start = starts[i];
        let end = starts.get(i + 1).copied().unwrap_or(total_frames);
        if end > start {
            slices.push(Slice { start_frame: start, end_frame: end });
        }
    }
    slices
}

/// Extract one slice as its own interleaved buffer.
pub fn render_slice(samples: &[f32], channels: u16, slice: &Slice) -> Vec<f32> {
    let ch = channels as usize;
    if ch == 0 {
        return Vec::new();
    }
    let total_frames = samples.len() / ch;
    let start = slice.start_frame.min(total_frames);
    let end = slice.end_frame.min(total_frames);
    if end <= start {
        return Vec::new();
    }
    samples[start * ch..end * ch].to_vec()
}

/// Detect BPM for a sample by mixing to mono and reusing stratum-dsp tempo
/// detection. Returns `None` when the tempo can't be reliably estimated.
pub fn detect_bpm(samples: &[f32], channels: u16, sample_rate: u32) -> Option<f64> {
    let ch = channels.max(1) as usize;
    let mono = mixdown_mono(samples, ch);
    crate::analysis::bpm::detect_bpm_key(&mono, sample_rate, 2.0).bpm
}

/// Average all channels of an interleaved buffer down to mono.
fn mixdown_mono(samples: &[f32], ch: usize) -> Vec<f32> {
    if ch <= 1 {
        return samples.to_vec();
    }
    let num_frames = samples.len() / ch;
    let mut mono = Vec::with_capacity(num_frames);
    for frame in 0..num_frames {
        let mut sum = 0.0f32;
        for c in 0..ch {
            sum += samples[frame * ch + c];
        }
        mono.push(sum / ch as f32);
    }
    mono
}

const WINDOW_SIZE: usize = 1024;
const HOP_SIZE: usize = 512;

/// Detect transient onset positions (in frames) via spectral-flux peak picking.
///
/// Computes a per-frame spectral flux (sum of positive magnitude increases
/// between consecutive STFT frames), then selects frames that both exceed an
/// adaptive threshold (mean + k·std, where k falls as sensitivity rises) and are
/// local maxima, enforcing a minimum gap so a single hit isn't split.
fn detect_onsets(mono: &[f32], sample_rate: u32, sensitivity: f32) -> Vec<usize> {
    if mono.len() < WINDOW_SIZE * 2 {
        // Too short for a meaningful STFT — treat as a single slice.
        return vec![0];
    }

    let mut planner = RealFftPlanner::<f32>::new();
    let fft = planner.plan_fft_forward(WINDOW_SIZE);

    let hann: Vec<f32> = (0..WINDOW_SIZE)
        .map(|i| {
            0.5 * (1.0 - (2.0 * std::f32::consts::PI * i as f32 / (WINDOW_SIZE - 1) as f32).cos())
        })
        .collect();

    // Spectral flux per hop, paired with the frame index of the *current* frame.
    let mut flux: Vec<f32> = Vec::new();
    let mut frame_positions: Vec<usize> = Vec::new();
    let mut prev_mag: Option<Vec<f32>> = None;

    let mut pos = 0;
    while pos + WINDOW_SIZE <= mono.len() {
        let mut windowed: Vec<f32> = mono[pos..pos + WINDOW_SIZE]
            .iter()
            .enumerate()
            .map(|(i, &s)| s * hann[i])
            .collect();
        let mut spectrum = fft.make_output_vec();
        if fft.process(&mut windowed, &mut spectrum).is_ok() {
            let mag: Vec<f32> = spectrum.iter().map(|c| c.norm()).collect();
            if let Some(ref prev) = prev_mag {
                let f: f32 = mag
                    .iter()
                    .zip(prev.iter())
                    .map(|(&c, &p)| (c - p).max(0.0))
                    .sum();
                flux.push(f);
                frame_positions.push(pos);
            }
            prev_mag = Some(mag);
        }
        pos += HOP_SIZE;
    }

    if flux.is_empty() {
        return vec![0];
    }

    let mean = flux.iter().sum::<f32>() / flux.len() as f32;
    let var = flux.iter().map(|&f| (f - mean) * (f - mean)).sum::<f32>() / flux.len() as f32;
    let std = var.sqrt();
    // sensitivity 0 -> mean + 1.5 std (few onsets); 1 -> mean + 0.3 std (many).
    let k = 1.5 - 1.2 * sensitivity.clamp(0.0, 1.0);
    let threshold = mean + k * std;

    // 50 ms: above a typical one-shot's attack span, so the several flux frames
    // a single hit spans collapse to one onset rather than splitting the hit.
    let min_gap_frames = (sample_rate as f32 * 0.05) as usize;
    // Absolute floor so a silent (all-zero-flux) buffer yields no onsets even
    // though its adaptive threshold is also zero.
    const FLUX_FLOOR: f32 = 1e-4;

    // Collect candidate peaks: above threshold and a strict local maximum.
    // Each candidate's boundary sits one hop *before* the frame whose flux rose,
    // so the cut lands ahead of the attack rather than a hop inside it (the flux
    // for the window at `pos` measures the energy jump from the window at
    // `pos - HOP_SIZE`, so the attack precedes `pos`).
    let mut candidates: Vec<(f32, usize)> = Vec::new();
    for i in 0..flux.len() {
        let f = flux[i];
        if f <= threshold || f < FLUX_FLOOR {
            continue;
        }
        let is_peak = (i == 0 || flux[i - 1] < f) && (i + 1 >= flux.len() || flux[i + 1] < f);
        if is_peak {
            candidates.push((f, frame_positions[i].saturating_sub(HOP_SIZE)));
        }
    }

    // Non-maximum suppression: take the strongest peaks first and drop any
    // weaker peak within `min_gap_frames` of one already kept. This keeps the
    // real (loud) transient when a soft pre-onset or flam sits beside it,
    // instead of greedily keeping whichever came first in time.
    candidates.sort_by(|a, b| b.0.partial_cmp(&a.0).unwrap_or(std::cmp::Ordering::Equal));
    let mut onsets: Vec<usize> = Vec::new();
    for (_flux, frame) in candidates {
        if onsets.iter().any(|&o| frame.abs_diff(o) < min_gap_frames) {
            continue;
        }
        onsets.push(frame);
    }
    onsets.sort_unstable();

    if onsets.is_empty() {
        onsets.push(0);
    }
    onsets
}

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

    #[test]
    fn equal_divisions_tiles_exactly() {
        // 1000 mono frames into 4 -> [0,250),[250,500),[500,750),[750,1000)
        let samples = vec![0.0f32; 1000];
        let slices = compute_slices(&samples, 1, 44100, &ChopMethod::EqualDivisions(4)).unwrap();
        assert_eq!(slices.len(), 4);
        assert_eq!(slices[0], Slice { start_frame: 0, end_frame: 250 });
        assert_eq!(slices[3], Slice { start_frame: 750, end_frame: 1000 });
        // Tiling: contiguous, no gaps, covers the whole buffer.
        assert_eq!(slices[0].start_frame, 0);
        assert_eq!(slices.last().unwrap().end_frame, 1000);
        for w in slices.windows(2) {
            assert_eq!(w[0].end_frame, w[1].start_frame);
        }
    }

    #[test]
    fn equal_divisions_stereo_frames() {
        // 8 stereo frames (16 samples) into 2 -> two 4-frame slices.
        let samples = vec![0.0f32; 16];
        let slices = compute_slices(&samples, 2, 44100, &ChopMethod::EqualDivisions(2)).unwrap();
        assert_eq!(slices.len(), 2);
        assert_eq!(slices[0], Slice { start_frame: 0, end_frame: 4 });
        assert_eq!(slices[1], Slice { start_frame: 4, end_frame: 8 });
    }

    #[test]
    fn bpm_grid_step_matches_tempo() {
        // 120 BPM at 48k = 0.5s/beat = 24000 frames/beat.
        let sample_rate = 48000;
        let samples = vec![0.0f32; 48000]; // 1 second
        let slices = compute_slices(
            &samples,
            1,
            sample_rate,
            &ChopMethod::BpmGrid { bpm: 120.0, subdivisions_per_beat: 1 },
        )
        .unwrap();
        // boundaries at 0 and 24000 -> two slices
        assert_eq!(slices.len(), 2);
        assert_eq!(slices[0], Slice { start_frame: 0, end_frame: 24000 });
        assert_eq!(slices[1], Slice { start_frame: 24000, end_frame: 48000 });
    }

    #[test]
    fn bpm_grid_subdivisions() {
        // Sixteenths at 120 BPM/48k: 24000/4 = 6000 frames per slice.
        let slices = compute_slices(
            &vec![0.0f32; 24000],
            1,
            48000,
            &ChopMethod::BpmGrid { bpm: 120.0, subdivisions_per_beat: 4 },
        )
        .unwrap();
        assert_eq!(slices.len(), 4);
        assert!(slices.iter().all(|s| s.len_frames() == 6000));
    }

    #[test]
    fn bpm_grid_rejects_nonpositive() {
        let r = compute_slices(
            &vec![0.0f32; 100],
            1,
            48000,
            &ChopMethod::BpmGrid { bpm: 0.0, subdivisions_per_beat: 1 },
        );
        assert!(r.is_err());
    }

    #[test]
    fn render_slice_extracts_stereo_range() {
        // 4 stereo frames: L/R interleaved.
        let samples = vec![1.0, -1.0, 2.0, -2.0, 3.0, -3.0, 4.0, -4.0];
        let out = render_slice(&samples, 2, &Slice { start_frame: 1, end_frame: 3 });
        assert_eq!(out, vec![2.0, -2.0, 3.0, -3.0]);
    }

    #[test]
    fn render_slice_clamps_out_of_range() {
        let samples = vec![1.0, 2.0, 3.0];
        let out = render_slice(&samples, 1, &Slice { start_frame: 2, end_frame: 99 });
        assert_eq!(out, vec![3.0]);
        let empty = render_slice(&samples, 1, &Slice { start_frame: 5, end_frame: 6 });
        assert!(empty.is_empty());
    }

    #[test]
    fn empty_input_yields_no_slices() {
        let slices = compute_slices(&[], 1, 44100, &ChopMethod::EqualDivisions(4)).unwrap();
        assert!(slices.is_empty());
    }

    #[test]
    fn zero_channels_errors() {
        assert!(compute_slices(&[0.0, 1.0], 0, 44100, &ChopMethod::EqualDivisions(2)).is_err());
    }

    #[test]
    fn transient_chop_finds_impulses() {
        // Build 4 impulses spaced 0.25s apart in a 1s 44.1k mono signal. Each
        // impulse is a short burst of full-scale noise-like content so the FFT
        // sees a broadband energy jump (a clear spectral-flux peak).
        let sr = 44100usize;
        let mut signal = vec![0.0f32; sr];
        let positions = [0usize, sr / 4, sr / 2, 3 * sr / 4];
        for &p in &positions {
            for k in 0..2000 {
                // Alternating sign = high-frequency burst, strong flux.
                let v = if k % 2 == 0 { 0.9 } else { -0.9 };
                signal[p + k] = v;
            }
        }
        let slices = compute_slices(&signal, 1, sr as u32, &ChopMethod::Transient { sensitivity: 0.5 }).unwrap();
        // Should recover roughly one slice per impulse (allow the detector some
        // slack but require it found the structure, not one big slice or dozens).
        assert!(
            (3..=6).contains(&slices.len()),
            "expected ~4 slices, got {}",
            slices.len()
        );
        // Slices tile the whole buffer.
        assert_eq!(slices[0].start_frame, 0);
        assert_eq!(slices.last().unwrap().end_frame, sr);
        for w in slices.windows(2) {
            assert_eq!(w[0].end_frame, w[1].start_frame);
        }
        // A detected boundary should land near the second impulse (sr/4).
        let near = slices
            .iter()
            .any(|s| (s.start_frame as i64 - (sr / 4) as i64).abs() < 3000);
        assert!(near, "no slice boundary near the 0.25s impulse: {slices:?}");
    }

    #[test]
    fn transient_chop_silence_is_single_slice() {
        let slices = compute_slices(&vec![0.0f32; 44100], 1, 44100, &ChopMethod::Transient { sensitivity: 0.5 }).unwrap();
        assert_eq!(slices.len(), 1);
        assert_eq!(slices[0], Slice { start_frame: 0, end_frame: 44100 });
    }
}