src/datasets.rs

//! Parsers for external emotion datasets mapped into DADOES labels.

use std::error::Error;
use std::fmt::{Display, Formatter};

use serde::Deserialize;
use serde::de::DeserializeOwned;

use crate::{ModelError, Mood, OwnedTrainingExample};

const DADOES_JSONL: &str = "dadoes-jsonl";
const EMPATHETIC_DIALOGUES: &str = "facebook/empathetic_dialogues";
const FIG_LONELINESS: &str = "FIG-Loneliness/FIG-Loneliness";
const LONELINESS_CAUSES: &str = "yael-katsman/Loneliness-Causes-and-Intensity";
const TEXT_EMOTION: &str = "smit18/text_emotion";
const PRAJWAL_TEXT_EMOTION: &str = "PrajwalNayaka/Text-Emotion";
const UM1NEKO_TEXT_EMOTION: &str = "Um1neko/text_emotion";

/// A recoverable external dataset parsing failure.
#[derive(Debug)]
pub enum DatasetParseError {
    /// A CSV file could not be decoded.
    Csv {
        /// Dataset source being parsed.
        dataset: &'static str,
        /// CSV parser failure.
        source: csv::Error,
    },
    /// A JSONL row could not be decoded.
    Json {
        /// Dataset source being parsed.
        dataset: &'static str,
        /// One-based input line number.
        line: usize,
        /// JSON parser failure.
        source: serde_json::Error,
    },
    /// A row contained a label not supported by that source mapping.
    UnknownLabel {
        /// Dataset source being parsed.
        dataset: &'static str,
        /// One-based input line number.
        line: usize,
        /// Raw label value.
        value: String,
    },
    /// A row did not contain enough information to produce a training example.
    InvalidRecord {
        /// Dataset source being parsed.
        dataset: &'static str,
        /// One-based input line number.
        line: usize,
        /// Stable reason string.
        reason: &'static str,
    },
    /// A mapped DADOES training example violated the local data model.
    InvalidExample {
        /// Dataset source being parsed.
        dataset: &'static str,
        /// One-based input line number.
        line: usize,
        /// Local data-model validation failure.
        source: ModelError,
    },
}

impl Display for DatasetParseError {
    fn fmt(&self, formatter: &mut Formatter<'_>) -> std::fmt::Result {
        match self {
            DatasetParseError::Csv { dataset, source } => {
                write!(formatter, "{dataset} CSV parse failed: {source}")
            }
            DatasetParseError::Json {
                dataset,
                line,
                source,
            } => write!(
                formatter,
                "{dataset} JSONL row {line} parse failed: {source}"
            ),
            DatasetParseError::UnknownLabel {
                dataset,
                line,
                value,
            } => write!(
                formatter,
                "{dataset} row {line} has unsupported label: {value}"
            ),
            DatasetParseError::InvalidRecord {
                dataset,
                line,
                reason,
            } => write!(formatter, "{dataset} row {line} is invalid: {reason}"),
            DatasetParseError::InvalidExample {
                dataset,
                line,
                source,
            } => write!(
                formatter,
                "{dataset} row {line} maps to an invalid DADOES example: {source}"
            ),
        }
    }
}

impl Error for DatasetParseError {}

/// Parses DADOES-native JSONL records.
///
/// Expected fields are `text`, `labels`, and optional `supervised_labels`.
/// `labels` may be empty for records that only supervise negative labels.
pub fn parse_dadoes_jsonl(input: &str) -> Result<Vec<OwnedTrainingExample>, DatasetParseError> {
    let records = parse_jsonl::<DadoesJsonlRecord>(input, DADOES_JSONL)?;
    let mut examples = Vec::new();

    for (line, record) in records {
        let labels = parse_mood_labels(DADOES_JSONL, line, &record.labels)?;
        let supervised = match record.supervised_labels {
            Some(labels) => parse_mood_labels(DADOES_JSONL, line, &labels)?,
            None => Mood::ALL.to_vec(),
        };
        examples.push(example(
            DADOES_JSONL,
            line,
            record.text,
            labels,
            supervised,
        )?);
    }

    Ok(examples)
}

