medh5 2.1.0

//! Auditing split claims across a cohort (§12.3).
//!
//! A split claim in a file is a **membership claim, not an authority**: the
//! dataset manifest decides.  Two findings matter and they are different:
//! *conflicting claims* (W906: one `set_id` against two manifest hashes, so a
//! file predates a re-split) and *subject leakage* (one unit of anatomy in two
//! partitions, or two folds, of one split).  Both are invisible in any one
//! file, so this is a cross-file operation.

use std::collections::{BTreeMap, BTreeSet, HashMap};
use std::path::Path;

use serde_json::{json, Map, Value};

use crate::collection::{open_any, AnyFile};
use crate::curation::identity::SplitClaim;
use crate::json::repr_str;
use crate::sample::Sample;
use crate::Result;

/// One file's claim to one partition of one split.
#[derive(Debug, Clone, PartialEq)]
pub struct Membership {
    pub path: String,
    pub sample_id: String,
    pub subject_id: String,
    pub group_id: String,
    pub claim: SplitClaim,
}

impl Membership {
    pub fn set_id(&self) -> &str {
        &self.claim.set_id
    }

    pub fn partition(&self) -> &str {
        &self.claim.partition
    }

    pub fn to_json(&self) -> Value {
        let mut out = Map::new();
        out.insert("path".into(), json!(self.path));
        out.insert("sample_id".into(), json!(self.sample_id));
        out.insert("subject_id".into(), json!(self.subject_id));
        out.insert("group_id".into(), json!(self.group_id));
        if let Value::Object(claim) = self.claim.to_json() {
            out.extend(claim);
        }
        Value::Object(out)
    }
}

/// One unit of anatomy appearing in more than one partition of one split.
///
/// The unit is a connected component of `subject_id` and grouping key: two
/// files are one unit if they share either, transitively.  A fold is a
/// partition of its own.
#[derive(Debug, Clone, PartialEq)]
pub struct Leak {
    pub set_id: String,
    pub group_id: String,
    /// Each place the unit's claims put it: the partition, as `train (fold 0)`
    /// where the claim names a fold.
    pub partitions: Vec<String>,
    pub paths: Vec<String>,
    pub groups: Vec<String>,
    pub subjects: Vec<String>,
}

impl Leak {
    pub fn to_json(&self) -> Value {
        let groups = if self.groups.is_empty() { vec![self.group_id.clone()] } else { self.groups.clone() };
        json!({
            "set_id": self.set_id,
            "group_id": self.group_id,
            "partitions": self.partitions,
            "paths": self.paths,
            "groups": groups,
            "subjects": self.subjects,
        })
    }
}

impl std::fmt::Display for Leak {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let joined = if self.groups.len() > 1 {
            format!(
                " (joined with {} by subject {})",
                self.groups[1..].iter().map(|g| repr_str(g)).collect::<Vec<_>>().join(", "),
                self.subjects.iter().map(|s| repr_str(s)).collect::<Vec<_>>().join(", ")
            )
        } else {
            String::new()
        };
        write!(
            f,
            "{}: group {}{joined} is in {} ({} files)",
            self.set_id,
            repr_str(&self.group_id),
            self.partitions.join(", "),
            self.paths.len()
        )
    }
}

/// One `set_id` claimed against more than one manifest (W906).
#[derive(Debug, Clone, PartialEq)]
pub struct Conflict {
    pub set_id: String,
    pub manifests: Vec<String>,
    pub paths_by_manifest: BTreeMap<String, Vec<String>>,
}

impl Conflict {
    pub fn to_json(&self) -> Value {
        json!({"set_id": self.set_id, "manifests": self.manifests, "paths_by_manifest": self.paths_by_manifest})
    }
}

impl std::fmt::Display for Conflict {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let short: Vec<String> = self.manifests.iter().map(|m| m.chars().take(12).collect()).collect();
        write!(f, "{}: {} different manifest hashes ({})", self.set_id, self.manifests.len(), short.join(", "))
    }
}

/// What a cohort's split claims say, and where they disagree.
#[derive(Debug, Clone, Default, PartialEq)]
pub struct SplitAudit {
    pub memberships: Vec<Membership>,
    pub conflicts: Vec<Conflict>,
    pub leaks: Vec<Leak>,
    /// Files carrying no split claim at all --- not an error, but easy to lose.
    pub unclaimed: Vec<String>,
    pub unreadable: Vec<(String, String)>,
}

