//! BPM and key accuracy against FSL10K ground truth.
//!
//! The pipeline emits two things a corpus can check it on, and this is where both
//! are scored: `bpm::detect_bpm_key`'s tempo and key. (The analysis bench used to
//! carry a third, accuracy of the single sample-class label, until that label was
//! retired; see `docs/ml_classifier.md`.)
//!
//! Ground truth comes from the Freesound Loop Dataset, which is the only corpus
//! here carrying expert tempo and key annotations. Zenodo publishes FSL10K as
//! CC-BY, but its sounds carry four different Freesound licences and 1,436 of
//! the 9,493 are CC-BY-NC or Sampling+. That is fine for what happens here --
//! this mode decodes audio, scores a number, and emits nothing derived from the
//! samples -- and it is not fine for anything that ships. Anything that selects
//! sounds out of this corpus rather than measuring over it reads the per-sound
//! licence from `metadata.json`; `scripts/corpus.py` records the mix in the
//! corpus MANIFEST.json.
//!
//! Fetch it with `scripts/corpus.py --datasets fsl10k`, then point this at the
//! extracted root:
//!
//!     AF_BENCH_FSL10K=/media/max/T9/af-corpus/_raw/fsl10k \
//!       cargo run --release -p audiofiles-bench -- accuracy
//!
//! Tempo is scored with the standard MIREX pair rather than a single number,
//! because a tempo estimator that is "wrong" is usually wrong by a factor of
//! two, and collapsing that into one accuracy figure hides which kind of wrong
//! it is:
//!   Acc1 = within 4% of ground truth.
//!   Acc2 = Acc1, or within 4% of truth scaled by 1/3, 1/2, 2, or 3.
//! A large Acc2-minus-Acc1 gap means octave errors, which are a different fix
//! from being diffusely inaccurate.

use std::collections::HashMap;
use std::path::{Path, PathBuf};

use audiofiles_core::analysis::{bpm, decode};
use rayon::prelude::*;

use crate::storage;

/// Ground truth for one sound, after consensus across annotators.
struct Truth {
    bpm: f64,
    /// Pitch class 0..=11 with mode, or None when the annotators marked the
    /// loop as having no key (drum loops, percussion, fx).
    key: Option<(u8, Mode)>,
}

#[derive(Clone, Copy, PartialEq, Eq, Debug)]
enum Mode {
    Major,
    Minor,
}

/// Parse a note name to a pitch class.
///
/// Comparing pitch classes rather than strings sidesteps enharmonic spelling
/// entirely: the annotations use sharps only, but nothing guarantees the
/// detector does, and "A#" and "Bb" are the same key.
fn pitch_class(note: &str) -> Option<u8> {
    let s = note.trim().to_lowercase();
    let mut chars = s.chars();
    let base = match chars.next()? {
        'c' => 0,
        'd' => 2,
        'e' => 4,
        'f' => 5,
        'g' => 7,
        'a' => 9,
        'b' => 11,
        _ => return None,
    };
    let accidental: i32 = match chars.next() {
        Some('#' | 's') => 1,
        Some('b' | 'f') => -1,
        None => 0,
        _ => return None,
    };
    Some(((base + accidental).rem_euclid(12)) as u8)
}

/// Parse a detected key string into a pitch class and mode.
///
/// Handles both shapes on purpose. `detect_bpm_key` now normalises to
/// "A minor" / "C major", but before that it passed through
/// `stratum_dsp::Key::name()` unchanged, which is "C" / "F#" for major and
/// "Am" / "C#m" for minor. Vaults analysed before migration 034 can still hold
/// the compact form, and accepting only one spelling silently scores 0%
/// forever rather than failing loudly, so both are parsed.
fn parse_app_key(s: &str) -> Option<(u8, Mode)> {
    let s = s.trim();
    if let Some((note, word)) = s.rsplit_once(' ') {
        let mode = match word.trim().to_lowercase().as_str() {
            "minor" | "min" => Mode::Minor,
            "major" | "maj" => Mode::Major,
            _ => return None,
        };
        return Some((pitch_class(note)?, mode));
    }
    // Bare form: trailing 'm' means minor, otherwise major. Checked after the
    // spaced form so "F major" is not read as a note named "F majo" + m.
    if let Some(note) = s.strip_suffix('m') {
        return Some((pitch_class(note)?, Mode::Minor));
    }
    Some((pitch_class(s)?, Mode::Major))
}

