kcode-speaker-v3-schema 0.1.0

use serde::{Deserialize, Deserializer, Serialize, Serializer};
use std::{collections::BTreeSet, error::Error, fmt, str::FromStr};

pub const MAX_AUDIO_DURATION_MS: u64 = 150_000;
pub const OGG_MEDIA_TYPE: &str = "audio/ogg";
pub const FEATURE_SCHEMA_REVISION: &str = "speaker-v3-features-24-r1";

pub const FEATURE_NAMES: [&str; 24] = [
    "median_f0_hz",
    "high_front_vowel_f1_hz",
    "high_back_vowel_f2_hz",
    "spectral_tilt_db_per_octave",
    "cepstral_peak_prominence_db",
    "foreign_accentedness_1_to_9",
    "dominant_rhotic_realization",
    "unstressed_vowel_reduction_percent",
    "high_front_vowel_f2_hz",
    "low_vowel_f1_hz",
    "h1_minus_h2_db",
    "rhotic_f3_minus_f2_hz",
    "word_initial_t_vot_ms",
    "dominant_lateral_realization",
    "monophthongization_percent",
    "vocal_gender_presentation",
    "low_vowel_f2_hz",
    "high_back_vowel_f1_hz",
    "mean_formant_dispersion_hz",
    "creaky_phonation_percent",
    "hypernasality_0_to_4",
    "sibilant_center_of_gravity_hz",
    "consonant_cluster_reduction_percent",
    "perceived_vocal_age_years",
];

const NUMERIC_FEATURE_NAMES: [&str; 22] = [
    "median_f0_hz",
    "high_front_vowel_f1_hz",
    "high_back_vowel_f2_hz",
    "spectral_tilt_db_per_octave",
    "cepstral_peak_prominence_db",
    "foreign_accentedness_1_to_9",
    "unstressed_vowel_reduction_percent",
    "high_front_vowel_f2_hz",
    "low_vowel_f1_hz",
    "h1_minus_h2_db",
    "rhotic_f3_minus_f2_hz",
    "word_initial_t_vot_ms",
    "monophthongization_percent",
    "vocal_gender_presentation",
    "low_vowel_f2_hz",
    "high_back_vowel_f1_hz",
    "mean_formant_dispersion_hz",
    "creaky_phonation_percent",
    "hypernasality_0_to_4",
    "sibilant_center_of_gravity_hz",
    "consonant_cluster_reduction_percent",
    "perceived_vocal_age_years",
];

#[derive(Debug, Clone, PartialEq, Eq)]
pub enum ValidationError {
    InvalidOgg,
    InvalidDuration(u64),
    ByteLengthOverflow,
    Blank(&'static str),
    InvalidSpeakerLabel(String),
    DuplicateSpeakerLabel(LocalSpeakerLabel),
    NonFiniteFeature(&'static str),
}

impl fmt::Display for ValidationError {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::InvalidOgg => formatter.write_str("audio is not a complete Ogg first page"),
            Self::InvalidDuration(value) => write!(formatter, "invalid audio duration: {value} ms"),
            Self::ByteLengthOverflow => formatter.write_str("audio byte length exceeds u64"),
            Self::Blank(field) => write!(formatter, "{field} is blank"),
            Self::InvalidSpeakerLabel(value) => write!(formatter, "invalid speaker label: {value}"),
            Self::DuplicateSpeakerLabel(value) => {
                write!(formatter, "duplicate speaker label: {value}")
            }
            Self::NonFiniteFeature(field) => write!(formatter, "{field} is not finite"),
        }
    }
}

impl Error for ValidationError {}

#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct OggAudioMetadata {
    duration_ms: u64,
    byte_length: u64,
    filename: Option<String>,
}