impl SplitAudit {
    pub fn ok(&self) -> bool {
        self.conflicts.is_empty() && self.leaks.is_empty() && self.unreadable.is_empty()
    }

    pub fn set_ids(&self) -> Vec<String> {
        self.memberships.iter().map(|m| m.claim.set_id.clone()).collect::<BTreeSet<_>>().into_iter().collect()
    }

    /// `partition -> sample ids` for one split.
    pub fn partitions(&self, set_id: &str) -> BTreeMap<String, Vec<String>> {
        let mut out: BTreeMap<String, Vec<String>> = BTreeMap::new();
        for member in self.memberships.iter().filter(|m| m.set_id() == set_id) {
            out.entry(member.partition().to_string()).or_default().push(member.sample_id.clone());
        }
        for ids in out.values_mut() {
            ids.sort();
        }
        out
    }

    pub fn counts(&self) -> BTreeMap<String, BTreeMap<String, usize>> {
        self.set_ids()
            .into_iter()
            .map(|set_id| {
                let counts = self.partitions(&set_id).into_iter().map(|(k, v)| (k, v.len())).collect();
                (set_id, counts)
            })
            .collect()
    }

    pub fn to_json(&self) -> Value {
        json!({
            "ok": self.ok(),
            "sets": self.set_ids(),
            "counts": self.counts(),
            "conflicts": self.conflicts.iter().map(Conflict::to_json).collect::<Vec<_>>(),
            "leaks": self.leaks.iter().map(Leak::to_json).collect::<Vec<_>>(),
            "unclaimed": self.unclaimed,
            "unreadable": self.unreadable.iter().map(|(p, e)| json!({"path": p, "error": e})).collect::<Vec<_>>(),
        })
    }
}

/// Read every file's claims and cross-check them (§12.3).
pub fn audit_splits<P: AsRef<Path>>(paths: &[P]) -> SplitAudit {
    let mut memberships = Vec::new();
    let mut unclaimed = Vec::new();
    let mut unreadable = Vec::new();
    for path in paths {
        let path = path.as_ref();
        let text = path.to_string_lossy().into_owned();
        match memberships_of(path) {
            Err(e) => unreadable.push((text, e.python_line())),
            Ok(found) => {
                if found.is_empty() {
                    unclaimed.push(text);
                }
                memberships.extend(found);
            }
        }
    }
    unclaimed.sort();
    unreadable.sort();
    let conflicts = conflicts(&memberships);
    let leaks = leaks(&memberships);
    SplitAudit { memberships, conflicts, leaks, unclaimed, unreadable }
}

/// Every claim in a file, whether it holds one sample or many.
pub fn memberships_of(path: &Path) -> Result<Vec<Membership>> {
    let text = path.to_string_lossy().into_owned();
    let samples: Vec<(String, Sample)> = match open_any(path, None)? {
        AnyFile::Sample(s) => vec![(String::new(), s)],
        AnyFile::Collection(c) => c.keys()?.into_iter().map(|k| Ok((k.clone(), c.get(&k)?))).collect::<Result<_>>()?,
    };
    let mut out = Vec::new();
    for (key, sample) in samples {
        let document = sample.document()?;
        let identity = &document.identity;
        let location = if key.is_empty() { text.clone() } else { format!("{text}::{key}") };
        for claim in &document.splits {
            out.push(Membership {
                path: location.clone(),
                sample_id: identity.sample_id.clone(),
                subject_id: identity.subject_id.clone(),
                group_id: document.cohort.grouping_key(&identity.subject_id).to_string(),
                claim: claim.clone(),
            });
        }
    }
    Ok(out)
}

fn conflicts(memberships: &[Membership]) -> Vec<Conflict> {
    let mut by_set: BTreeMap<String, BTreeMap<String, Vec<String>>> = BTreeMap::new();
    for member in memberships {
        let Some(digest) = member.claim.manifest_sha256.as_deref().filter(|d| !d.is_empty()) else { continue };
        by_set
            .entry(member.claim.set_id.clone())
            .or_default()
            .entry(digest.to_string())
            .or_default()
            .push(member.path.clone());
    }
    by_set
        .into_iter()
        .filter(|(_, by_manifest)| by_manifest.len() > 1)
        .map(|(set_id, by_manifest)| Conflict {
            set_id,
            manifests: by_manifest.keys().cloned().collect(),
            paths_by_manifest: by_manifest
                .into_iter()
                .map(|(k, mut v)| {
                    v.sort();
                    (k, v)
                })
                .collect(),
        })
        .collect()
}

