//! Optional trained head (Phase 6 of the hybrid classifier).
//!
//! For large libraries, brute-force k-NN ([`super::exemplar`]) costs O(N·d) per query.
//! This module *distills* the same exemplar store into a compact, per-library
//! **one-vs-rest logistic-regression** model: per tag, a [`NUM_FEATURES`]-weight vector +
//! bias scoring the standardized feature vector through a sigmoid. Inference is then
//! O(tags·d), independent of library size, and the whole model is a few KB.
//!
//! The head is a **local, regenerable cache** (the `trained_head` table, not synced):
//! it rebuilds from `sample_features` + `tags` on any device, so there is nothing to sync
//! and nothing data-derived to ship. It is *optional*, the k-NN layer stays the source of
//! truth and the only path offering neighbor explainability; the head trades that for speed
//! on big libraries. Training is fully deterministic (zero init, fixed-iteration full-batch
//! gradient descent, no RNG), so two devices distilling the same exemplars agree.

use std::collections::HashMap;

use crate::db::Database;
use crate::error::{CoreError, Result, unix_now};
use tracing::instrument;

use super::classify::{FEATURE_VERSION, NUM_FEATURES};
use super::exemplar::{
    self, MlOutcome, TagScore, apply_policy, load_query_vector, sample_tag_set, standardize_vec,
};

/// A tag needs at least this many positive exemplars to earn a distilled model; rarer
/// tags stay with k-NN (too few points to fit a stable linear boundary).
pub const MIN_POSITIVE_EXAMPLES: usize = 5;

/// Below this many labeled exemplars, brute-force k-NN is already fast; the head is
/// recommended only once the library grows past it. Advisory, callers may train anyway.
pub const RECOMMENDED_MIN_EXEMPLARS: usize = 2_000;

// Training hyperparameters, deterministic full-batch gradient descent.
const ITERATIONS: usize = 300;
const LEARNING_RATE: f64 = 0.5;
const L2: f64 = 1e-3;

/// A per-tag linear scorer: `sigmoid(weights · x_std + bias)`.
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct TagModel {
    pub tag: String,
    pub weights: Vec<f64>,
    pub bias: f64,
}

/// A compact, per-library classifier distilled from the exemplar store.
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct TrainedHead {
    /// Feature version the model was trained against; stale versions are unusable.
    pub feat_version: u32,
    /// Standardization params (over the training exemplars), the head expects raw
    /// feature vectors and standardizes them itself, matching k-NN.
    pub means: Vec<f64>,
    pub stds: Vec<f64>,
    pub classes: Vec<TagModel>,
    /// Number of labeled exemplars the model was trained on (staleness/worthwhile UI).
    pub exemplar_count: usize,
    pub trained_at: i64,
}

impl TrainedHead {
    /// Whether the model matches the current feature layout. A mismatch means the vectors
    /// were recomputed under a new `FEATURE_VERSION`; the head must be retrained.
    pub fn is_current(&self) -> bool {
        self.feat_version == FEATURE_VERSION
    }

    /// Score every modeled tag for a raw (un-standardized) query vector, highest first.
    /// Neighbors are empty: the head trades k-NN's explainability for speed.
    pub fn score(&self, raw_query: &[f64]) -> Vec<TagScore> {
        if raw_query.len() != NUM_FEATURES {
            return Vec::new();
        }
        let x = standardize_vec(raw_query, &self.means, &self.stds);
        let mut scores: Vec<TagScore> = self
            .classes
            .iter()
            .filter(|m| m.weights.len() == x.len())
            .map(|m| {
                let z = dot(&m.weights, &x) + m.bias;
                TagScore {
                    tag: m.tag.clone(),
                    score: sigmoid(z),
                    neighbors: Vec::new(),
                }
            })
            .collect();
        // total_cmp never panics on NaN; the head is trained from the same
        // finiteness-filtered exemplar set, so this is a backstop.
        scores.sort_by(|a, b| b.score.total_cmp(&a.score).then(a.tag.cmp(&b.tag)));
        scores
    }
}

fn sigmoid(z: f64) -> f64 {
    1.0 / (1.0 + (-z).exp())
}

fn dot(a: &[f64], b: &[f64]) -> f64 {
    a.iter().zip(b).map(|(x, y)| x * y).sum()
}

