//! Scoreable corpus rows, and the in-memory indexes built from them.
//!
//! Extracted from [`crate::layer_eval`] when [`crate::stability`] arrived and
//! needed the same four things: read the analysed vault back as (vector, labels,
//! truth) rows projected onto a label space, split them without an RNG, and build
//! a real `ExemplarIndex` from an arbitrary subset. Same reason [`crate::labelled`]
//! exists: the step is subtle enough that two copies would drift, and the
//! subtleties are all about what is silently excluded.
//!
//! What lives here is label-space-agnostic and policy-free. The gate belongs to
//! `layer_eval`, the bar to `stability`; neither is a property of a row.

use std::collections::{BTreeSet, HashMap};

use audiofiles_core::analysis::afcl::{
    self, Afcl, AfclExemplar, AfclManifest, DEFAULT_IMPORT_WEIGHT,
};
use audiofiles_core::analysis::features::{FEATURE_VERSION, NUM_FEATURES};
use audiofiles_core::db::Database;

use crate::families::{self, LabelSpace};
use crate::labelled;

/// One corpus sample: its vector, its labels, and the class it is scored against.
pub(crate) struct Row {
    pub(crate) hash: String,
    pub(crate) vector: Vec<f64>,
    /// Labels in the evaluation's space, so the index under test carries the tags
    /// being graded.
    pub(crate) tags: Vec<String>,
    /// The single class this row is ground truth for, or `None` when the corpus
    /// gave it more than one. Content-addressed import collapses a file that
    /// appears in two class folders into one row carrying both tags; its true
    /// class is undecidable, so it trains but is never tested.
    pub(crate) truth: Option<String>,
    /// The corpus folder(s) this row came from, before projection. Kept so a
    /// family's members can be broken out by where they came from: `low` is kick
    /// plus tom, and whether the layer recovers a tom as readily as a kick is the
    /// open split `af-coarse-families` asks about.
    pub(crate) origin: String,
    /// The sample's filename, as imported. Only the filename rules read it, and
    /// only to answer whether this sample is in the population the layer exists
    /// to serve: `starter_rules` already labels 97.7% of this corpus off its
    /// name, and every accuracy number so far was measured on exactly that
    /// population. See [`crate::stability`].
    pub(crate) name: String,
}

/// What [`load_rows`] set aside, so no exclusion is silent.
pub(crate) struct Dropped {
    /// Rows whose every corpus tag projects nowhere in this label space.
    pub(crate) unprojectable: usize,
    /// Corpus labels those rows came from.
    pub(crate) origins: BTreeSet<String>,
}