/// Load and reconcile annotations.
///
/// Half the annotated sounds in FSL10K carry more than one annotator's
/// judgement, and annotators disagree. Rather than picking one arbitrarily,
/// disagreements are dropped and counted: scoring a detector against a
/// contested label measures the disagreement, not the detector.
fn load_truth(annotations: &Path) -> (HashMap<String, Truth>, usize, usize) {
    let mut by_sound: HashMap<String, Vec<serde_json::Value>> = HashMap::new();
    let mut discarded = 0usize;

    let mut stack = vec![annotations.to_path_buf()];
    while let Some(dir) = stack.pop() {
        let Ok(entries) = std::fs::read_dir(&dir) else {
            continue;
        };
        for entry in entries.flatten() {
            let path = entry.path();
            if path.is_dir() {
                stack.push(path);
                continue;
            }
            let Some(name) = path.file_name().and_then(|n| n.to_str()) else {
                continue;
            };
            let Some(id) = name
                .strip_prefix("sound-")
                .and_then(|n| n.strip_suffix(".json"))
            else {
                continue;
            };
            let Ok(text) = std::fs::read_to_string(&path) else {
                continue;
            };
            let Ok(v) = serde_json::from_str::<serde_json::Value>(&text) else {
                continue;
            };
            if v.get("discard").and_then(serde_json::Value::as_bool) == Some(true) {
                discarded += 1;
                continue;
            }
            by_sound.entry(id.to_string()).or_default().push(v);
        }
    }

    let mut truth = HashMap::new();
    let mut disputed = 0usize;

    for (id, annots) in by_sound {
        // BPM is stored as a string in these files.
        let bpms: Vec<f64> = annots
            .iter()
            .filter_map(|a| a.get("bpm").and_then(serde_json::Value::as_str))
            .filter_map(|s| s.trim().parse::<f64>().ok())
            .filter(|b| *b > 0.0)
            .collect();
        if bpms.is_empty() {
            continue;
        }
        // Annotators agree if they are all within 4% of the first, the same
        // tolerance the scoring uses.
        let first = bpms[0];
        if bpms.iter().any(|b| (b - first).abs() / first > 0.04) {
            disputed += 1;
            continue;
        }
        let bpm_truth = bpms.iter().sum::<f64>() / bpms.len() as f64;

        let keys: Vec<Option<(u8, Mode)>> = annots
            .iter()
            .map(|a| {
                let k = a.get("key").and_then(serde_json::Value::as_str)?;
                let m = a.get("mode").and_then(serde_json::Value::as_str)?;
                let mode = match m {
                    "min" => Mode::Minor,
                    "maj" => Mode::Major,
                    _ => return None,
                };
                Some((pitch_class(k)?, mode))
            })
            .collect();
        // "none"/"unknown" parse to None, which is itself a meaningful answer
        // (the loop has no key), so only disagreement between Some values is a
        // dispute.
        let key_truth = if keys.iter().all(|k| *k == keys[0]) {
            keys[0]
        } else {
            None
        };

        truth.insert(
            id,
            Truth {
                bpm: bpm_truth,
                key: key_truth,
            },
        );
    }

    (truth, discarded, disputed)
}

/// Index audio files by Freesound sound ID (the filename stem).
fn index_audio(root: &Path) -> HashMap<String, PathBuf> {
    let mut out = HashMap::new();
    let mut stack = vec![root.to_path_buf()];
    while let Some(dir) = stack.pop() {
        let Ok(entries) = std::fs::read_dir(&dir) else {
            continue;
        };
        for entry in entries.flatten() {
            let path = entry.path();
            if path.is_dir() {
                stack.push(path);
                continue;
            }
            let is_audio = path.extension().and_then(|e| e.to_str()).is_some_and(|e| {
                matches!(
                    e.to_lowercase().as_str(),
                    "wav" | "flac" | "mp3" | "ogg" | "aif" | "aiff"
                )
            });
            if !is_audio {
                continue;
            }
            if let Some(stem) = path.file_stem().and_then(|s| s.to_str()) {
                // FSL10K names audio by sound ID, sometimes with a suffix after
                // the ID; take the leading digit run.
                let id: String = stem.chars().take_while(char::is_ascii_digit).collect();
                if !id.is_empty() {
                    out.insert(id, path);
                }
            }
        }
    }
    out
}

fn within(a: f64, b: f64, tol: f64) -> bool {
    b > 0.0 && (a - b).abs() / b <= tol
}

