farmap 0.2.1

A library for working with Farcaster label datasets
Documentation
use crate::{user::InvalidInputError, utils::distribution_from_counts};
use chrono::NaiveDate;
use serde::{Deserialize, Serialize};

#[derive(Clone, Copy, Debug, PartialEq, Serialize, Deserialize)]
pub enum SpamScore {
    Zero,
    One,
    Two,
}

pub type SpamRecord = (SpamScore, NaiveDate);

impl TryFrom<usize> for SpamScore {
    type Error = InvalidInputError;

    fn try_from(value: usize) -> Result<Self, Self::Error> {
        match value {
            0 => Ok(Self::Zero),
            1 => Ok(Self::One),
            2 => Ok(Self::Two),
            _ => Err(InvalidInputError::SpamScoreError { label: value }),
        }
    }
}

#[derive(Serialize, Debug)]
pub struct SpamScoreCount {
    date: NaiveDate,
    nonspam: u64,
    maybe: u64,
    spam: u64,
}

impl SpamScoreCount {
    pub fn new(date: NaiveDate, spam_count: u64, maybe_count: u64, nonspam_count: u64) -> Self {
        Self {
            date,
            nonspam: nonspam_count,
            maybe: maybe_count,
            spam: spam_count,
        }
    }

    pub fn date(&self) -> NaiveDate {
        self.date
    }

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

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

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

    pub fn add(&mut self, score: &SpamScore) {
        match score {
            SpamScore::Zero => self.spam += 1,
            SpamScore::One => self.maybe += 1,
            SpamScore::Two => self.nonspam += 1,
        }
    }

    pub fn total(&self) -> u64 {
        self.spam + self.maybe + self.nonspam
    }

    pub fn distributions(&self) -> Option<[f32; 3]> {
        distribution_from_counts(&[self.spam, self.maybe, self.nonspam])
    }
}

#[derive(Serialize, Debug)]
pub struct SpamScoreDistribution {
    date: NaiveDate,
    nonspam: f64,
    maybe: f64,
    spam: f64,
}

impl SpamScoreDistribution {
    pub fn new(date: NaiveDate, spam: f64, maybe: f64, nonspam: f64) -> Result<Self, String> {
        let sum = spam + maybe + nonspam;
        if !(0.99..=1.01).contains(&sum) {
            Err("provided values are not a distribution".to_string())
        } else {
            Ok(Self {
                date,
                nonspam,
                maybe,
                spam,
            })
        }
    }

    pub fn date(&self) -> NaiveDate {
        self.date
    }

    pub fn spam(&self) -> f64 {
        self.spam
    }

    pub fn maybe_spam(&self) -> f64 {
        self.maybe
    }

    pub fn non_spam(&self) -> f64 {
        self.nonspam
    }
}

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

    pub fn basic_spam_score_count() -> SpamScoreCount {
        let date = NaiveDate::from_ymd_opt(2024, 1, 1).unwrap();
        SpamScoreCount::new(date, 100, 150, 200)
    }

    #[test]
    pub fn test_label_value_invalid_input() {
        assert!(SpamScore::try_from(0).is_ok());
        assert!(SpamScore::try_from(1).is_ok());
        assert!(SpamScore::try_from(2).is_ok());
        assert!(SpamScore::try_from(3).is_err());
        assert!(SpamScore::try_from(100).is_err());
    }

    #[test]
    pub fn basic_spam_score_count_test() {
        let count = basic_spam_score_count();
        assert_eq!(count.spam(), 100);
        assert_eq!(count.maybe_spam(), 150);
        assert_eq!(count.non_spam(), 200);
        assert_eq!(count.total(), 100 + 150 + 200);
    }
}