impl OggAudioMetadata {
    pub fn from_bytes(
        bytes: &[u8],
        duration_ms: u64,
        filename: Option<String>,
    ) -> Result<Self, ValidationError> {
        validate_duration(duration_ms)?;
        validate_optional_text(filename.as_deref(), "filename")?;
        if bytes.len() < 27 || &bytes[..4] != b"OggS" || bytes[4] != 0 {
            return Err(ValidationError::InvalidOgg);
        }
        let body_start = 27 + bytes[26] as usize;
        if bytes.len() < body_start {
            return Err(ValidationError::InvalidOgg);
        }
        let body_length: usize = bytes[27..body_start]
            .iter()
            .map(|value| *value as usize)
            .sum();
        if bytes.len() < body_start + body_length {
            return Err(ValidationError::InvalidOgg);
        }
        Ok(Self {
            duration_ms,
            byte_length: u64::try_from(bytes.len())
                .map_err(|_| ValidationError::ByteLengthOverflow)?,
            filename,
        })
    }

    pub fn validate(&self) -> Result<(), ValidationError> {
        validate_duration(self.duration_ms)?;
        validate_optional_text(self.filename.as_deref(), "filename")?;
        (self.byte_length >= 27)
            .then_some(())
            .ok_or(ValidationError::InvalidOgg)
    }

    pub fn duration_ms(&self) -> u64 {
        self.duration_ms
    }

    pub fn byte_length(&self) -> u64 {
        self.byte_length
    }

    pub fn filename(&self) -> Option<&str> {
        self.filename.as_deref()
    }

    pub fn media_type(&self) -> &'static str {
        OGG_MEDIA_TYPE
    }
}

fn validate_duration(duration_ms: u64) -> Result<(), ValidationError> {
    (1..=MAX_AUDIO_DURATION_MS)
        .contains(&duration_ms)
        .then_some(())
        .ok_or(ValidationError::InvalidDuration(duration_ms))
}

fn validate_optional_text(value: Option<&str>, field: &'static str) -> Result<(), ValidationError> {
    if value.is_some_and(|text| text.trim().is_empty()) {
        return Err(ValidationError::Blank(field));
    }
    Ok(())
}

fn validate_text(value: &str, field: &'static str) -> Result<(), ValidationError> {
    (!value.trim().is_empty())
        .then_some(())
        .ok_or(ValidationError::Blank(field))
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct LocalSpeakerLabel(u32);

impl LocalSpeakerLabel {
    pub fn new(number: u32) -> Result<Self, ValidationError> {
        (number > 0)
            .then_some(Self(number))
            .ok_or_else(|| ValidationError::InvalidSpeakerLabel("Speaker 0".into()))
    }

    pub fn number(self) -> u32 {
        self.0
    }
}

impl fmt::Display for LocalSpeakerLabel {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(formatter, "Speaker {}", self.0)
    }
}

impl FromStr for LocalSpeakerLabel {
    type Err = ValidationError;

    fn from_str(value: &str) -> Result<Self, Self::Err> {
        let number = value
            .strip_prefix("Speaker ")
            .and_then(|value| value.parse::<u32>().ok())
            .filter(|value| *value > 0)
            .ok_or_else(|| ValidationError::InvalidSpeakerLabel(value.into()))?;
        Ok(Self(number))
    }
}

impl Serialize for LocalSpeakerLabel {
    fn serialize<S: Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
        serializer.serialize_str(&self.to_string())
    }
}