/// Parses exported EmpatheticDialogues CSV rows.
///
/// The parser expects `context`, `prompt`, and `utterance` columns from the
/// Hugging Face dataset export.
pub fn parse_empathetic_dialogues_csv(
    input: &str,
) -> Result<Vec<OwnedTrainingExample>, DatasetParseError> {
    let records = parse_csv::<EmpatheticDialoguesRecord>(input, EMPATHETIC_DIALOGUES)?;
    let mut examples = Vec::new();

    for (line, record) in records {
        let Some(labels) = empathetic_moods(&record.context) else {
            return Err(DatasetParseError::UnknownLabel {
                dataset: EMPATHETIC_DIALOGUES,
                line,
                value: record.context,
            });
        };
        examples.push(example(
            EMPATHETIC_DIALOGUES,
            line,
            join_text_parts([record.prompt.as_str(), record.utterance.as_str()]),
            labels.to_vec(),
            empathetic_supervised_moods(),
        )?);
    }

    Ok(examples)
}

/// Parses FIG-Loneliness JSONL exported from the Hugging Face Arrow dataset.
///
/// The parser expects `text` and `lonely` fields. `lonely` may be a boolean, a
/// numeric class id, or a two-element `[non_lonely, lonely]` vector.
pub fn parse_fig_loneliness_jsonl(
    input: &str,
) -> Result<Vec<OwnedTrainingExample>, DatasetParseError> {
    let records = parse_jsonl::<FigLonelinessRecord>(input, FIG_LONELINESS)?;
    let mut examples = Vec::new();

    for (line, record) in records {
        let labels = if record.lonely.is_positive() {
            vec![Mood::Lonely]
        } else {
            Vec::new()
        };
        examples.push(example(
            FIG_LONELINESS,
            line,
            record.text,
            labels,
            vec![Mood::Lonely],
        )?);
    }

    Ok(examples)
}

/// Parses Loneliness Causes and Intensity CSV rows.
///
/// The parser expects `title`, `text`, `t1_label`, and `t2_label` columns.
pub fn parse_loneliness_causes_csv(
    input: &str,
) -> Result<Vec<OwnedTrainingExample>, DatasetParseError> {
    let records = parse_csv::<LonelinessCausesRecord>(input, LONELINESS_CAUSES)?;
    let mut examples = Vec::new();

    for (line, record) in records {
        let labels = if loneliness_causes_positive(&record) {
            vec![Mood::Lonely]
        } else {
            Vec::new()
        };
        examples.push(example(
            LONELINESS_CAUSES,
            line,
            join_text_parts([record.title.as_str(), record.text.as_str()]),
            labels,
            vec![Mood::Lonely],
        )?);
    }

    Ok(examples)
}

/// Parses the `smit18/text_emotion` CSV file.
///
/// The parser expects `content` and `sentiment` columns.
pub fn parse_text_emotion_csv(input: &str) -> Result<Vec<OwnedTrainingExample>, DatasetParseError> {
    let records = parse_csv::<TextEmotionRecord>(input, TEXT_EMOTION)?;
    let mut examples = Vec::new();

    for (line, record) in records {
        let Some(labels) = text_emotion_moods(&record.sentiment) else {
            return Err(DatasetParseError::UnknownLabel {
                dataset: TEXT_EMOTION,
                line,
                value: record.sentiment,
            });
        };
        examples.push(example(
            TEXT_EMOTION,
            line,
            record.content,
            labels.to_vec(),
            text_emotion_supervised_moods(),
        )?);
    }

    Ok(examples)
}

/// Parses the `PrajwalNayaka/Text-Emotion` CSV file.
///
/// The parser expects `text` and `emotion` columns.
pub fn parse_prajwal_text_emotion_csv(
    input: &str,
) -> Result<Vec<OwnedTrainingExample>, DatasetParseError> {
    let records = parse_csv::<PrajwalTextEmotionRecord>(input, PRAJWAL_TEXT_EMOTION)?;
    let mut examples = Vec::new();

    for (line, record) in records {
        let Some(labels) = prajwal_moods(&record.emotion) else {
            return Err(DatasetParseError::UnknownLabel {
                dataset: PRAJWAL_TEXT_EMOTION,
                line,
                value: record.emotion,
            });
        };
        examples.push(example(
            PRAJWAL_TEXT_EMOTION,
            line,
            record.text,
            labels.to_vec(),
            prajwal_supervised_moods(),
        )?);
    }

    Ok(examples)
}