struct Scored {
    bpm_acc1: bool,
    bpm_acc2: bool,
    /// Set only when both truth and detection carry a key.
    key_exact: Option<bool>,
    key_related: Option<bool>,
    /// Tracked separately so a zero key score can be attributed. "No key was
    /// scored" has three very different causes -- the loop has no key in the
    /// ground truth, the detector returned nothing, or the detected string
    /// failed to parse -- and collapsing them hides a parser bug as a result.
    had_truth_key: bool,
    had_detected_key: bool,
    detected_key_unparsed: bool,
}

fn score(detected_bpm: Option<f64>, detected_key: Option<&str>, truth: &Truth) -> Scored {
    let (acc1, acc2) = match detected_bpm {
        Some(d) => {
            let a1 = within(d, truth.bpm, 0.04);
            // Octave and triplet confusions: the classic failure mode of
            // autocorrelation tempo estimators.
            let a2 = a1
                || [1.0 / 3.0, 0.5, 2.0, 3.0]
                    .iter()
                    .any(|f| within(d, truth.bpm * f, 0.04));
            (a1, a2)
        }
        None => (false, false),
    };

    let (key_exact, key_related) = match (truth.key, detected_key.and_then(parse_app_key)) {
        (Some((tp, tm)), Some((dp, dm))) => {
            let exact = tp == dp && tm == dm;
            // Related = relative major/minor, or a perfect fifth apart in the
            // same mode. These are the confusions a chroma-based estimator makes
            // because the pitch content genuinely overlaps.
            let relative = match tm {
                Mode::Major => dm == Mode::Minor && dp == (tp + 9) % 12,
                Mode::Minor => dm == Mode::Major && dp == (tp + 3) % 12,
            };
            let fifth = tm == dm && (dp == (tp + 7) % 12 || dp == (tp + 5) % 12);
            (Some(exact), Some(exact || relative || fifth))
        }
        _ => (None, None),
    };

    Scored {
        bpm_acc1: acc1,
        bpm_acc2: acc2,
        key_exact,
        key_related,
        had_truth_key: truth.key.is_some(),
        had_detected_key: detected_key.is_some(),
        detected_key_unparsed: detected_key.is_some_and(|k| parse_app_key(k).is_none()),
    }
}

