sensorlm/data/captioning/

semantic.rs

//! Level-3 (semantic) caption generation.
//!
//! Produces high-level, human-readable descriptions of labelled **activities**,
//! **sleep periods**, and **mood** events that occurred during the measurement
//! window.
//!
//! # Caption structure
//!
//! ```text
//! Walking was detected between minutes 480 and 540.
//! Running occurred from minute 600 to 660.
//! Sleep during minutes 0 to 440.
//! The person logged their mood as calm at minute 300.
//! ```
//!
//! # Reference correspondence
//!
//! This is a direct port of `captioning.py::generate_semantic_caption`.

use rand::{seq::SliceRandom, Rng};

use crate::data::captioning::templates::{ACTIVITY_TEMPLATES, MOOD_TEMPLATES};

// ---------------------------------------------------------------------------
// Public types
// ---------------------------------------------------------------------------

/// A labelled activity interval.
#[derive(Debug, Clone)]
pub struct ActivityEvent {
    /// Activity name as it will appear in the caption (e.g. `"Walking"`).
    pub name: String,
    /// Start offset in minutes from midnight (0–1439).
    pub start_minute: usize,
    /// End offset in minutes from midnight (0–1439).
    pub end_minute: usize,
}

impl ActivityEvent {
    /// Duration of the event in minutes.
    pub fn duration(&self) -> usize {
        self.end_minute.saturating_sub(self.start_minute) + 1
    }
}

/// A sleep interval.
#[derive(Debug, Clone)]
pub struct SleepEvent {
    /// Start of sleep in minutes from midnight.
    pub start_minute: usize,
    /// End of sleep in minutes from midnight.
    pub end_minute: usize,
}

/// A self-reported mood entry.
#[derive(Debug, Clone)]
pub struct MoodEvent {
    /// Mood label, e.g. `"calm"`, `"stressed"`.
    pub mood: String,
    /// Time of report in minutes from midnight.
    pub minute: usize,
}

// ---------------------------------------------------------------------------
// Caption generator
// ---------------------------------------------------------------------------

/// Generate a level-3 semantic caption.
///
/// # Arguments
///
/// * `activities`           – Labelled activity events.
/// * `sleep`                – Sleep interval events.
/// * `moods`                – Mood log entries (may be empty).
/// * `top_k_activity`       – Keep at most this many activities (by duration,
///   shortest first so the longest — most significant — are included).
/// * `top_k_sleep`          – Keep at most this many sleep periods.
/// * `min_activity_duration`– Activities shorter than this many minutes are
///   discarded (default: 20 min, matching reference).
/// * `rng`                  – Random number generator for template selection.
///
/// # Returns
///
/// A multi-line string that combines the activity, sleep, and mood captions.
pub fn generate_semantic_caption<R: Rng>(
    activities: &[ActivityEvent],
    sleep: &[SleepEvent],
    moods: &[MoodEvent],
    top_k_activity: usize,
    top_k_sleep: usize,
    min_activity_duration: usize,
    rng: &mut R,
) -> String {
    let mut parts = Vec::<String>::new();

    // -----------------------------------------------------------------------
    // Activities
    // -----------------------------------------------------------------------
    let mut valid_acts: Vec<&ActivityEvent> = activities
        .iter()
        .filter(|a| a.duration() >= min_activity_duration)
        .collect();

    // Sort by duration ascending so the last `top_k_activity` are the longest.
    valid_acts.sort_by_key(|a| a.duration());

    // Keep top-k (the ones at the *end* of the sorted list are longest).
    let selected_acts: Vec<&&ActivityEvent> = valid_acts
        .iter()
        .rev()
        .take(top_k_activity)
        .collect();

    let mut act_sentences = Vec::new();
    for &&act in &selected_acts {
        act_sentences.push(describe_activity(&act.name, act.start_minute, act.end_minute, rng));
    }
    if !act_sentences.is_empty() {
        parts.push(act_sentences.join(" "));
    }

    // -----------------------------------------------------------------------
    // Sleep
    // -----------------------------------------------------------------------
    let mut sleep_sorted: Vec<&SleepEvent> = sleep.iter().collect();
    sleep_sorted.sort_by_key(|s| s.end_minute.saturating_sub(s.start_minute));

    let mut sleep_sentences = Vec::new();
    for &s in sleep_sorted.iter().rev().take(top_k_sleep) {
        sleep_sentences.push(describe_activity(
            "Sleep",
            s.start_minute,
            s.end_minute,
            rng,
        ));
    }
    if !sleep_sentences.is_empty() {
        parts.push(sleep_sentences.join(" "));
    }

    // -----------------------------------------------------------------------
    // Mood
    // -----------------------------------------------------------------------
    for m in moods {
        parts.push(describe_mood(&m.mood, m.minute, rng));
    }

    parts.join("\n")
}

