Graph basedgraph

5 files changed, 501 insertions, 0 deletions

diff --git a/src/config.rs b/src/config.rs
new file mode 100644
index 0000000..0be1d09
--- /dev/null
+++ b/src/config.rs
@@ -0,0 +1,48 @@
+use directories::{ProjectDirs, UserDirs};
+use serde::{Deserialize, Serialize};
+
+#[derive(Serialize, Deserialize)]
+pub struct Config {
+    pub music_roots: Vec<String>,
+    pub audio_exts: Vec<String>,
+    pub skip_window_secs: u64,
+    pub learning_rate: f32,
+    pub decay_rate: f32,
+    pub min_score_thresh: f32,
+    pub temperature: f32,
+}
+
+pub fn default_state_path() -> String {
+    ProjectDirs::from("qualifier", "organisation", "mood")
+        .and_then(|pd| {
+            pd.config_dir()
+                .join("mood.ron")
+                .to_str()
+                .map(|x| x.to_string())
+        })
+        .unwrap_or("~/.config/mood/mood.ron".to_string())
+}
+
+impl Default for Config {
+    fn default() -> Config {
+        Config {
+            music_roots: vec![
+                UserDirs::new()
+                    .and_then(|u| {
+                        u.audio_dir()
+                            .and_then(|p| p.to_str().map(|x| x.to_string()))
+                    })
+                    .unwrap_or("~/Music".to_string()),
+            ],
+            audio_exts: [".mp3", ".flac", ".wav", ".m4a", ".ogg", ".vorbis"]
+                .into_iter()
+                .map(String::from)
+                .collect(),
+            skip_window_secs: 60,
+            learning_rate: 0.1,
+            decay_rate: 0.001,
+            min_score_thresh: 1e-5,
+            temperature: 0.6,
+        }
+    }
+}
diff --git a/src/learner.rs b/src/learner.rs
new file mode 100644
index 0000000..91b24fe
--- /dev/null
+++ b/src/learner.rs
@@ -0,0 +1,161 @@
+use std::collections::{HashMap, HashSet};
+
+use rand::{self, seq::IndexedRandom};
+use serde::{Deserialize, Serialize};
+
+use crate::{config::Config, trajectory::Action};
+
+fn key(a: &str, b: &str) -> (String, String) {
+    let a = a.to_string();
+    let b = b.to_string();
+    if a < b { (a, b) } else { (b, a) }
+}
+
+fn compute_sum(
+    cand: &str,
+    known_keys: &HashSet<&String>,
+    map: &HashMap<(String, String), f32>,
+) -> f32 {
+    known_keys
+        .iter()
+        .filter_map(|k| map.get(&key(k, cand)))
+        .sum()
+}
+
+fn softmax_sample(items: &[(&String, f32)], temperature: f32) -> Option<(String, f32)> {
+    if items.is_empty() {
+        return None;
+    }
+    let mut rng = rand::rng();
+
+    let max_val = items
+        .iter()
+        .map(|(_, v)| *v)
+        .fold(f32::NEG_INFINITY, f32::max);
+
+    let exp_values: Vec<_> = items
+        .iter()
+        .map(|&(k, val)| (k, ((val - max_val) / temperature).exp()))
+        .collect();
+
+    let sum: f32 = exp_values.iter().map(|(_, v)| *v).sum();
+
+    let norm_values: Vec<_> = exp_values.iter().map(|&(k, v)| (k, v / sum)).collect();
+
+    norm_values
+        .choose_weighted(&mut rng, |item| item.1)
+        .map(|p| (p.0.clone(), p.1))
+        .ok()
+}
+
+#[derive(Serialize, Deserialize, Default)]
+pub struct Learner {
+    similar: HashMap<(String, String), f32>,
+    different: HashMap<(String, String), f32>,
+}
+
+#[derive(Debug)]
+pub enum Learning {
+    SkipExtend,
+    MoreExtend(Vec<String>, String),
+    SkipToMore(Vec<String>, String),
+    MoreToSkip(Vec<String>, String),
+}
+
+impl Learner {
+    pub fn prune(&mut self, valid: &HashSet<String>) {
+        self.similar
+            .retain(|(a, b), _v| valid.contains(a) && valid.contains(b));
+        self.different
+            .retain(|(a, b), _v| valid.contains(a) && valid.contains(b));
+    }
+
+    pub fn learn(&mut self, config: &Config, learning: &Learning) {
+        // global decay
+        for (_k, v) in self.similar.iter_mut().chain(self.different.iter_mut()) {
+            *v *= 1.0 - config.decay_rate;
+        }
+        // update
+        match learning {
+            Learning::SkipExtend => {
+                // This is the least information carrying case, it does not
+                // follow that the trajectory contains similar or different
+                // items, we may simply be seeking something in particular.
+            }
+            Learning::MoreExtend(trajectory, new) => {
+                // We're continuing a good run so `new` is similar to everything
+                // in `trajectory`, but to account for the possibilty that our
+                // mood has changed over the course of this streak we damp
+                // sub-linearly that update by a proxy of temporal distance.
+                for (distance, e) in trajectory.iter().rev().enumerate() {
+                    let damp = ((distance + 1) as f32).powf(-0.5);
+                    let v = self.similar.entry(key(e, new)).or_default();
+                    *v = (1.0 - config.learning_rate) * *v + config.learning_rate * damp;
+                }
+            }
+            Learning::SkipToMore(trajectory, new) => {
+                // We have learnt that `new` is different to everything in
+                // `trajectory`, the strongest signal we have.
+                for e in trajectory {
+                    let v = self.different.entry(key(e, new)).or_default();
+                    *v = (1.0 - config.learning_rate) * *v + config.learning_rate;
+                }
+            }
+            Learning::MoreToSkip(trajectory, new) => {
+                // `new` could be different to everything in `trajectory`, or we
+                // simply changed our minds, so we have only weak evidence of
+                // difference. The positive coherence of trajectory was taken
+                // care of during MoreExtend above.
+                let damp = (trajectory.len() + 1) as f32;
+                for e in trajectory {
+                    let v = self.different.entry(key(e, new)).or_default();
+                    *v = (1.0 - config.learning_rate) * *v + config.learning_rate / damp;
+                }
+            }
+        };
+        // global threshold drop
+        self.similar = self
+            .similar
+            .extract_if(|_k, v| (*v).abs() > config.min_score_thresh)
+            .collect();
+        self.different = self
+            .different
+            .extract_if(|_k, v| (*v).abs() > config.min_score_thresh)
+            .collect();
+    }
+
+    pub fn sample(
+        &self,
+        trajectory: &[String],
+        action: &Action,
+        candidates: &HashSet<String>,
+        temperature: f32,
+    ) -> Option<(String, f32)> {
+        let trajectory: HashSet<_> = trajectory.iter().collect();
+        let normalisation: f32 = f32::max(trajectory.len() as f32, 1.0);
+
+        let candidates: Vec<_> = candidates
+            .iter()
+            .filter(|c| !trajectory.contains(c))
+            .collect();
+
+        if candidates.is_empty() {
+            return None;
+        }
+
+        let items = candidates
+            .iter()
+            .map(|&c| {
+                let sim = compute_sum(c, &trajectory, &self.similar);
+                let dif = compute_sum(c, &trajectory, &self.different);
+                let score = match action {
+                    Action::Skip => dif - sim,
+                    Action::More => sim - dif,
+                };
+                (c, score / normalisation)
+            })
+            .collect::<Vec<_>>();
+
+        softmax_sample(&items, temperature)
+    }
+}
diff --git a/src/main.rs b/src/main.rs
new file mode 100644
index 0000000..74f90c8
--- /dev/null
+++ b/src/main.rs
@@ -0,0 +1,61 @@
+use std::process::Command;
+
+mod config;
+mod learner;
+mod state;
+mod trajectory;
+
+use crate::{config::default_state_path, state::State};
+
+fn mpc_load_and_play(path: &str) {
+    _ = Command::new("mpc")
+        .arg("clear")
+        .output()
+        .expect("failed to execute mpc clear");
+
+    _ = Command::new("mpc")
+        .args(["add", path])
+        .output()
+        .expect("failed to execute mpc add");
+
+    _ = Command::new("mpc")
+        .arg("play")
+        .output()
+        .expect("failed to execute mpc play");
+}
+
+fn notify(text: &str) {
+    if let Err(e) = Command::new("notify-send")
+        .args(["-t", "5000", "-a", "mood", &text])
+        .status()
+    {
+        eprintln!("Failed to show notification: {e}");
+    }
+}
+
+fn main() {
+    let state = State::load(&default_state_path());
+    let mut state = match state {
+        Err(e) => {
+            eprintln!("Failed to load state: {e}");
+            State::default()
+        }
+        Ok(s) => s,
+    };
+
+    let ur = state.update_albums();
+    println!("Update result: {ur}");
+
+    if let Some((next_album, prob)) = state.next() {
+        let text = format!("Now playing: {next_album} ({:.2}%)", prob * 100.0);
+        println!("{text}");
+        mpc_load_and_play(&next_album);
+        notify(&text);
+    } else {
+        println!("No next album available.");
+    }
+
+    if let Err(e) = state.try_save() {
+        eprintln!("Failed to save state: {e}");
+    }
+}
diff --git a/src/state.rs b/src/state.rs
new file mode 100644
index 0000000..b9a6618
--- /dev/null
+++ b/src/state.rs
@@ -0,0 +1,133 @@
+use std::{collections::HashSet, fs};
+
+use derive_more::Display;
+use expanduser::expanduser;
+use serde::{Deserialize, Serialize};
+use walkdir::WalkDir;
+
+use crate::{
+    config::{Config, default_state_path},
+    learner::Learner,
+    trajectory::Trajectory,
+};
+
+#[derive(Serialize, Deserialize, Default)]
+pub struct State {
+    pub config: Config,
+    trajectory: Trajectory,
+    learner: Learner,
+    albums: HashSet<String>,
+}
+
+impl State {
+    pub fn next(&mut self) -> Option<(String, f32)> {
+        let learning = self.trajectory.step(self.config.skip_window_secs);
+        if let Some(ref learning) = learning {
+            println!("Obtained learning: {learning:?}");
+            self.learner.learn(&self.config, learning);
+        }
+        if let Some((new_album, prob)) = self.learner.sample(
+            self.trajectory.slice(),
+            &learning.as_ref().into(),
+            &self.albums,
+            self.config.temperature,
+        ) {
+            self.trajectory.log(&new_album);
+            Some((new_album, prob))
+        } else {
+            None
+        }
+    }
+}
+
+#[derive(Display)]
+#[display("+{added}/-{removed}, total {new_total}")]
+pub struct UpdateResult {
+    added: usize,
+    removed: usize,
+    new_total: usize,
+}
+
+impl State {
+    pub fn load(path: &str) -> Result<State, String> {
+        expanduser(path)
+            .map_err(|e| e.to_string())
+            .and_then(|p| fs::read_to_string(p).map_err(|e| e.to_string()))
+            .and_then(|contents| ron::from_str(&contents).map_err(|e| e.to_string()))
+    }
+
+    pub fn try_save(&self) -> Result<(), String> {
+        let str_path = default_state_path();
+        let path = expanduser(&str_path).map_err(|e| e.to_string())?;
+        _ = path.parent().map(|parent| std::fs::create_dir_all(parent));
+        let serialised = ron::ser::to_string_pretty(self, ron::ser::PrettyConfig::default())
+            .map_err(|e| format!("Failed to serialize state: {}", e))?;
+        fs::write(&path, serialised)
+            .map_err(|e| format!("Failed to write state file at {}: {}", str_path, e))?;
+        Ok(())
+    }
+
+    pub fn update_albums(&mut self) -> UpdateResult {
+        let roots: Vec<_> = self
+            .config
+            .music_roots
+            .iter()
+            .filter_map(|x| expanduser(x).ok().filter(|p| p.is_dir()))
+            .collect();
+
+        let mut new_albums: HashSet<String> = HashSet::new();
+
+        for root_path in roots {
+            for entry in WalkDir::new(&root_path) {
+                let entry = match entry {
+                    Ok(e) if e.path().is_dir() => e,
+                    _ => continue,
+                };
+
+                if !self.contains_audio(&entry) {
+                    continue;
+                }
+
+                if let Some(rel_str) = entry
+                    .path()
+                    .strip_prefix(&root_path)
+                    .ok()
+                    .and_then(|p| p.to_str())
+                {
+                    new_albums.insert(rel_str.to_string());
+                }
+            }
+        }
+
+        let added = new_albums.difference(&self.albums).count();
+        let removed = self.albums.difference(&new_albums).count();
+        self.albums.clone_from(&new_albums);
+
+        self.trajectory.prune(&self.albums);
+        self.learner.prune(&self.albums);
+
+        UpdateResult {
+            added,
+            removed,
+            new_total: self.albums.len(),
+        }
+    }
+
+    fn contains_audio(&self, entry: &walkdir::DirEntry) -> bool {
+        fs::read_dir(entry.path()).ok().is_some_and(|mut entries| {
+            entries.any(|e| {
+                e.as_ref()
+                    .is_ok_and(|entry| self.is_audio_file(&entry.path()))
+            })
+        })
+    }
+
+    fn is_audio_file(&self, path: &std::path::Path) -> bool {
+        !path.is_dir()
+            && path
+                .extension()
+                .and_then(|ext| ext.to_str())
+                .map(|s| s.to_lowercase())
+                .is_some_and(|ext| self.config.audio_exts.contains(&format!(".{ext}")))
+    }
+}
diff --git a/src/trajectory.rs b/src/trajectory.rs
new file mode 100644
index 0000000..214e16c
--- /dev/null
+++ b/src/trajectory.rs
@@ -0,0 +1,98 @@
+use serde::{Deserialize, Serialize};
+use std::{collections::HashSet, time::SystemTime};
+
+use crate::learner::Learning;
+
+#[derive(Serialize, Deserialize, Debug)]
+struct LastData {
+    timestamp: SystemTime,
+    album: String,
+}
+
+#[derive(Serialize, Deserialize, Default)]
+pub struct Trajectory {
+    last_data: Option<LastData>,
+    streak_kind: Option<Action>,
+    history: Vec<String>,
+}
+
+#[derive(Serialize, Deserialize, Clone, PartialEq, Debug)]
+pub enum Action {
+    Skip,
+    More,
+}
+
+impl From<Option<&Learning>> for Action {
+    fn from(learning: Option<&Learning>) -> Action {
+        if let Some(ts) = learning {
+            match ts {
+                Learning::SkipExtend | Learning::MoreToSkip(_, _) => Action::Skip,
+                Learning::MoreExtend(_, _) | Learning::SkipToMore(_, _) => Action::More,
+            }
+        } else {
+            Action::Skip
+        }
+    }
+}
+
+impl Trajectory {
+    pub fn prune(&mut self, valid: &HashSet<String>) {
+        self.history.retain_mut(|x| valid.contains(x));
+
+        if let Some(ref last_data) = self.last_data
+            && !valid.contains(&last_data.album)
+        {
+            self.last_data = None
+        }
+    }
+
+    pub fn slice(&self) -> &[String] {
+        &self.history
+    }
+
+    pub fn log(&mut self, new_album: &str) {
+        self.last_data = Some(LastData {
+            timestamp: SystemTime::now(),
+            album: new_album.to_string(),
+        })
+    }
+
+    pub fn step(&mut self, skip_window_secs: u64) -> Option<Learning> {
+        if let Some(ref last_data) = self.last_data {
+            let action = if SystemTime::now()
+                .duration_since(last_data.timestamp)
+                .ok()
+                .is_some_and(|d| d.as_secs() < skip_window_secs)
+            {
+                Action::Skip
+            } else {
+                Action::More
+            };
+
+            let last_album = last_data.album.clone();
+            let current_streak = self.history.clone();
+
+            let learning = if let Some(ref kind) = self.streak_kind
+                && kind == &action
+            {
+                self.history.push(last_album.clone());
+                match action {
+                    Action::Skip => Learning::SkipExtend,
+                    Action::More => Learning::MoreExtend(current_streak, last_album),
+                }
+            } else {
+                self.history.clear();
+                self.history.push(last_album.clone());
+                match action {
+                    Action::Skip => Learning::MoreToSkip(current_streak, last_album),
+                    Action::More => Learning::SkipToMore(current_streak, last_album),
+                }
+            };
+
+            self.streak_kind = Some(action);
+            Some(learning)
+        } else {
+            None
+        }
+    }
+}