path: root/pkgs/by-name/mp/mpdpopm/src/dj/algorithms.rs

use std::{
    collections::HashSet,
    time::{Duration, SystemTime},
};

use anyhow::{Context, Result};
use rand::{Rng, distr};
use tracing::info;

use crate::{clients::Client, storage};

pub(crate) trait Algorithm {
    async fn next_track(&mut self, client: &mut Client) -> Result<String>;
}

/// Generates generic discovery playlist, that fulfills following requirements:
///  - Will (eventually) include every not-played song. (So it can be used to rank a library)
///  - Returns liked songs more often then not-played or negative songs.
pub struct Discovery {
    already_done: HashSet<String>,
    negative_chance: f64,
    neutral_chance: f64,
    positive_chance: f64,
}

impl Algorithm for Discovery {
    async fn next_track(&mut self, client: &mut Client) -> Result<String> {
        macro_rules! take {
            ($rng:expr, $from:expr) => {{
                info!(concat!(
                    "Trying to select a `",
                    stringify!($from),
                    "` track."
                ));

                assert!(!$from.is_empty());

                let normalized_weights = {
                    // We normalize the weights here, because negative values don't work for the
                    // distribution function we use below.
                    //    "-5" "-3" "1" "6" "19" | +5
                    // ->  "0"  "2" "6" "11" "24"
                    let mut weights = $from.iter().map(|(_, w)| *w).collect::<Vec<_>>();

                    weights.sort_by_key(|w| *w);

                    let first = *weights.first().expect(
                        "the value to exist, because we never run `take!` with an empty vector",
                    );

                    if first.is_negative() {
                        weights
                            .into_iter()
                            .rev()
                            .map(|w| w + first.abs())
                            .collect::<Vec<_>>()
                    } else {
                        weights
                    }
                };

                let sample = $rng.sample(
                    distr::weighted::WeightedIndex::new(normalized_weights.iter())
                        .expect("to be okay, because the weights are normalized"),
                );

                let output = $from.remove(sample);

                info!(
                    concat!(
                        "(",
                        stringify!($from),
                        ") Selected `{}` with weight: `{}` (normalized to `{}`)"
                    ),
                    output.0, output.1, normalized_weights[sample]
                );

                Ok::<_, anyhow::Error>(output)
            }};
        }

        let mut rng = rand::rng();
        let (mut positive, mut neutral, mut negative) = {
            let tracks = {
                let mut base = client
                    .get_all_songs()
                    .await?
                    .into_iter()
                    .filter(|song| !self.already_done.contains(song))
                    .collect::<Vec<_>>();

                if base.is_empty() {
                    // We could either have no tracks in the library,
                    // or we actually already listed to everything.
                    self.already_done = HashSet::new();

                    info!("Resetting already done songs, as we have no more to choose from");

                    base = client.get_all_songs().await?;
                }

                base
            };

            let mut sorted_tracks = Vec::with_capacity(tracks.len());
            for track in tracks {
                let weight = Self::weight_track(client, &track).await?;

                sorted_tracks.push((track, weight));
            }

            sorted_tracks.sort_by_key(|(_, weight)| *weight);

            let len = sorted_tracks.len() / 3;

            // We split the tracks into three thirds, so that we can also force a pick from e.g.
            // the lower third (the negative ones).
            let negative = sorted_tracks.drain(..len).collect::<Vec<_>>();
            let neutral = sorted_tracks.drain(..len).collect::<Vec<_>>();
            let positive = sorted_tracks;

            assert_eq!(negative.len(), neutral.len());

            (positive, neutral, negative)
        };

        let pick = rng.sample(
            distr::weighted::WeightedIndex::new(
                [
                    self.positive_chance,
                    self.neutral_chance,
                    self.negative_chance,
                ]
                .iter(),
            )
            .expect("to be valid, as hardcoded"),
        );

        let next = match pick {
            0 if !positive.is_empty() => take!(rng, positive),
            1 if !neutral.is_empty() => take!(rng, neutral),
            2 if !negative.is_empty() => take!(rng, negative),
            0..=2 => {
                // We couldn't actually satisfy the request, because we didn't have the required
                // track. So we just use the first non-empty one.
                if !positive.is_empty() {
                    take!(rng, positive)
                } else if !neutral.is_empty() {
                    take!(rng, neutral)
                } else if !negative.is_empty() {
                    take!(rng, negative)
                } else {
                    assert!(positive.is_empty() && neutral.is_empty() && negative.is_empty());
                    todo!("No songs available to select from, I don't know how to select one.");
                }
            }
            _ => unreachable!("These indexes are not possible"),
        }?;

        self.already_done.insert(next.0.to_owned());

        Ok(next.0)
    }
}

impl Discovery {
    pub(crate) fn new(positive_chance: f64, neutral_chance: f64, negative_chance: f64) -> Self {
        Self {
            already_done: HashSet::new(),
            positive_chance,
            neutral_chance,
            negative_chance,
        }
    }

    /// Calculate a recommendation score for a track.
    ///
    /// The algorithm maps tracks, that the user likes to a high score and songs that the user
    /// dislikes to a lower number.
    /// Currently, only the rating, skip count and play count are considered. Similarity scores,
    /// fetched from e.g. last.fm should be included in the future.
    pub async fn weight_track(client: &mut Client, track: &str) -> Result<i64> {
        let last_played_delta = {
            let last_played = storage::last_played::get(client, track).await?.unwrap_or(0);
            let now = SystemTime::now()
                .duration_since(SystemTime::UNIX_EPOCH)
                .expect("to be before")
                .as_secs();

            let played_seconds_ago = now - last_played;

            const HOUR: u64 = Duration::from_hours(1).as_secs();
            const DAY: u64 = Duration::from_hours(24).as_secs();
            const MONTH: u64 = Duration::from_hours(24 * 30).as_secs();

            match played_seconds_ago {
                ..HOUR => {
                    // it was played in the last hour already
                    -3
                }
                HOUR..DAY => {
                    // it was not played in the last hour, but in the last day
                    -2
                }
                DAY..MONTH => {
                    // it was not played in the last day, but in the last month
                    -1
                }
                MONTH.. => {
                    // it was not played in a month
                    1
                }
            }
        };

        let rating = i32::from(storage::rating::get(client, track).await?.unwrap_or(0));
        let play_count = i32::try_from(storage::play_count::get(client, track).await?.unwrap_or(0))
            .context("`play_count` too big")?;
        let skip_count = i32::try_from(storage::skip_count::get(client, track).await?.unwrap_or(0))
            .context("`skip_count` too big")?;

        let output: f64 = 1.0 * f64::from(rating)
            + 0.3 * f64::from(play_count)
            + -0.6 * f64::from(skip_count)
            + 0.65 * f64::from(last_played_delta);

        let weight = output.round() as i64;

        // info!("`{track}`: {weight}");

        Ok(weight)
    }
}