// ---------------------------------------------------------------------------
// Helpers
// ---------------------------------------------------------------------------

fn describe_activity<R: Rng>(
    activity: &str,
    start_minute: usize,
    end_minute: usize,
    rng: &mut R,
) -> String {
    let tmpl = ACTIVITY_TEMPLATES
        .choose(rng)
        .copied()
        .unwrap_or(ACTIVITY_TEMPLATES[0]);
    tmpl.replace("{activity}", activity)
        .replace("{start_minute}", &start_minute.to_string())
        .replace("{end_minute}", &end_minute.to_string())
}

fn describe_mood<R: Rng>(mood: &str, minute: usize, rng: &mut R) -> String {
    let tmpl = MOOD_TEMPLATES
        .choose(rng)
        .copied()
        .unwrap_or(MOOD_TEMPLATES[0]);
    tmpl.replace("{mood}", mood)
        .replace("{time}", &minute.to_string())
}

/// Parse a `"HH:MM:SS"` or ISO-8601 time string into minutes from midnight.
///
/// Returns `None` if parsing fails.
pub fn parse_minutes(time_str: &str) -> Option<usize> {
    // Try "HH:MM:SS"
    let parts: Vec<&str> = time_str.trim_start_matches(|c: char| !c.is_ascii_digit())
        .split(':')
        .collect();
    if parts.len() >= 2 {
        let h: usize = parts[0].parse().ok()?;
        let m: usize = parts[1].parse().ok()?;
        return Some(h * 60 + m);
    }
    None
}

#[cfg(test)]
mod tests {
    use super::*;
    use rand::{rngs::StdRng, SeedableRng};

    #[test]
    fn test_semantic_caption_basic() {
        let mut rng = StdRng::seed_from_u64(1);
        let activities = vec![
            ActivityEvent { name: "Walking".into(), start_minute: 480, end_minute: 540 },
            ActivityEvent { name: "Running".into(), start_minute: 600, end_minute: 660 },
        ];
        let sleep = vec![SleepEvent { start_minute: 0, end_minute: 440 }];
        let moods = vec![MoodEvent { mood: "calm".into(), minute: 300 }];
        let cap = generate_semantic_caption(&activities, &sleep, &moods, 8, 2, 20, &mut rng);
        assert!(cap.contains("Walking") || cap.contains("Running"),
            "Caption should mention activities: {cap}");
        assert!(cap.contains("Sleep"), "Caption should mention sleep: {cap}");
        assert!(cap.contains("calm"), "Caption should mention mood: {cap}");
    }

    #[test]
    fn test_min_duration_filter() {
        let mut rng = StdRng::seed_from_u64(2);
        let activities = vec![
            ActivityEvent { name: "Short".into(), start_minute: 0, end_minute: 10 }, // 11 min < 20
            ActivityEvent { name: "Long".into(),  start_minute: 0, end_minute: 60 }, // 61 min > 20
        ];
        let cap = generate_semantic_caption(&activities, &[], &[], 8, 2, 20, &mut rng);
        assert!(!cap.contains("Short"), "Short activity should be filtered");
        assert!(cap.contains("Long"), "Long activity should appear");
    }

    #[test]
    fn test_parse_minutes() {
        assert_eq!(parse_minutes("08:30:00"), Some(510));
        assert_eq!(parse_minutes("00:00:00"), Some(0));
    }
}