/// Grouping key -> every grouping key sharing anatomy with it, sorted.
///
/// `pairs` are `(subject_id, group_id)`.  The unit is the connected component
/// of subjects and keys (union--find): two visits of one subject curated under
/// two `group_id` values are one unit, which grouping by the key alone misses.
pub fn anatomy_units<'a>(pairs: impl IntoIterator<Item = (&'a str, &'a str)>) -> BTreeMap<String, Vec<String>> {
    type Node = (bool, String); // (is_subject, name); subjects sort after groups like ("group" < "subject")
    let mut parent: HashMap<Node, Node> = HashMap::new();
    let mut order: Vec<Node> = Vec::new();
    fn find(parent: &mut HashMap<Node, Node>, order: &mut Vec<Node>, node: Node) -> Node {
        if !parent.contains_key(&node) {
            parent.insert(node.clone(), node.clone());
            order.push(node.clone());
        }
        let mut node = node;
        loop {
            let up = parent[&node].clone();
            if up == node {
                return node;
            }
            let grand = parent[&up].clone();
            parent.insert(node.clone(), grand.clone());
            node = grand;
        }
    }
    for (subject, group) in pairs {
        let a = find(&mut parent, &mut order, (true, subject.to_string()));
        let b = find(&mut parent, &mut order, (false, group.to_string()));
        if a != b {
            let (lo, hi) = if a < b { (a, b) } else { (b, a) };
            parent.insert(hi, lo);
        }
    }
    // Each root's unit, found by hashing: a linear search per key made 40,000
    // singleton units take 2.5 s.
    let mut members: HashMap<Node, Vec<String>> = HashMap::new();
    for node in order.clone() {
        if !node.0 {
            let root = find(&mut parent, &mut order, node.clone());
            members.entry(root).or_default().push(node.1);
        }
    }
    let mut out = BTreeMap::new();
    for mut unit in members.into_values() {
        unit.sort();
        for group in &unit {
            out.insert(group.clone(), unit.clone());
        }
    }
    out
}