impl<'de> Deserialize<'de> for LocalSpeakerLabel {
    fn deserialize<D: Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
        String::deserialize(deserializer)?
            .parse()
            .map_err(serde::de::Error::custom)
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum VocalGenderPresentation {
    StronglyFeminine,
    Feminine,
    Androgynous,
    Masculine,
    StronglyMasculine,
}

impl VocalGenderPresentation {
    pub fn numeric_value(self) -> f64 {
        match self {
            Self::StronglyFeminine => -2.0,
            Self::Feminine => -1.0,
            Self::Androgynous => 0.0,
            Self::Masculine => 1.0,
            Self::StronglyMasculine => 2.0,
        }
    }
}

#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, Default)]
pub struct FeatureVector24 {
    pub median_f0_hz: Option<f64>,
    pub high_front_vowel_f1_hz: Option<f64>,
    pub high_back_vowel_f2_hz: Option<f64>,
    pub spectral_tilt_db_per_octave: Option<f64>,
    pub cepstral_peak_prominence_db: Option<f64>,
    pub foreign_accentedness_1_to_9: Option<f64>,
    pub dominant_rhotic_realization: Option<String>,
    pub unstressed_vowel_reduction_percent: Option<f64>,
    pub high_front_vowel_f2_hz: Option<f64>,
    pub low_vowel_f1_hz: Option<f64>,
    pub h1_minus_h2_db: Option<f64>,
    pub rhotic_f3_minus_f2_hz: Option<f64>,
    pub word_initial_t_vot_ms: Option<f64>,
    pub dominant_lateral_realization: Option<String>,
    pub monophthongization_percent: Option<f64>,
    pub vocal_gender_presentation: Option<VocalGenderPresentation>,
    pub low_vowel_f2_hz: Option<f64>,
    pub high_back_vowel_f1_hz: Option<f64>,
    pub mean_formant_dispersion_hz: Option<f64>,
    pub creaky_phonation_percent: Option<f64>,
    pub hypernasality_0_to_4: Option<f64>,
    pub sibilant_center_of_gravity_hz: Option<f64>,
    pub consonant_cluster_reduction_percent: Option<f64>,
    pub perceived_vocal_age_years: Option<f64>,
}

impl FeatureVector24 {
    pub fn validate(&self) -> Result<(), ValidationError> {
        for (name, value) in NUMERIC_FEATURE_NAMES.into_iter().zip(self.numeric_values()) {
            if value.is_some_and(|number| !number.is_finite()) {
                return Err(ValidationError::NonFiniteFeature(name));
            }
        }
        validate_optional_text(
            self.dominant_rhotic_realization.as_deref(),
            "dominant_rhotic_realization",
        )?;
        validate_optional_text(
            self.dominant_lateral_realization.as_deref(),
            "dominant_lateral_realization",
        )
    }

    pub fn numeric_values(&self) -> [Option<f64>; 22] {
        [
            self.median_f0_hz,
            self.high_front_vowel_f1_hz,
            self.high_back_vowel_f2_hz,
            self.spectral_tilt_db_per_octave,
            self.cepstral_peak_prominence_db,
            self.foreign_accentedness_1_to_9,
            self.unstressed_vowel_reduction_percent,
            self.high_front_vowel_f2_hz,
            self.low_vowel_f1_hz,
            self.h1_minus_h2_db,
            self.rhotic_f3_minus_f2_hz,
            self.word_initial_t_vot_ms,
            self.monophthongization_percent,
            self.vocal_gender_presentation
                .map(VocalGenderPresentation::numeric_value),
            self.low_vowel_f2_hz,
            self.high_back_vowel_f1_hz,
            self.mean_formant_dispersion_hz,
            self.creaky_phonation_percent,
            self.hypernasality_0_to_4,
            self.sibilant_center_of_gravity_hz,
            self.consonant_cluster_reduction_percent,
            self.perceived_vocal_age_years,
        ]
    }

    pub fn nominal_values(&self) -> [Option<&str>; 2] {
        [
            self.dominant_rhotic_realization.as_deref(),
            self.dominant_lateral_realization.as_deref(),
        ]
    }

    pub fn present_feature_count(&self) -> u8 {
        let numeric = self
            .numeric_values()
            .into_iter()
            .filter(Option::is_some)
            .count();
        let nominal = self
            .nominal_values()
            .into_iter()
            .filter(Option::is_some)
            .count();
        (numeric + nominal) as u8
    }
}

#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct StructuredSpeaker {
    pub speaker: LocalSpeakerLabel,
    pub language: String,
    pub features: FeatureVector24,
    pub features_usable_for_training: bool,
}

impl StructuredSpeaker {
    pub fn validate(&self) -> Result<(), ValidationError> {
        validate_text(&self.language, "language")?;
        self.features.validate()
    }
}

#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct StructuredAnalysis {
    pub transcript: String,
    pub speakers: Vec<StructuredSpeaker>,
}