/// Parses the `Um1neko/text_emotion` JSONL files.
///
/// The parser expects `text` and integer `labels` fields.
pub fn parse_um1neko_text_emotion_jsonl(
    input: &str,
) -> Result<Vec<OwnedTrainingExample>, DatasetParseError> {
    let records = parse_jsonl::<Um1nekoTextEmotionRecord>(input, UM1NEKO_TEXT_EMOTION)?;
    let mut examples = Vec::new();

    for (line, record) in records {
        let mut labels = Vec::new();
        for label in record.labels {
            let Some(mapped) = um1neko_moods(label) else {
                return Err(DatasetParseError::UnknownLabel {
                    dataset: UM1NEKO_TEXT_EMOTION,
                    line,
                    value: label.to_string(),
                });
            };
            push_unique_moods(&mut labels, mapped);
        }
        examples.push(example(
            UM1NEKO_TEXT_EMOTION,
            line,
            record.text,
            labels,
            um1neko_supervised_moods(),
        )?);
    }

    Ok(examples)
}

#[derive(Debug, Deserialize)]
struct DadoesJsonlRecord {
    text: String,
    labels: Vec<String>,
    supervised_labels: Option<Vec<String>>,
}

#[derive(Debug, Deserialize)]
struct EmpatheticDialoguesRecord {
    context: String,
    prompt: String,
    utterance: String,
}

#[derive(Debug, Deserialize)]
struct FigLonelinessRecord {
    text: String,
    lonely: BinaryLonelyValue,
}

#[derive(Debug, Deserialize)]
#[serde(untagged)]
enum BinaryLonelyValue {
    Bool(bool),
    Number(u8),
    Vector(Vec<u8>),
}

impl BinaryLonelyValue {
    fn is_positive(&self) -> bool {
        match self {
            BinaryLonelyValue::Bool(value) => *value,
            BinaryLonelyValue::Number(value) => *value != 0,
            BinaryLonelyValue::Vector(values) => {
                let non_lonely = values.first().copied().unwrap_or(0);
                let lonely = values.get(1).copied().unwrap_or(0);
                lonely > non_lonely || lonely > 0
            }
        }
    }
}

#[derive(Debug, Deserialize)]
struct LonelinessCausesRecord {
    title: String,
    text: String,
    t1_label: String,
    t2_label: f32,
}

#[derive(Debug, Deserialize)]
struct TextEmotionRecord {
    sentiment: String,
    content: String,
}

#[derive(Debug, Deserialize)]
struct PrajwalTextEmotionRecord {
    text: String,
    emotion: String,
}

#[derive(Debug, Deserialize)]
struct Um1nekoTextEmotionRecord {
    text: String,
    labels: Vec<usize>,
}