/// Read the analysed corpus back out of the scratch vault as scoreable rows,
/// projected onto `space`.
///
/// The projection happens here rather than in [`labelled`] on purpose. The vault
/// is ground truth at the finest resolution the corpus carries, which is what
/// `afcl_gen` exports from; grading at a coarser resolution is a property of the
/// evaluation, not of the corpus, so it belongs on the read side. That also keeps
/// the exported layer byte-identical and leaves the retired instrument question
/// runnable instead of deleted.
pub(crate) fn load_rows(db: &Database, space: LabelSpace) -> Result<(Vec<Row>, Dropped), String> {
    let conn = db.conn();

    let mut tags_by_hash: HashMap<String, Vec<String>> = HashMap::new();
    {
        let mut stmt = conn
            .prepare("SELECT sample_hash, tag FROM tags")
            .map_err(|e| format!("tags query: {e}"))?;
        let rows = stmt
            .query_map([], |row| {
                Ok((row.get::<_, String>(0)?, row.get::<_, String>(1)?))
            })
            .map_err(|e| format!("tags query: {e}"))?;
        for r in rows {
            let (hash, tag) = r.map_err(|e| format!("tags row: {e}"))?;
            tags_by_hash.entry(hash).or_default().push(tag);
        }
    }

    // Names are a left join in spirit: a row with no `samples` entry is one this
    // module fabricated in a test, and it still has to be scoreable.
    let mut name_by_hash: HashMap<String, String> = HashMap::new();
    {
        let mut stmt = conn
            .prepare("SELECT hash, original_name FROM samples")
            .map_err(|e| format!("names query: {e}"))?;
        let rows = stmt
            .query_map([], |row| {
                Ok((row.get::<_, String>(0)?, row.get::<_, String>(1)?))
            })
            .map_err(|e| format!("names query: {e}"))?;
        for r in rows {
            let (hash, name) = r.map_err(|e| format!("names row: {e}"))?;
            name_by_hash.insert(hash, name);
        }
    }

    let mut out = Vec::new();
    let mut dropped = Dropped {
        unprojectable: 0,
        origins: BTreeSet::new(),
    };
    let mut stmt = conn
        .prepare("SELECT hash, vector FROM sample_features WHERE feat_version = ?1 ORDER BY hash")
        .map_err(|e| format!("features query: {e}"))?;
    let rows = stmt
        .query_map([FEATURE_VERSION], |row| {
            Ok((row.get::<_, String>(0)?, row.get::<_, String>(1)?))
        })
        .map_err(|e| format!("features query: {e}"))?;
    for r in rows {
        let (hash, json) = r.map_err(|e| format!("features row: {e}"))?;
        let Some(mut corpus_tags) = tags_by_hash.remove(&hash) else {
            continue;
        };
        corpus_tags.sort();
        corpus_tags.dedup();
        let vector: Vec<f64> =
            serde_json::from_str(&json).map_err(|e| format!("vector for {hash}: {e}"))?;
        // The index drops these too (`is_usable_vector`), so counting them as
        // testable would score the layer on samples it never sees.
        if vector.len() != NUM_FEATURES || !vector.iter().all(|x| x.is_finite()) {
            continue;
        }

        // Project, then dedup again: two instrument tags landing in one family is
        // not ambiguity, it is the coarser question being easier. A file in both
        // the kick and tom folders has no instrument truth and a perfectly good
        // family one.
        let mut tags: Vec<String> = corpus_tags
            .iter()
            .filter_map(|t| families::project(space, t))
            .map(str::to_string)
            .collect();
        tags.sort();
        tags.dedup();
        if tags.is_empty() {
            dropped.unprojectable += 1;
            dropped.origins.extend(
                corpus_tags
                    .iter()
                    .map(|t| labelled::label_for_tag(t).to_string()),
            );
            continue;
        }

        let truth = (tags.len() == 1).then(|| tags[0].clone());
        let origin = corpus_tags
            .iter()
            .map(|t| labelled::label_for_tag(t))
            .collect::<Vec<_>>()
            .join("+");
        let name = name_by_hash.get(&hash).cloned().unwrap_or_default();
        out.push(Row {
            hash,
            vector,
            tags,
            truth,
            origin,
            name,
        });
    }
    Ok((out, dropped))
}

/// Assign each row a fold, stratified by class.
///
/// Round-robin down each class's hash-sorted list rather than a shuffle: the
/// content hash is already an arbitrary order with respect to the source pack a
/// file came from, so this needs no RNG and two runs over the same corpus produce
/// the same folds. Ambiguous rows get no fold; they train everywhere.
pub(crate) fn assign_folds(rows: &[Row], folds: usize) -> Vec<Option<usize>> {
    let mut seen_per_class: HashMap<&str, usize> = HashMap::new();
    rows.iter()
        .map(|r| {
            let truth = r.truth.as_deref()?;
            let n = seen_per_class.entry(truth).or_default();
            let fold = *n % folds;
            *n += 1;
            Some(fold)
        })
        .collect()
}

/// Build an in-memory vault holding just `local` rows, so the index under test is
/// built by the same `build_index` the app calls.
///
/// Every exemplar here is `local` (weight 1.0) where the shipped ones arrive from
/// an imported layer at a lower weight. That does not change a score: the weight
/// is a constant multiplier inside a sum that is then divided by its own total,
/// so a uniform weight cancels. It matters only against a mixed index of user
/// labels plus the layer, which is [`mixed_db`].
pub(crate) fn local_db(local: &[&Row]) -> Result<Database, String> {
    mixed_db(local, &[])
}