impl StructuredAnalysis {
    pub fn validate(&self) -> Result<(), ValidationError> {
        validate_text(&self.transcript, "transcript")?;
        let mut labels = BTreeSet::new();
        for speaker in &self.speakers {
            speaker.validate()?;
            if !labels.insert(speaker.speaker) {
                return Err(ValidationError::DuplicateSpeakerLabel(speaker.speaker));
            }
        }
        Ok(())
    }
}

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

    fn ogg(body_length: u8) -> Vec<u8> {
        let mut bytes = vec![0; 28 + body_length as usize];
        bytes[..4].copy_from_slice(b"OggS");
        bytes[4] = 0;
        bytes[26] = 1;
        bytes[27] = body_length;
        bytes
    }

    fn speaker(number: u32) -> StructuredSpeaker {
        StructuredSpeaker {
            speaker: LocalSpeakerLabel::new(number).unwrap(),
            language: "English".into(),
            features: FeatureVector24::default(),
            features_usable_for_training: true,
        }
    }

    #[test]
    fn ogg_metadata_checks_page_and_duration() {
        let bytes = ogg(3);
        let metadata =
            OggAudioMetadata::from_bytes(&bytes, MAX_AUDIO_DURATION_MS, Some("voice.ogg".into()))
                .unwrap();
        assert_eq!(metadata.media_type(), OGG_MEDIA_TYPE);
        assert_eq!(metadata.byte_length(), bytes.len() as u64);
        assert_eq!(metadata.filename(), Some("voice.ogg"));
        assert_eq!(
            OggAudioMetadata::from_bytes(&bytes, 0, None),
            Err(ValidationError::InvalidDuration(0))
        );
        assert_eq!(
            OggAudioMetadata::from_bytes(&bytes[..bytes.len() - 1], 1, None),
            Err(ValidationError::InvalidOgg)
        );
        let mut wrong = bytes;
        wrong[4] = 1;
        assert_eq!(
            OggAudioMetadata::from_bytes(&wrong, 1, None),
            Err(ValidationError::InvalidOgg)
        );
    }

    #[test]
    fn speaker_labels_have_one_exact_form() {
        let label = LocalSpeakerLabel::new(12).unwrap();
        assert_eq!(label.to_string(), "Speaker 12");
        assert_eq!("Speaker 12".parse(), Ok(label));
        assert!("speaker 12".parse::<LocalSpeakerLabel>().is_err());
        assert_eq!(serde_json::to_string(&label).unwrap(), "\"Speaker 12\"");
        assert_eq!(
            serde_json::from_str::<LocalSpeakerLabel>("\"Speaker 12\"").unwrap(),
            label
        );
    }

    #[test]
    fn features_validate_and_project_in_frozen_order() {
        let features = FeatureVector24 {
            median_f0_hz: Some(100.0),
            dominant_rhotic_realization: Some("tap".into()),
            vocal_gender_presentation: Some(VocalGenderPresentation::Masculine),
            perceived_vocal_age_years: Some(30.0),
            ..FeatureVector24::default()
        };
        assert_eq!(features.numeric_values()[0], Some(100.0));
        assert_eq!(features.numeric_values()[13], Some(1.0));
        assert_eq!(features.numeric_values()[21], Some(30.0));
        assert_eq!(features.nominal_values(), [Some("tap"), None]);
        assert_eq!(features.present_feature_count(), 4);
        assert!(features.validate().is_ok());

        let invalid = FeatureVector24 {
            hypernasality_0_to_4: Some(f64::NAN),
            ..FeatureVector24::default()
        };
        assert_eq!(
            invalid.validate(),
            Err(ValidationError::NonFiniteFeature("hypernasality_0_to_4"))
        );

        let blank = FeatureVector24 {
            dominant_lateral_realization: Some(" ".into()),
            ..FeatureVector24::default()
        };
        assert_eq!(
            blank.validate(),
            Err(ValidationError::Blank("dominant_lateral_realization"))
        );
    }

    #[test]
    fn structured_analysis_validates_and_round_trips() {
        let analysis = StructuredAnalysis {
            transcript: "[high] Speaker 1: hello world".into(),
            speakers: vec![speaker(1)],
        };
        analysis.validate().unwrap();
        let encoded = serde_json::to_vec(&analysis).unwrap();
        let decoded: StructuredAnalysis = serde_json::from_slice(&encoded).unwrap();
        assert_eq!(decoded, analysis);
        assert!(decoded.speakers[0].features_usable_for_training);

        let duplicate = StructuredAnalysis {
            transcript: "speech".into(),
            speakers: vec![speaker(1), speaker(1)],
        };
        assert!(matches!(
            duplicate.validate(),
            Err(ValidationError::DuplicateSpeakerLabel(_))
        ));
    }
}