//! Разбиение реплик по говорящим по отпечаткам голоса.
//!
//! Отпечаток — вектор, у одного голоса векторы смотрят в одну сторону
//! (косинусная близость высокая), у разных — в разные. Сколько людей в записи,
//! заранее неизвестно, поэтому разбиение — снизу вверх: каждая реплика сначала
//! сама по себе, потом сливаются две самые похожие группы, и так, пока похожих
//! не останется (или пока групп не станет столько, сколько сказал человек).
//!
//! Короткая реплика («Ага», «Да») даёт шумный отпечаток: в слиянии участвуют
//! только реплики от [`MIN_ANCHOR_S`] секунд, остальные потом просто
//! прикладываются к ближайшему говорящему.

/// Сколько говорящих искать.
#[derive(Debug, Clone, Copy, PartialEq)]
pub enum Count {
    /// Сколько получится: группы сливаются, пока средняя близость между ними не
    /// ниже порога.
    Auto { threshold: f32 },
    /// Ровно столько (если реплик хватает).
    Fixed(usize),
}

/// Порог по умолчанию для модели WeSpeaker ResNet34 на отпечатках без
/// общего сдвига записи (см. `centered`). Подобран на подкасте с тремя
/// людьми в студии (`examples/voices.rs`): весь выпуск 1 ч 52 мин — трое
/// (49 %, 37 %, 14 % речи); при −0,05 двое сливаются, при 0,05 — пятеро.
pub const DEFAULT_THRESHOLD: f32 = 0.0;

/// Реплики короче — не опора для разбиения, их отпечаток шумный.
pub const MIN_ANCHOR_S: f32 = 1.2;

/// Больше опор не берём: слияние — квадрат от их числа по памяти и больше по
/// времени. Двухчасовая встреча — около 2000 реплик; 800 самых длинных
/// хватает, чтобы услышать каждого, кто говорил хоть минуту.
const MAX_ANCHORS: usize = 800;

/// Говорящий, у которого меньше этой доли речи, в автоматическом режиме
/// растворяется в остальных: это чаще обрывок чужого голоса, смех, музыка.
const MIN_SHARE: f32 = 0.03;

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

fn normalized(mut v: Vec<f32>) -> Vec<f32> {
    let n = dot(&v, &v).sqrt();
    if n > 0.0 {
        v.iter_mut().for_each(|x| *x /= n);
    }
    v
}

/// Номер говорящего для каждой реплики: 0, 1, 2… по порядку первого
/// появления (реплики — по времени). `embeddings` — отпечатки, `seconds` —
/// длина реплик. Пустой отпечаток (реплика слишком короткая, чтобы его
/// снять) получает говорящего соседней реплики.
pub fn assign(embeddings: &[Vec<f32>], seconds: &[f32], count: Count) -> Vec<usize> {
    assert_eq!(embeddings.len(), seconds.len());
    let n = embeddings.len();
    if n == 0 {
        return Vec::new();
    }
    let emb = centered(embeddings, seconds);
    let usable = |i: usize| !emb[i].is_empty();

    // Опоры: самые длинные реплики с отпечатком.
    let mut anchors: Vec<usize> = (0..n)
        .filter(|&i| usable(i) && seconds[i] >= MIN_ANCHOR_S)
        .collect();
    if anchors.is_empty() {
        anchors = (0..n).filter(|&i| usable(i)).collect();
    }
    if anchors.is_empty() {
        return vec![0; n];
    }
    anchors.sort_by(|&a, &b| seconds[b].total_cmp(&seconds[a]));
    anchors.truncate(MAX_ANCHORS);

    let groups = merge(&emb, &anchors, count);
    let mut centroids = centroids(&emb, seconds, &groups);
    if let Count::Auto { .. } = count {
        dissolve_small(&mut centroids, &emb, seconds, &anchors);
    }

    // Каждую реплику — к ближайшему говорящему; два прохода уточнения, как в
    // k-средних: центры пересчитываются по всем репликам, а не только опорам.
    let mut labels = vec![usize::MAX; n];
    for _ in 0..2 {
        for i in (0..n).filter(|&i| usable(i)) {
            labels[i] = nearest(&centroids, &emb[i]);
        }
        let groups: Vec<Vec<usize>> = (0..centroids.len())
            .map(|k| (0..n).filter(|&i| labels[i] == k).collect())
            .collect();
        let fresh = centroids_of(&emb, seconds, &groups);
        for (c, f) in centroids.iter_mut().zip(fresh) {
            if let Some(f) = f {
                *c = f;
            }
        }
    }
    for i in (0..n).filter(|&i| usable(i)) {
        labels[i] = nearest(&centroids, &emb[i]);
    }
    fill_gaps(&mut labels);
    renumber(&labels)
}