pub(crate) fn run(root: &Path, limit: Option<usize>) {
    println!("━━━ BPM + KEY ACCURACY (FSL10K ground truth) ━━━");
    println!();

    // The zip lays out FSL10K/audio/...; accept either the zip root or the
    // directory that directly contains audio/ and annotations/.
    let annotations = [
        "annotations",
        "FSL10K/annotations",
        "../fsl10k-annotations/annotations",
    ]
    .iter()
    .map(|p| root.join(p))
    .find(|p| p.is_dir());
    let audio = ["FSL10K/audio", "audio", "FSL10K/FSL10K/audio"]
        .iter()
        .map(|p| root.join(p))
        .find(|p| p.is_dir());

    let (Some(annotations), Some(audio)) = (annotations, audio) else {
        eprintln!(
            "could not find annotations/ and audio/ under {}",
            root.display()
        );
        eprintln!("fetch with: ./scripts/corpus.py --datasets fsl10k");
        return;
    };

    // Printed but not serialised: this mode emits no timings, so the drive
    // cannot change its results. It is recorded for provenance, so a scorecard
    // says which copy of the dataset produced the numbers.
    let audio_storage = storage::describe(&audio);
    storage::print_conditions(&[("audio", &audio_storage)], None);

    let (truth, discarded, disputed) = load_truth(&annotations);
    let index = index_audio(&audio);
    println!(
        "  annotations: {} sounds with agreed ground truth",
        truth.len()
    );
    println!(
        "  {discarded} annotation(s) flagged discard, {disputed} sound(s) dropped as disputed"
    );
    println!("  audio files indexed: {}", index.len());

    let mut pairs: Vec<(&String, &Truth, &PathBuf)> = truth
        .iter()
        .filter_map(|(id, t)| index.get(id).map(|p| (id, t, p)))
        .collect();
    pairs.sort_by_key(|(id, _, _)| (*id).clone());
    if let Some(lim) = limit {
        pairs.truncate(lim);
    }

    if pairs.is_empty() {
        eprintln!("no annotated sound matched an audio file");
        return;
    }
    println!("  scoring {} loops", pairs.len());
    println!();

    let results: Vec<Scored> = pairs
        .par_iter()
        .filter_map(|(_, t, path)| {
            let decoded = decode::decode_to_mono(path).ok()?;
            // 2.0 matches the analysis pipeline's min_duration gate, so this
            // measures what the app actually runs rather than a variant.
            let r = bpm::detect_bpm_key(&decoded.samples, decoded.sample_rate, 2.0);
            Some(score(r.bpm, r.key.as_deref(), t))
        })
        .collect();

    let n = results.len() as f64;
    let acc1 = results.iter().filter(|r| r.bpm_acc1).count() as f64 / n * 100.0;
    let acc2 = results.iter().filter(|r| r.bpm_acc2).count() as f64 / n * 100.0;

    println!("  TEMPO ({} scored)", results.len());
    println!("    Acc1 (within 4%)            {acc1:>6.1}%");
    println!("    Acc2 (octave/triplet ok)    {acc2:>6.1}%");
    let gap = acc2 - acc1;
    println!("    octave-error gap            {gap:>6.1}%");
    if gap > 15.0 {
        println!("    ^ large gap: most misses are half/double time, not noise");
    }
    println!();

    let with_truth = results.iter().filter(|r| r.had_truth_key).count();
    let with_detected = results.iter().filter(|r| r.had_detected_key).count();
    let unparsed = results.iter().filter(|r| r.detected_key_unparsed).count();

    let keyed: Vec<&Scored> = results.iter().filter(|r| r.key_exact.is_some()).collect();
    println!("  KEY coverage");
    println!(
        "    loops with ground-truth key   {with_truth:>5} / {}",
        results.len()
    );
    println!(
        "    loops with detected key       {with_detected:>5} / {}",
        results.len()
    );
    if unparsed > 0 {
        println!("    detected but UNPARSED         {unparsed:>5}  <- parser/format mismatch");
    }
    println!("    scorable (both present)       {:>5}", keyed.len());

    // Precision on keyless material, which the accuracy figures cannot show.
    // Most loops in a sample library are drums, and annotators mark those as
    // having no key. A detector that emits a key anyway is not inaccurate by
    // the measures above -- those only score loops that have a key -- but it
    // fills the library with confident wrong labels, and a key filter that
    // returns drum loops is a user-visible bug.
    let keyless = results.len() - with_truth;
    let spurious = results
        .iter()
        .filter(|r| !r.had_truth_key && r.had_detected_key)
        .count();
    if keyless > 0 {
        let rate = spurious as f64 / keyless as f64 * 100.0;
        println!("    key emitted on keyless loops  {spurious:>5} / {keyless}  ({rate:.1}%)");
    }
    println!();

    if keyed.is_empty() {
        println!("  KEY: nothing scorable. Check the coverage lines above for why.");
    } else {
        let kn = keyed.len() as f64;
        let exact = keyed.iter().filter(|r| r.key_exact == Some(true)).count() as f64 / kn * 100.0;
        let related =
            keyed.iter().filter(|r| r.key_related == Some(true)).count() as f64 / kn * 100.0;
        println!("  KEY ({} scored, of {} total)", keyed.len(), results.len());
        println!("    exact                       {exact:>6.1}%");
        println!("    exact or relative/fifth     {related:>6.1}%");
        // Chance is 1/24 for exact over 12 pitch classes x 2 modes.
        println!("    (chance for exact is 4.2%)");
    }
}

#[cfg(test)]
mod tests {
    use super::{Mode, parse_app_key, pitch_class};

    #[test]
    fn parses_stratum_key_format() {
        // The format detect_bpm_key actually emits.
        assert_eq!(parse_app_key("C"), Some((0, Mode::Major)));
        assert_eq!(parse_app_key("F#"), Some((6, Mode::Major)));
        assert_eq!(parse_app_key("Am"), Some((9, Mode::Minor)));
        assert_eq!(parse_app_key("C#m"), Some((1, Mode::Minor)));
    }

    #[test]
    fn parses_documented_key_format() {
        assert_eq!(parse_app_key("A minor"), Some((9, Mode::Minor)));
        assert_eq!(parse_app_key("C major"), Some((0, Mode::Major)));
    }

    #[test]
    fn enharmonics_are_the_same_pitch_class() {
        assert_eq!(pitch_class("a#"), pitch_class("bb"));
        assert_eq!(pitch_class("c#"), pitch_class("db"));
    }
}