/// Fit one balanced binary logistic regression by deterministic full-batch gradient
/// descent with L2 regularization. Per-sample class weights counter the heavy
/// positive/negative imbalance of multi-label tags (a tag is rare among all exemplars).
fn train_one(xs: &[&[f64]], labels: &[bool], dims: usize) -> (Vec<f64>, f64) {
    let n = xs.len() as f64;
    let pos = labels.iter().filter(|y| **y).count() as f64;
    let neg = n - pos;
    // Balance so positives and negatives contribute equal total weight.
    let w_pos = if pos > 0.0 { n / (2.0 * pos) } else { 0.0 };
    let w_neg = if neg > 0.0 { n / (2.0 * neg) } else { 0.0 };

    let mut w = vec![0.0; dims];
    let mut b = 0.0;
    for _ in 0..ITERATIONS {
        let mut grad_w = vec![0.0; dims];
        let mut grad_b = 0.0;
        for (x, &y) in xs.iter().zip(labels) {
            let p = sigmoid(dot(&w, x) + b);
            let sw = if y { w_pos } else { w_neg };
            let err = sw * (p - if y { 1.0 } else { 0.0 });
            for d in 0..dims {
                grad_w[d] += err * x[d];
            }
            grad_b += err;
        }
        for d in 0..dims {
            // Mean gradient (weights sum to n) + L2 pull toward zero.
            w[d] -= LEARNING_RATE * (grad_w[d] / n + L2 * w[d]);
        }
        b -= LEARNING_RATE * (grad_b / n);
    }
    (w, b)
}

/// Distill the current exemplar store into a trained head **without persisting it**.
/// Returns `None` when there is nothing worth training (no labeled exemplars, or no tag
/// clears [`MIN_POSITIVE_EXAMPLES`]).
#[instrument(skip_all)]
pub fn train(db: &Database) -> Result<Option<TrainedHead>> {
    let index = exemplar::build_index(db)?;
    if index.is_empty() {
        return Ok(None);
    }

    // Shared standardized design matrix + per-row tag sets (borrow the index once).
    let xs: Vec<&[f64]> = index.rows().map(|(v, _)| v).collect();
    let row_tags: Vec<&[String]> = index.rows().map(|(_, t)| t).collect();

    // Tags with enough positives to model; sorted for deterministic class order.
    let mut counts: HashMap<&str, usize> = HashMap::new();
    for tags in &row_tags {
        for t in *tags {
            *counts.entry(t.as_str()).or_default() += 1;
        }
    }
    let mut candidates: Vec<&str> = counts
        .iter()
        .filter(|(_, c)| **c >= MIN_POSITIVE_EXAMPLES)
        .map(|(t, _)| *t)
        .collect();
    candidates.sort_unstable();
    if candidates.is_empty() {
        return Ok(None);
    }

    let classes = candidates
        .iter()
        .map(|&tag| {
            let labels: Vec<bool> = row_tags
                .iter()
                .map(|tags| tags.iter().any(|t| t == tag))
                .collect();
            let (weights, bias) = train_one(&xs, &labels, NUM_FEATURES);
            TagModel {
                tag: tag.to_string(),
                weights,
                bias,
            }
        })
        .collect();

    Ok(Some(TrainedHead {
        feat_version: FEATURE_VERSION,
        means: index.means().to_vec(),
        stds: index.stds().to_vec(),
        classes,
        exemplar_count: index.len(),
        trained_at: unix_now(),
    }))
}

/// Distill and persist the head (the background-worker entry point). Persists the fresh
/// model, or [`clear`]s the stored head when there is nothing to train. Returns what was
/// stored (or `None`).
#[instrument(skip_all)]
pub fn train_and_save(db: &Database) -> Result<Option<TrainedHead>> {
    match train(db)? {
        Some(head) => {
            save(db, &head)?;
            Ok(Some(head))
        }
        None => {
            clear(db)?;
            Ok(None)
        }
    }
}

/// Persist the singleton trained head, replacing any prior model.
#[instrument(skip_all)]
pub fn save(db: &Database, head: &TrainedHead) -> Result<()> {
    let json = serde_json::to_string(head).map_err(|e| CoreError::Serialization(e.to_string()))?;
    db.conn().execute(
        "INSERT INTO trained_head (id, feat_version, exemplar_count, model, trained_at)
         VALUES (1, ?1, ?2, ?3, ?4)
         ON CONFLICT(id) DO UPDATE SET
            feat_version = ?1, exemplar_count = ?2, model = ?3, trained_at = ?4",
        rusqlite::params![
            head.feat_version,
            head.exemplar_count as i64,
            json,
            head.trained_at
        ],
    )?;
    Ok(())
}