/// Отпечатки без общего для всей записи: из каждого вычитается средний.
/// У записи есть своё «звучание» — комната, микрофон, сжатие, фоновая
/// музыка, — и оно одинаково сдвигает отпечатки всех говорящих: на подкасте
/// близость любых двух реплик была 0,5–0,8. Без общего сдвига остаются
/// различия голосов, и один порог годится для разных записей.
fn centered(embeddings: &[Vec<f32>], seconds: &[f32]) -> Vec<Vec<f32>> {
    let unit: Vec<Vec<f32>> = embeddings.iter().map(|e| normalized(e.clone())).collect();
    let Some(dim) = unit.iter().map(Vec::len).find(|&d| d > 0) else {
        return unit;
    };
    let mut mean = vec![0f32; dim];
    let mut total = 0f32;
    for (e, s) in unit.iter().zip(seconds).filter(|(e, _)| e.len() == dim) {
        let w = s.max(0.1);
        mean.iter_mut().zip(e).for_each(|(m, x)| *m += x * w);
        total += w;
    }
    mean.iter_mut().for_each(|m| *m /= total);
    unit.into_iter()
        .map(|e| {
            if e.len() != dim {
                return Vec::new();
            }
            normalized(e.iter().zip(&mean).map(|(x, m)| x - m).collect())
        })
        .collect()
}

/// Слияние снизу вверх со средней связью: близость двух групп — средняя
/// близость их реплик.
fn merge(emb: &[Vec<f32>], anchors: &[usize], count: Count) -> Vec<Vec<usize>> {
    let m = anchors.len();
    let mut sim = vec![0f32; m * m];
    for a in 0..m {
        for b in a + 1..m {
            let s = dot(&emb[anchors[a]], &emb[anchors[b]]);
            sim[a * m + b] = s;
            sim[b * m + a] = s;
        }
    }
    let mut members: Vec<Vec<usize>> = anchors.iter().map(|&i| vec![i]).collect();
    let mut alive = vec![true; m];
    let mut left = m;
    let target = match count {
        Count::Fixed(k) => k.max(1),
        Count::Auto { .. } => 1,
    };
    while left > target {
        let mut best = (f32::MIN, 0, 0);
        for a in (0..m).filter(|&a| alive[a]) {
            for b in (a + 1..m).filter(|&b| alive[b]) {
                if sim[a * m + b] > best.0 {
                    best = (sim[a * m + b], a, b);
                }
            }
        }
        let (s, a, b) = best;
        if let Count::Auto { threshold } = count {
            if s < threshold {
                break;
            }
        }
        let (na, nb) = (members[a].len() as f32, members[b].len() as f32);
        for k in (0..m).filter(|&k| alive[k] && k != a && k != b) {
            let v = (na * sim[a * m + k] + nb * sim[b * m + k]) / (na + nb);
            sim[a * m + k] = v;
            sim[k * m + a] = v;
        }
        let moved = std::mem::take(&mut members[b]);
        members[a].extend(moved);
        alive[b] = false;
        left -= 1;
    }
    (0..m)
        .filter(|&a| alive[a])
        .map(|a| std::mem::take(&mut members[a]))
        .collect()
}

fn centroids_of(emb: &[Vec<f32>], seconds: &[f32], groups: &[Vec<usize>]) -> Vec<Option<Vec<f32>>> {
    groups
        .iter()
        .map(|g| {
            let first = g.first()?;
            let mut c = vec![0f32; emb[*first].len()];
            for &i in g {
                for (x, y) in c.iter_mut().zip(&emb[i]) {
                    *x += y * seconds[i].max(0.1);
                }
            }
            Some(normalized(c))
        })
        .collect()
}

fn centroids(emb: &[Vec<f32>], seconds: &[f32], groups: &[Vec<usize>]) -> Vec<Vec<f32>> {
    centroids_of(emb, seconds, groups).into_iter().flatten().collect()
}

fn nearest(centroids: &[Vec<f32>], e: &[f32]) -> usize {
    centroids
        .iter()
        .enumerate()
        .map(|(k, c)| (k, dot(c, e)))
        .max_by(|a, b| a.1.total_cmp(&b.1))
        .map(|(k, _)| k)
        .unwrap_or(0)
}