fn parse_csv<T>(input: &str, dataset: &'static str) -> Result<Vec<(usize, T)>, DatasetParseError>
where
    T: DeserializeOwned,
{
    let mut reader = csv::Reader::from_reader(input.as_bytes());
    let mut records = Vec::new();

    for (position, record) in reader.deserialize::<T>().enumerate() {
        let line = position + 2;
        let record = record.map_err(|source| DatasetParseError::Csv { dataset, source })?;
        records.push((line, record));
    }

    Ok(records)
}

fn parse_jsonl<T>(input: &str, dataset: &'static str) -> Result<Vec<(usize, T)>, DatasetParseError>
where
    T: DeserializeOwned,
{
    let mut records = Vec::new();
    for (position, row) in input.lines().enumerate() {
        let line = position + 1;
        if row.trim().is_empty() {
            continue;
        }
        let record = serde_json::from_str::<T>(row).map_err(|source| DatasetParseError::Json {
            dataset,
            line,
            source,
        })?;
        records.push((line, record));
    }

    Ok(records)
}

fn example(
    dataset: &'static str,
    line: usize,
    text: String,
    labels: Vec<Mood>,
    supervised: Vec<Mood>,
) -> Result<OwnedTrainingExample, DatasetParseError> {
    OwnedTrainingExample::new_with_supervised(text, labels, supervised).map_err(|source| {
        DatasetParseError::InvalidExample {
            dataset,
            line,
            source,
        }
    })
}

fn parse_mood_labels(
    dataset: &'static str,
    line: usize,
    labels: &[String],
) -> Result<Vec<Mood>, DatasetParseError> {
    let mut moods = Vec::new();
    for label in labels {
        let Some(mood) = Mood::from_label(label.trim()) else {
            return Err(DatasetParseError::UnknownLabel {
                dataset,
                line,
                value: label.to_owned(),
            });
        };
        push_unique_mood(&mut moods, mood);
    }
    Ok(moods)
}

fn join_text_parts<const N: usize>(parts: [&str; N]) -> String {
    parts
        .into_iter()
        .map(str::trim)
        .filter(|part| !part.is_empty())
        .collect::<Vec<_>>()
        .join("\n")
}

fn loneliness_causes_positive(record: &LonelinessCausesRecord) -> bool {
    !record.t1_label.to_ascii_lowercase().contains("not lonely") || record.t2_label >= 2.0
}

fn label_key(label: &str) -> String {
    label.trim().to_ascii_lowercase().replace([' ', '-'], "_")
}

fn empathetic_moods(label: &str) -> Option<&'static [Mood]> {
    match label_key(label).as_str() {
        "afraid" | "anxious" | "apprehensive" | "terrified" => Some(&[Mood::Anxious]),
        "angry" | "furious" | "disgusted" => Some(&[Mood::Angry]),
        "annoyed" => Some(&[Mood::Frustrated]),
        "anticipating" => Some(&[Mood::Excited, Mood::Hopeful]),
        "ashamed" | "embarrassed" | "guilty" => Some(&[Mood::Sad, Mood::Anxious]),
        "caring" | "content" | "grateful" | "impressed" | "proud" => {
            Some(&[Mood::Satisfied, Mood::Happy])
        }
        "confident" | "faithful" | "prepared" | "trusting" => {
            Some(&[Mood::Hopeful, Mood::Satisfied])
        }
        "devastated" | "nostalgic" | "sad" | "sentimental" => Some(&[Mood::Sad]),
        "disappointed" => Some(&[Mood::Sad, Mood::Frustrated]),
        "excited" | "surprised" => Some(&[Mood::Excited]),
        "jealous" => Some(&[Mood::Angry, Mood::Sad]),
        "joyful" => Some(&[Mood::Happy]),
        "lonely" => Some(&[Mood::Lonely]),
        "hopeful" => Some(&[Mood::Hopeful]),
        _ => None,
    }
}

fn empathetic_supervised_moods() -> Vec<Mood> {
    vec![
        Mood::Happy,
        Mood::Satisfied,
        Mood::Excited,
        Mood::Anxious,
        Mood::Frustrated,
        Mood::Sad,
        Mood::Angry,
        Mood::Lonely,
        Mood::Hopeful,
    ]
}

fn text_emotion_moods(label: &str) -> Option<&'static [Mood]> {
    match label_key(label).as_str() {
        "happiness" => Some(&[Mood::Happy]),
        "love" => Some(&[Mood::Happy, Mood::Satisfied]),
        "sadness" => Some(&[Mood::Sad]),
        "worry" => Some(&[Mood::Anxious]),
        "hate" => Some(&[Mood::Angry]),
        _ => None,
    }
}

fn text_emotion_supervised_moods() -> Vec<Mood> {
    vec![
        Mood::Happy,
        Mood::Satisfied,
        Mood::Anxious,
        Mood::Sad,
        Mood::Angry,
    ]
}

fn prajwal_moods(label: &str) -> Option<&'static [Mood]> {
    match label_key(label).as_str() {
        "neutral" => Some(&[Mood::Neutral]),
        "happiness" => Some(&[Mood::Happy]),
        "love" => Some(&[Mood::Happy, Mood::Satisfied]),
        "empty" | "sadness" => Some(&[Mood::Sad]),
        "hate" | "anger" => Some(&[Mood::Angry]),
        "enthusiasm" => Some(&[Mood::Excited]),
        "relief" => Some(&[Mood::Satisfied]),
        "fun" => Some(&[Mood::Happy, Mood::Excited]),
        "surprise" => Some(&[Mood::Excited]),
        _ => None,
    }
}