/// Build an in-memory vault holding `local` rows at weight 1.0 and `imported`
/// rows as an enabled `.afcl` layer at [`DEFAULT_IMPORT_WEIGHT`].
///
/// This is the deployment shape, and it is not a uniform index: the weight only
/// cancels when every exemplar carries the same one. Here the ratio is real, and
/// so is the second-order effect that makes this worth simulating rather than
/// reasoning about — `build_index` fits the standardization means and standard
/// deviations over local **and** imported exemplars together, so adding a user's
/// labels moves the space every distance is measured in.
///
/// The layer goes in through `afcl::import` rather than a hand-written INSERT so
/// the weight, the enabled flag and the `feat_version` gate are the ones the app
/// applies. Imported exemplars carry vector + tags and no hash, exactly as a
/// shipped layer does.
pub(crate) fn mixed_db(local: &[&Row], imported: &[&Row]) -> Result<Database, String> {
    let db = Database::open_in_memory().map_err(|e| format!("open_in_memory: {e}"))?;
    db.transaction(|_tx| {
        for row in local {
            db.conn().execute(
                "INSERT INTO samples (hash, original_name, file_extension, file_size, import_date, last_modified)
                 VALUES (?1, ?2, 'wav', 1, 0, 0)",
                rusqlite::params![&row.hash, &row.name],
            )?;
            let json = serde_json::to_string(&row.vector).unwrap_or_default();
            db.conn().execute(
                "INSERT INTO sample_features (hash, feat_version, vector, computed_at) VALUES (?1, ?2, ?3, 0)",
                rusqlite::params![&row.hash, FEATURE_VERSION, json],
            )?;
        }
        Ok(())
    })
    .map_err(|e| format!("seeding the index: {e}"))?;
    // Tags outside the transaction: `add_tag` is the public path and manages its
    // own writes, and this is an in-memory DB where the fsync cost it avoids does
    // not exist.
    for row in local {
        for tag in &row.tags {
            audiofiles_core::tags::add_tag(&db, &row.hash, tag)
                .map_err(|e| format!("tagging the index: {e}"))?;
        }
    }

    if !imported.is_empty() {
        let exemplars: Vec<AfclExemplar> = imported
            .iter()
            .map(|r| AfclExemplar {
                vector: r.vector.clone(),
                tags: r.tags.clone(),
            })
            .collect();
        let layer = Afcl {
            manifest: AfclManifest {
                afcl_version: afcl::AFCL_VERSION,
                feat_version: FEATURE_VERSION,
                name: "simulated imported layer".into(),
                description: "built by audiofiles-bench, never written to disk".into(),
                kind: afcl::LayerKind::Official.as_str().to_string(),
                created_at: 0,
                license_note: String::new(),
                exemplar_count: exemplars.len(),
                rule_count: 0,
                policy_count: 0,
            },
            exemplars,
            rules: Vec::new(),
            policy: Vec::new(),
        };
        let summary = afcl::import(&db, &layer, Some(afcl::BUNDLED_SOURCE))
            .map_err(|e| format!("importing the simulated layer: {e}"))?;
        // An exemplar dropped at import is one the measurement thinks it has and
        // does not. `load_rows` already applied the same finiteness filter, so a
        // mismatch here means the two disagree and the run is not what it says.
        if summary.exemplars != imported.len() {
            return Err(format!(
                "the simulated layer imported {} of {} exemplars",
                summary.exemplars,
                imported.len()
            ));
        }
    }
    Ok(db)
}