/// Говорящие с долей речи меньше [`MIN_SHARE`] растворяются в остальных.
fn dissolve_small(centroids: &mut Vec<Vec<f32>>, emb: &[Vec<f32>], seconds: &[f32], anchors: &[usize]) {
    loop {
        if centroids.len() <= 1 {
            return;
        }
        let mut talk = vec![0f32; centroids.len()];
        for &i in anchors {
            talk[nearest(centroids, &emb[i])] += seconds[i];
        }
        let total: f32 = talk.iter().sum();
        let (k, t) = talk
            .iter()
            .enumerate()
            .min_by(|a, b| a.1.total_cmp(b.1))
            .map(|(k, t)| (k, *t))
            .unwrap();
        if t >= total * MIN_SHARE {
            return;
        }
        centroids.remove(k);
    }
}

/// Реплика без отпечатка — говорящий предыдущей (или следующей, если она первая).
fn fill_gaps(labels: &mut [usize]) {
    let first = labels.iter().copied().find(|&l| l != usize::MAX).unwrap_or(0);
    let mut prev = first;
    for l in labels.iter_mut() {
        if *l == usize::MAX {
            *l = prev;
        } else {
            prev = *l;
        }
    }
}

/// 0, 1, 2… по порядку первого появления: «Спикер 1» — тот, кто заговорил первым.
fn renumber(labels: &[usize]) -> Vec<usize> {
    let mut map = std::collections::HashMap::new();
    labels
        .iter()
        .map(|l| {
            let next = map.len();
            *map.entry(*l).or_insert(next)
        })
        .collect()
}

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

    /// Детерминированный «шум» без внешних крейтов.
    fn noise(seed: u32, dim: usize, amp: f32) -> Vec<f32> {
        let mut x = seed.wrapping_mul(2_654_435_761).wrapping_add(1);
        (0..dim)
            .map(|_| {
                x ^= x << 13;
                x ^= x >> 17;
                x ^= x << 5;
                ((x % 10_000) as f32 / 10_000.0 - 0.5) * amp
            })
            .collect()
    }

    /// Голос — направление; реплика — направление плюс шум.
    fn voice(v: usize, seed: u32) -> Vec<f32> {
        let base = noise(1000 + v as u32, 64, 2.0);
        base.iter()
            .zip(noise(seed, 64, 1.1))
            .map(|(a, b)| a + b)
            .collect()
    }

    fn talk(order: &[usize]) -> (Vec<Vec<f32>>, Vec<f32>) {
        let emb = order
            .iter()
            .enumerate()
            .map(|(i, &v)| voice(v, i as u32))
            .collect();
        (emb, vec![4.0; order.len()])
    }

    #[test]
    fn finds_three_voices_and_numbers_them_by_first_word() {
        let order = [2, 2, 0, 1, 0, 2, 1, 1, 0, 2, 0, 1, 2, 0];
        let (emb, sec) = talk(&order);
        let got = assign(
            &emb,
            &sec,
            Count::Auto {
                threshold: DEFAULT_THRESHOLD,
            },
        );
        // Кто заговорил первым — 0: голос 2 → 0, голос 0 → 1, голос 1 → 2.
        let want: Vec<usize> = order.iter().map(|v| [1, 2, 0][*v]).collect();
        assert_eq!(got, want);
    }

    #[test]
    fn fixed_count_is_respected() {
        let order = [0, 1, 2, 0, 1, 2, 0, 1, 2];
        let (emb, sec) = talk(&order);
        let two = assign(&emb, &sec, Count::Fixed(2));
        assert_eq!(two.iter().max(), Some(&1));
        let one = assign(&emb, &sec, Count::Fixed(1));
        assert!(one.iter().all(|&l| l == 0));
    }

    #[test]
    fn a_stray_second_does_not_become_a_speaker() {
        let mut order = vec![0; 20];
        order.extend([1; 20]);
        let (mut emb, mut sec) = talk(&order);
        // Одна реплика чужого голоса на 1,5 с — меньше 3 % речи.
        emb.push(voice(7, 99));
        sec.push(1.5);
        let got = assign(
            &emb,
            &sec,
            Count::Auto {
                threshold: DEFAULT_THRESHOLD,
            },
        );
        assert_eq!(got.iter().max(), Some(&1), "{got:?}");
    }

    #[test]
    fn short_and_empty_replies_follow_their_neighbours() {
        let order = [0, 0, 1, 1];
        let (mut emb, mut sec) = talk(&order);
        emb.insert(2, Vec::new());
        sec.insert(2, 0.3);
        let got = assign(
            &emb,
            &sec,
            Count::Auto {
                threshold: DEFAULT_THRESHOLD,
            },
        );
        assert_eq!(got, [0, 0, 0, 1, 1]);
        assert!(assign(&[], &[], Count::Fixed(2)).is_empty());
        assert_eq!(assign(&[Vec::new()], &[0.2], Count::Fixed(2)), [0]);
    }
}