/// Load the stored head, if any. A model trained under a stale `FEATURE_VERSION` is
/// dropped (returns `None`), callers should retrain.
#[instrument(skip_all)]
pub fn load(db: &Database) -> Result<Option<TrainedHead>> {
    let json: Option<String> = db
        .conn()
        .query_row("SELECT model FROM trained_head WHERE id = 1", [], |row| {
            row.get(0)
        })
        .ok();
    let Some(json) = json else { return Ok(None) };
    let head: TrainedHead =
        serde_json::from_str(&json).map_err(|e| CoreError::Serialization(e.to_string()))?;
    if !head.is_current() {
        return Ok(None);
    }
    Ok(Some(head))
}

/// Remove the stored head.
#[instrument(skip_all)]
pub fn clear(db: &Database) -> Result<()> {
    db.conn()
        .execute("DELETE FROM trained_head WHERE id = 1", [])?;
    Ok(())
}

/// Count current-version labeled exemplars cheaply (drives the worthwhile/staleness
/// decision without building the full index).
#[instrument(skip_all)]
pub fn labeled_exemplar_count(db: &Database) -> Result<usize> {
    let n: i64 = db.conn().query_row(
        "SELECT COUNT(DISTINCT t.sample_hash) FROM tags t
         JOIN sample_features f ON f.hash = t.sample_hash AND f.feat_version = ?1
         WHERE NOT EXISTS (
            SELECT 1 FROM tag_provenance p
            WHERE p.sample_hash = t.sample_hash AND p.tag = t.tag AND p.source = 'ml')",
        [FEATURE_VERSION],
        |row| row.get(0),
    )?;
    Ok(n as usize)
}

/// Whether a trained head is worth building for the current library size.
#[instrument(skip_all)]
pub fn is_worthwhile(db: &Database) -> Result<bool> {
    Ok(labeled_exemplar_count(db)? >= RECOMMENDED_MIN_EXEMPLARS)
}

/// Score a stored sample against the head under each tag's policy, writing nothing.
/// Mirrors [`exemplar::preview_sample`] but via the distilled model.
#[instrument(skip_all)]
pub fn preview_sample(db: &Database, hash: &str, head: &TrainedHead) -> Result<MlOutcome> {
    let Some(raw) = load_query_vector(db, hash)? else {
        return Ok(MlOutcome::default());
    };
    let scores = head.score(&raw);
    let existing = sample_tag_set(db, hash)?;
    apply_policy(db, scores, &existing)
}

/// Score a stored sample via the head and **auto-apply** above-threshold tags
/// (`source = 'ml'`), returning what was applied and what's pending review.
#[instrument(skip_all)]
pub fn apply_ml_suggestions(db: &Database, hash: &str, head: &TrainedHead) -> Result<MlOutcome> {
    let outcome = preview_sample(db, hash, head)?;
    db.transaction_core(|tx| {
        for s in &outcome.auto_applied {
            crate::rules::apply_tag_sourced(db, tx, hash, &s.tag, "ml")?;
        }
        Ok(())
    })?;
    Ok(outcome)
}