fn prajwal_supervised_moods() -> Vec<Mood> {
    vec![
        Mood::Happy,
        Mood::Satisfied,
        Mood::Excited,
        Mood::Sad,
        Mood::Angry,
        Mood::Neutral,
    ]
}

fn um1neko_moods(label: usize) -> Option<&'static [Mood]> {
    match label {
        0 => Some(&[Mood::Neutral]),
        1..=3 => Some(&[Mood::Angry]),
        4..=6 => Some(&[Mood::Hopeful]),
        7..=9 => Some(&[Mood::Angry]),
        10..=12 => Some(&[Mood::Anxious]),
        13..=15 => Some(&[Mood::Happy, Mood::Satisfied]),
        16..=18 => Some(&[Mood::Sad]),
        19..=21 => Some(&[Mood::Excited]),
        22..=24 => Some(&[Mood::Satisfied, Mood::Hopeful]),
        _ => None,
    }
}

fn um1neko_supervised_moods() -> Vec<Mood> {
    vec![
        Mood::Happy,
        Mood::Satisfied,
        Mood::Excited,
        Mood::Anxious,
        Mood::Sad,
        Mood::Angry,
        Mood::Hopeful,
        Mood::Neutral,
    ]
}

fn push_unique_moods(target: &mut Vec<Mood>, moods: &[Mood]) {
    for mood in moods {
        push_unique_mood(target, *mood);
    }
}

fn push_unique_mood(target: &mut Vec<Mood>, mood: Mood) {
    if !target.contains(&mood) {
        target.push(mood);
    }
}

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

    #[test]
    fn dadoes_jsonl_accepts_negative_partial_supervision() {
        let input = "{\"text\":\"I worked alone but felt calm.\",\"labels\":[],\"supervised_labels\":[\"lonely\"]}\n";
        let examples = parse_dadoes_jsonl(input);
        assert!(examples.is_ok());

        if let Ok(examples) = examples {
            assert_eq!(examples.len(), 1);
            let example = examples.first();
            assert!(example.is_some());
            if let Some(example) = example {
                assert!(example.labels().is_empty());
                assert_eq!(example.supervised_labels(), &[Mood::Lonely]);
            }
        }
    }

    #[test]
    fn loneliness_causes_maps_not_lonely_as_negative() {
        let input = concat!(
            "title,text,t1_label,t2_label\n",
            "quiet day,Nothing social happened,['Not lonely'],1.0\n"
        );
        let examples = parse_loneliness_causes_csv(input);
        assert!(examples.is_ok());

        if let Ok(examples) = examples {
            let example = examples.first();
            assert!(example.is_some());
            if let Some(example) = example {
                assert!(example.labels().is_empty());
                assert_eq!(example.supervised_labels(), &[Mood::Lonely]);
            }
        }
    }

    #[test]
    fn empathetic_dialogues_maps_lonely_context() {
        let input = concat!(
            "context,prompt,utterance\n",
            "lonely,I waited all day alone,Nobody came by\n"
        );
        let examples = parse_empathetic_dialogues_csv(input);
        assert!(examples.is_ok());

        if let Ok(examples) = examples {
            let example = examples.first();
            assert!(example.is_some());
            if let Some(example) = example {
                assert!(example.labels().contains(&Mood::Lonely));
            }
        }
    }

    #[test]
    fn um1neko_jsonl_maps_numeric_labels() {
        let input = "{\"text\":\"The result made me glad and secure.\",\"labels\":[13,22]}\n";
        let examples = parse_um1neko_text_emotion_jsonl(input);
        assert!(examples.is_ok());

        if let Ok(examples) = examples {
            let example = examples.first();
            assert!(example.is_some());
            if let Some(example) = example {
                assert!(example.labels().contains(&Mood::Happy));
                assert!(example.labels().contains(&Mood::Satisfied));
                assert!(example.labels().contains(&Mood::Hopeful));
            }
        }
    }
}