/// The weight an imported layer carries beside the user's own labels, restated
/// here so a report can print the ratio it was measured at.
pub(crate) const IMPORT_WEIGHT: f64 = DEFAULT_IMPORT_WEIGHT;

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

    fn row(hash: &str, tag: &str) -> Row {
        Row {
            hash: hash.into(),
            vector: vec![0.5; NUM_FEATURES],
            tags: vec![tag.into()],
            truth: Some(tag.into()),
            origin: labelled::label_for_tag(tag).into(),
            name: format!("{hash}.wav"),
        }
    }

    #[test]
    fn folds_are_stratified_and_ambiguous_rows_get_none() {
        let mut rows: Vec<Row> = (0..10)
            .map(|i| row(&format!("k{i}"), "instrument.drum.kick"))
            .collect();
        rows.extend((0..4).map(|i| row(&format!("s{i}"), "instrument.drum.snare")));
        rows.push(Row {
            hash: "dup".into(),
            vector: vec![0.0; NUM_FEATURES],
            tags: Vec::new(),
            truth: None,
            origin: String::new(),
            name: "dup.wav".into(),
        });

        let folds = assign_folds(&rows, 5);
        assert_eq!(
            folds.last().copied().flatten(),
            None,
            "ambiguous row trains everywhere"
        );

        // Every fold holds two of the ten kicks: stratification, not a global
        // round-robin that would leave a fold without any snares.
        for f in 0..5 {
            let kicks = rows
                .iter()
                .zip(&folds)
                .filter(|(r, fold)| {
                    **fold == Some(f) && r.truth.as_deref() == Some("instrument.drum.kick")
                })
                .count();
            assert_eq!(kicks, 2, "fold {f}");
        }
        // Four snares over five folds: one fold is short, and none holds two.
        for f in 0..5 {
            let snares = rows
                .iter()
                .zip(&folds)
                .filter(|(r, fold)| {
                    **fold == Some(f) && r.truth.as_deref() == Some("instrument.drum.snare")
                })
                .count();
            assert!(snares <= 1, "fold {f} holds {snares} snares");
        }
    }

    #[test]
    fn folds_are_deterministic() {
        let rows: Vec<Row> = (0..20)
            .map(|i| row(&format!("k{i}"), "instrument.drum.kick"))
            .collect();
        assert_eq!(assign_folds(&rows, 5), assign_folds(&rows, 5));
    }

    #[test]
    fn a_fold_index_scores_a_held_out_sample_from_its_neighbours() {
        // End to end over the real index: two tight clusters, one held out
        // sample, and the index must answer with the cluster it sits in.
        let mut rows = Vec::new();
        for i in 0..5 {
            let mut r = row(&format!("k{i}"), "instrument.drum.kick");
            r.vector = vec![0.01 * f64::from(i); NUM_FEATURES];
            rows.push(r);
        }
        for i in 0..5 {
            let mut r = row(&format!("s{i}"), "instrument.drum.snare");
            r.vector = vec![100.0 + f64::from(i); NUM_FEATURES];
            rows.push(r);
        }
        let refs: Vec<&Row> = rows.iter().collect();
        let db = local_db(&refs).unwrap();
        let index = audiofiles_core::analysis::exemplar::build_index(&db).unwrap();
        assert_eq!(index.len(), 10);

        let query = vec![0.02; NUM_FEATURES];
        let scored = index.score(&query, audiofiles_core::analysis::exemplar::DEFAULT_K, None);
        assert_eq!(
            scored.first().map(|s| s.tag.as_str()),
            Some("instrument.drum.kick")
        );
    }

    #[test]
    fn load_rows_marks_a_multi_tagged_sample_ambiguous() {
        let rows = [
            row("a", "instrument.drum.kick"),
            Row {
                hash: "b".into(),
                vector: vec![0.5; NUM_FEATURES],
                tags: vec![
                    "instrument.drum.kick".into(),
                    "instrument.drum.snare".into(),
                ],
                truth: None,
                origin: "kick+snare".into(),
                name: "b.wav".into(),
            },
        ];
        let refs: Vec<&Row> = rows.iter().collect();
        let db = local_db(&refs).unwrap();

        let (read, dropped) = load_rows(&db, LabelSpace::Instrument).unwrap();
        assert_eq!(read.len(), 2);
        assert_eq!(dropped.unprojectable, 0);
        let b = read.iter().find(|r| r.hash == "b").unwrap();
        assert_eq!(b.truth, None, "two class tags means no ground truth");
        assert_eq!(b.tags.len(), 2, "but it still trains with both labels");
    }

    #[test]
    fn the_family_projection_resolves_an_instrument_level_ambiguity() {
        // A file in both the kick and tom folders has no instrument truth and a
        // perfectly good family one: the coarser question is the easier question,
        // and that has to show up as a testable row rather than a dropped one.
        let rows = [Row {
            hash: "a".into(),
            vector: vec![0.5; NUM_FEATURES],
            tags: vec!["instrument.drum.kick".into(), "instrument.drum.tom".into()],
            truth: None,
            origin: "kick+tom".into(),
            name: "a.wav".into(),
        }];
        let refs: Vec<&Row> = rows.iter().collect();
        let db = local_db(&refs).unwrap();

        let (read, _) = load_rows(&db, LabelSpace::Family).unwrap();
        assert_eq!(read[0].truth.as_deref(), Some("family.low"));
        assert_eq!(read[0].origin, "kick+tom");
    }

    #[test]
    fn percussion_leaves_the_family_run_entirely() {
        // Not merely untested: an exemplar carrying a guessed family label would
        // pull the index toward that guess for every other sample too.
        let rows = [Row {
            hash: "p".into(),
            vector: vec![0.5; NUM_FEATURES],
            tags: vec!["instrument.percussion".into()],
            truth: Some("instrument.percussion".into()),
            origin: "percussion".into(),
            name: "p.wav".into(),
        }];
        let refs: Vec<&Row> = rows.iter().collect();
        let db = local_db(&refs).unwrap();

        let (read, dropped) = load_rows(&db, LabelSpace::Family).unwrap();
        assert!(read.is_empty(), "it must not train either");
        assert_eq!(dropped.unprojectable, 1);
        assert!(dropped.origins.contains("percussion"));
    }

    #[test]
    fn load_rows_carries_the_filename_back() {
        // The rule-miss population is keyed off this, and a silently empty name
        // would read as "no filename rule fires" for the whole corpus, which is
        // the opposite of the truth.
        let rows = [row("a", "instrument.drum.kick")];
        let refs: Vec<&Row> = rows.iter().collect();
        let db = local_db(&refs).unwrap();
        let (read, _) = load_rows(&db, LabelSpace::Instrument).unwrap();
        assert_eq!(read[0].name, "a.wav");
    }

    #[test]
    fn a_mixed_index_holds_both_populations() {
        let local: Vec<Row> = (0..3)
            .map(|i| row(&format!("l{i}"), "instrument.drum.kick"))
            .collect();
        let imported: Vec<Row> = (0..4)
            .map(|i| row(&format!("i{i}"), "instrument.drum.snare"))
            .collect();
        let l: Vec<&Row> = local.iter().collect();
        let i: Vec<&Row> = imported.iter().collect();
        let db = mixed_db(&l, &i).unwrap();

        let index = audiofiles_core::analysis::exemplar::build_index(&db).unwrap();
        assert_eq!(index.len(), 7, "local and imported both reach the index");

        let layers = afcl::list_layers(&db).unwrap();
        assert_eq!(layers.len(), 1);
        assert!(layers[0].enabled, "a disabled layer contributes nothing");
        assert!((layers[0].weight - IMPORT_WEIGHT).abs() < 1e-9);
    }

    #[test]
    fn an_imported_exemplar_outvoted_by_local_ones_shows_the_weight_is_live() {
        // The property `mixed_db` exists for. Two exemplars equidistant from the
        // query, one local and one imported: the local tag must score 1/(1+0.5)
        // of the neighbourhood rather than half of it. If the weight were being
        // cancelled away this would come back 0.5.
        let mut near_local = row("l0", "instrument.drum.kick");
        near_local.vector = vec![0.0; NUM_FEATURES];
        let mut near_imported = row("i0", "instrument.drum.snare");
        near_imported.vector = vec![0.0; NUM_FEATURES];

        let db = mixed_db(&[&near_local], &[&near_imported]).unwrap();
        let index = audiofiles_core::analysis::exemplar::build_index(&db).unwrap();
        let scored = index.score(&vec![0.0; NUM_FEATURES], 15, None);

        let kick = scored
            .iter()
            .find(|s| s.tag == "instrument.drum.kick")
            .expect("the local exemplar is in the neighbourhood");
        assert!(
            (kick.score - (1.0 / (1.0 + IMPORT_WEIGHT))).abs() < 1e-9,
            "local should carry {:.3} of the weight, carried {:.3}",
            1.0 / (1.0 + IMPORT_WEIGHT),
            kick.score
        );
    }
}