fn leaks(memberships: &[Membership]) -> Vec<Leak> {
    let units = anatomy_units(memberships.iter().map(|m| (m.subject_id.as_str(), m.group_id.as_str())));
    let mut by_unit: BTreeMap<(String, Vec<String>), Vec<&Membership>> = BTreeMap::new();
    for member in memberships {
        let unit = units.get(&member.group_id).cloned().unwrap_or_else(|| vec![member.group_id.clone()]);
        by_unit.entry((member.claim.set_id.clone(), unit)).or_default().push(member);
    }
    let mut out = Vec::new();
    for ((set_id, groups), members) in by_unit {
        // A fold is a partition of its own: k-fold claims are `train` in
        // every fold but the validation one, so a unit in folds 0 and 3 read
        // `train` twice and was no leak (`Split::leaks` holds the same rule).
        let placements: BTreeSet<(&str, Option<i64>)> =
            members.iter().map(|m| (m.claim.partition.as_str(), m.claim.fold_int())).collect();
        if placements.len() > 1 {
            let partitions: Vec<String> = placements
                .into_iter()
                .map(|(partition, fold)| match fold {
                    Some(fold) => format!("{partition} (fold {fold})"),
                    None => partition.to_string(),
                })
                .collect();
            let mut paths: Vec<String> = members.iter().map(|m| m.path.clone()).collect();
            paths.sort();
            out.push(Leak {
                set_id,
                group_id: groups[0].clone(),
                partitions,
                paths,
                subjects: members.iter().map(|m| m.subject_id.clone()).collect::<BTreeSet<_>>().into_iter().collect(),
                groups,
            });
        }
    }
    out
}

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

    #[test]
    fn units_join_keys_through_a_shared_subject() {
        let units = anatomy_units([("s1", "g1"), ("s1", "g2"), ("s2", "g3")]);
        assert_eq!(units["g1"], vec!["g1", "g2"]);
        assert_eq!(units["g2"], vec!["g1", "g2"]);
        assert_eq!(units["g3"], vec!["g3"]);
    }

    #[test]
    fn s12_3_units_are_the_connected_components_of_subjects_and_keys() {
        // Pairs that chain, fan out and stand alone, against a flood fill.
        let mut state = 7u64;
        let pairs: Vec<(String, String)> = (0..800)
            .map(|_| {
                state = state.wrapping_mul(6_364_136_223_846_793_005).wrapping_add(1_442_695_040_888_963_407);
                (format!("s{}", (state >> 33) % 700), format!("g{}", (state >> 13) % 900))
            })
            .collect();
        let units = anatomy_units(pairs.iter().map(|(s, g)| (s.as_str(), g.as_str())));
        let mut by_subject: HashMap<&str, Vec<&str>> = HashMap::new();
        let mut by_group: HashMap<&str, Vec<&str>> = HashMap::new();
        for (subject, group) in &pairs {
            by_subject.entry(subject).or_default().push(group);
            by_group.entry(group).or_default().push(subject);
        }
        assert_eq!(units.len(), by_group.len());
        let mut sizes = BTreeSet::new();
        for start in by_group.keys() {
            let (mut groups, mut subjects, mut todo) = (BTreeSet::from([*start]), BTreeSet::new(), vec![*start]);
            while let Some(group) = todo.pop() {
                for subject in &by_group[group] {
                    if subjects.insert(*subject) {
                        todo.extend(by_subject[subject].iter().filter(|g| groups.insert(**g)));
                    }
                }
            }
            sizes.insert(groups.len());
            assert_eq!(units[*start], groups.into_iter().collect::<Vec<_>>(), "{start}");
        }
        assert!(sizes.len() > 3, "{sizes:?}");
    }

    #[test]
    fn s12_3_the_units_of_a_large_cohort_are_found_without_a_search_per_key() {
        // Each key's unit was found by a linear search through the units so
        // far: 40,000 singleton units took 2.5 s, and this many a minute.
        let n = 200_000;
        let pairs: Vec<(String, String)> = (0..n).map(|i| (format!("s{i}"), format!("g{i}"))).collect();
        let started = std::time::Instant::now();
        let units = anatomy_units(pairs.iter().map(|(s, g)| (s.as_str(), g.as_str())));
        // Generous, for unoptimised builds on shared CI runners: still half
        // what the search per key took at this size.
        assert!(started.elapsed() < std::time::Duration::from_secs(30), "{:?}", started.elapsed());
        assert_eq!(units.len(), n);
        assert_eq!(units["g7"], ["g7"]);
    }

    fn claimed(path: &str, subject: &str, partition: &str, fold: Option<i64>) -> Membership {
        Membership {
            path: path.into(),
            sample_id: path.into(),
            subject_id: subject.into(),
            group_id: subject.into(),
            claim: SplitClaim {
                set_id: "cv".into(),
                partition: partition.into(),
                fold: fold.map(serde_json::Number::from),
                assigned_by: None,
                assigned_at: None,
                manifest_sha256: None,
            },
        }
    }

    #[test]
    fn s12_3_a_unit_in_two_folds_is_a_leak() {
        // Fold 2 validates: a subject in folds 0 and 3 is `train` in both
        // claims, and trains the model that validates on it in either fold.
        let found = leaks(&[claimed("a.medh5", "p1", "train", Some(0)), claimed("b.medh5", "p1", "train", Some(3))]);
        assert_eq!(found.len(), 1);
        assert_eq!(found[0].partitions, ["train (fold 0)", "train (fold 3)"]);
        assert_eq!(found[0].paths, ["a.medh5", "b.medh5"]);
        assert_eq!(found[0].subjects, ["p1"]);
        let found = leaks(&[claimed("a.medh5", "p1", "val", Some(2)), claimed("b.medh5", "p1", "train", Some(0))]);
        assert_eq!(found[0].partitions, ["train (fold 0)", "val (fold 2)"]);
        // One fold is one place.
        assert!(
            leaks(&[claimed("a.medh5", "p1", "train", Some(0)), claimed("b.medh5", "p1", "train", Some(0))]).is_empty()
        );
        // Without folds, partitions are named as they were.
        let found = leaks(&[claimed("a.medh5", "p1", "train", None), claimed("b.medh5", "p1", "test", None)]);
        assert_eq!(found[0].partitions, ["test", "train"]);
        assert!(leaks(&[claimed("a.medh5", "p1", "train", None), claimed("b.medh5", "p1", "train", None)]).is_empty());
    }
}