/// Auto-apply above-threshold head tags to every analyzed, non-deleted sample, returning
/// the total tags applied. The head-routed counterpart to [`exemplar::auto_apply_library`]
/// O(tags) per sample, independent of library size.
#[instrument(skip_all)]
pub fn auto_apply_library(db: &Database, head: &TrainedHead) -> Result<usize> {
    let hashes: Vec<String> = {
        let mut stmt = db.conn().prepare(
            "SELECT s.hash FROM live_samples s
             JOIN sample_features f ON f.hash = s.hash AND f.feat_version = ?1",
        )?;
        stmt.query_map([FEATURE_VERSION], |row| row.get(0))?
            .collect::<std::result::Result<Vec<_>, _>>()?
    };
    let mut applied = 0;
    for hash in hashes {
        applied += apply_ml_suggestions(db, &hash, head)?.auto_applied.len();
    }
    Ok(applied)
}

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

    fn insert(db: &Database, hash: &str, vec: &[f64; NUM_FEATURES], tags: &[&str]) {
        db.conn()
            .execute(
                "INSERT INTO samples (hash, original_name, file_extension, file_size, import_date, last_modified) \
                 VALUES (?1, ?1, 'wav', 1, 0, 0)",
                [hash],
            )
            .unwrap();
        let json = serde_json::to_string(&vec.to_vec()).unwrap();
        db.conn()
            .execute(
                "INSERT INTO sample_features (hash, feat_version, vector, computed_at) VALUES (?1, ?2, ?3, 0)",
                rusqlite::params![hash, FEATURE_VERSION, json],
            )
            .unwrap();
        for t in tags {
            crate::tags::add_tag(db, hash, t).unwrap();
        }
    }

    /// Build two separated clusters of kicks and snares with `n` each, varied per index so
    /// the standardizer sees non-zero variance.
    fn seed_two_classes(db: &Database, n: usize) {
        for i in 0..n {
            let k = 0.01 * i as f64;
            insert(
                db,
                &format!("k{i}"),
                &[k; NUM_FEATURES],
                &["instrument.drum.kick"],
            );
            insert(
                db,
                &format!("s{i}"),
                &[10.0 + k; NUM_FEATURES],
                &["instrument.drum.snare"],
            );
        }
    }

    #[test]
    fn no_exemplars_trains_nothing() {
        let db = Database::open_in_memory().unwrap();
        assert!(train(&db).unwrap().is_none());
    }

    #[test]
    fn too_few_positives_trains_nothing() {
        let db = Database::open_in_memory().unwrap();
        // Only 2 of each tag (< MIN_POSITIVE_EXAMPLES).
        seed_two_classes(&db, 2);
        assert!(train(&db).unwrap().is_none());
    }

    #[test]
    fn distinguishes_two_classes() {
        let db = Database::open_in_memory().unwrap();
        seed_two_classes(&db, 8);
        let head = train(&db).unwrap().expect("should train");
        assert_eq!(head.classes.len(), 2);
        assert!(head.is_current());

        // A query in the kick cluster scores kick above snare.
        let near_kick = [0.04; NUM_FEATURES];
        let scores = head.score(&near_kick);
        let kick = scores
            .iter()
            .find(|s| s.tag == "instrument.drum.kick")
            .unwrap();
        let snare = scores
            .iter()
            .find(|s| s.tag == "instrument.drum.snare")
            .unwrap();
        assert!(
            kick.score > snare.score,
            "kick {} vs snare {}",
            kick.score,
            snare.score
        );
        assert!(kick.score > 0.5);
    }

    #[test]
    fn save_load_round_trip_and_clear() {
        let db = Database::open_in_memory().unwrap();
        seed_two_classes(&db, 6);
        let head = train_and_save(&db).unwrap().unwrap();
        let loaded = load(&db).unwrap().expect("persisted");
        assert_eq!(loaded.classes.len(), head.classes.len());
        assert_eq!(loaded.exemplar_count, head.exemplar_count);
        clear(&db).unwrap();
        assert!(load(&db).unwrap().is_none());
    }

    #[test]
    fn stale_feature_version_is_ignored_on_load() {
        let db = Database::open_in_memory().unwrap();
        seed_two_classes(&db, 6);
        let mut head = train(&db).unwrap().unwrap();
        head.feat_version = FEATURE_VERSION + 1;
        save(&db, &head).unwrap();
        // Stored row exists but is stale → load returns None.
        assert!(load(&db).unwrap().is_none());
    }

    #[test]
    fn deterministic_across_runs() {
        let train_once = || {
            let db = Database::open_in_memory().unwrap();
            seed_two_classes(&db, 7);
            train(&db).unwrap().unwrap().classes
        };
        let a = train_once();
        let b = train_once();
        assert_eq!(a.len(), b.len());
        for (ca, cb) in a.iter().zip(&b) {
            assert_eq!(ca.tag, cb.tag);
            assert_eq!(
                ca.bias.to_bits(),
                cb.bias.to_bits(),
                "bias differs for {}",
                ca.tag
            );
            for (wa, wb) in ca.weights.iter().zip(&cb.weights) {
                assert_eq!(wa.to_bits(), wb.to_bits());
            }
        }
    }

    #[test]
    fn apply_auto_tags_via_head() {
        let db = Database::open_in_memory().unwrap();
        seed_two_classes(&db, 10);
        let head = train(&db).unwrap().unwrap();
        // Unlabeled query in the kick cluster.
        insert(&db, "q", &[0.05; NUM_FEATURES], &[]);
        let outcome = apply_ml_suggestions(&db, "q", &head).unwrap();
        let applied: Vec<&str> = outcome
            .auto_applied
            .iter()
            .map(|s| s.tag.as_str())
            .collect();
        assert!(applied.contains(&"instrument.drum.kick"), "got {applied:?}");
        // Provenance is 'ml'.
        let prov = crate::rules::sample_tag_provenance(&db, "q").unwrap();
        assert!(
            prov.iter()
                .any(|(t, s, _)| t == "instrument.drum.kick" && s == "ml")
        );
    }

    #[test]
    fn labeled_count_and_worthwhile() {
        let db = Database::open_in_memory().unwrap();
        seed_two_classes(&db, 5); // 10 labeled exemplars
        assert_eq!(labeled_exemplar_count(&db).unwrap(), 10);
        assert!(!is_worthwhile(&db).unwrap()); // well under RECOMMENDED_MIN_EXEMPLARS
    }
}