about summary refs log tree commit diff stats
path: root/pkgs/by-name/mp/mpdpopm/src/dj/algorithms.rs
blob: 37c914707eb3a1db4787d9f648e4bfa9709772fd (plain) (blame)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
use std::collections::HashSet;

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(crate) struct Discovery {
    already_done: HashSet<String>,
}

impl Algorithm for Discovery {
    async fn next_track(&mut self, client: &mut Client) -> Result<String> {
        macro_rules! take {
            ($first:expr, $second:expr, $third:expr) => {{
                $first.pop().map_or_else(
                    || {
                        $second.pop().map_or_else(
                            || {
                                $third.pop().map_or_else(
                                    || {
                                        unreachable!(
                                            "This means that there are no songs in the libary"
                                        )
                                    },
                                    |val| {
                                        tracing::info!(
                                            "Selecting a `{}` track for the next entry in the queue",
                                            stringify!($third)
                                        );
                                        Ok::<_, anyhow::Error>(val)
                                    },
                                )
                            },
                            |val| {
                                tracing::info!(
                                    "Selecting a `{}` track for the next entry in the queue",
                                    stringify!($second)
                                );
                                Ok::<_, anyhow::Error>(val)
                            },
                        )
                    },
                    |val| {
                        tracing::info!(
                            "Selecting a `{}` track for the next entry in the queue",
                            stringify!($first)
                        );
                        Ok::<_, anyhow::Error>(val)
                    },
                )
            }};
        }

        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 positive = vec![];
            let mut neutral = vec![];
            let mut negative = vec![];

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

                match weight {
                    1..=i64::MAX => positive.push(track),
                    0 => neutral.push(track),
                    i64::MIN..0 => negative.push(track),
                }
            }

            (positive, neutral, negative)
        };

        let pick = {
            let mut rng = rand::rng();
            rng.sample(
                distr::weighted::WeightedIndex::new([0.65, 0.5, 0.2].iter())
                    .expect("to be valid, as hardcoded"),
            )
        };

        let next = match pick {
            0 => take!(positive, neutral, negative),
            1 => take!(neutral, positive, negative),
            2 => take!(negative, neutral, positive),
            _ => unreachable!("These indexes are not possible"),
        }?;

        self.already_done.insert(next.clone());

        Ok(next)
    }
}

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

    /// 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.
    async fn weight_track(client: &mut Client, track: &str) -> Result<i64> {
        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);

        let weight = output.round() as i64;

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

        Ok(weight)
    }
}