jap_learn/src/core/srs.rs

use crate::core::deck::DeckData;
use rand::distr::weighted::WeightedIndex;
use rand::prelude::ThreadRng;
use rand::prelude::{Distribution, SliceRandom};
use rand::rng;
use crate::data::entities::CardStatistics;

const MAX_HISTORY_LEN: usize = 20;
const MAX_HISTORY_LEN_PART: f32 = 0.33;
pub const MAX_SCORE: u8 = 25;
pub const FADE_PER_DAY: f32 = 0.95;

#[derive(Clone, PartialEq, Copy, Eq)]
pub enum OrderMode {
    Default,
    TrainWorstFirst,
    FullRandom,
}
#[derive(Clone, Copy, Eq, PartialEq, Default)]
pub enum AppendMode {
    Full,
    #[default]
    Manual,
}
#[derive(Clone)]
pub enum OrderModule {
    SemiRandomSRS(SemiRandomSRSModule),
    RandomSRS(RandomSRSModule),
    WorstWordsSRS(WorstWordsSRSModule),
}
pub trait SRSModule {
    fn next(&mut self, set: &mut DeckData) -> usize;
    fn open(&mut self, status: WordOpenMode, index: usize, updated_word: &CardStatistics);
    fn init(&mut self, set: &mut DeckData);
}
#[derive(Clone)]
pub struct RandomSRSModule {
    basket: Vec<usize>,
    pub(crate) initialized: bool,
}
impl RandomSRSModule {
    pub(crate) fn new() -> RandomSRSModule {
        Self {
            basket: vec![],
            initialized: false,
        }
    }
}
impl SRSModule for RandomSRSModule {
    fn next(&mut self, set: &mut DeckData) -> usize {
        match self.basket.pop() {
            Some(index) => index,
            None => {
                self.basket = (0..set.words.len()).collect::<Vec<usize>>();
                self.basket.shuffle(&mut rand::rng());
                self.basket.pop().expect("empty basket")
            }
        }
    }

    fn open(&mut self, _: WordOpenMode, _: usize, _: &CardStatistics) {}

    fn init(&mut self, set: &mut DeckData) {
        self.initialized = true;
        self.basket = (0..set.set.len()).collect::<Vec<usize>>();
        self.basket.shuffle(&mut rand::rng())
    }
}
#[derive(Clone)]
pub struct SemiRandomSRSModule {
    history: Vec<usize>,
    last_weights: WeightedIndex<f32>,
    generator: ThreadRng,
    pub(crate) initialized: bool,
}
impl SemiRandomSRSModule {
    pub(crate) fn new() -> SemiRandomSRSModule {
        SemiRandomSRSModule {
            history: vec![],
            last_weights: WeightedIndex::new([1.0]).unwrap(),
            generator: rng(),
            initialized: false,
        }
    }
}
impl SRSModule for SemiRandomSRSModule {
    fn next(&mut self, set: &mut DeckData) -> usize {
        let index = self.last_weights.sample(&mut self.generator);

        if self.history.contains(&index) && self.history.len() != 1 {
            return self.next(set);
        }

        if self.history.len() == self.history_len(set) {
            self.history.remove(0);
        }
        self.history.push(index);

        index
    }
    fn open(&mut self, _: WordOpenMode, index: usize, word: &CardStatistics) {
        let new_weight = (100.0 / word.calculated_score()).powf(2.0);
        self.last_weights
            .update_weights(&[(index, &new_weight)])
            .expect("update weights error");
    }
    fn init(&mut self, set: &mut DeckData) {
        self.initialized = true;
        let weights = set
            .set
            .iter()
            .map(|s| (100.0 / s.calculated_score()).powf(2.0) * 2.0)
            .collect::<Vec<f32>>();
        self.last_weights = WeightedIndex::new(weights).unwrap();
    }
}
impl SemiRandomSRSModule {
    fn history_len(&self, set: &DeckData) -> usize {
        ((set.len() as f32 * MAX_HISTORY_LEN_PART) as usize).clamp(1, MAX_HISTORY_LEN)
    }
}
#[derive(Clone)]
pub struct WorstWordsSRSModule {
    pub(crate) initialized: bool,
    pool_size: usize,
    rounds_remaining: u8,
    pool: Vec<usize>,
    queue: Vec<usize>,
    rounds_count: u8,
}
impl SRSModule for WorstWordsSRSModule {
    fn next(&mut self, set: &mut DeckData) -> usize {
        if self.rounds_remaining == 0 {
            self.fill_pool(set);
            self.rounds_remaining = self.rounds_count;
        }

        match self.queue.pop() {
            Some(index) => index,
            None => {
                self.queue.append(&mut self.pool.clone());
                self.rounds_remaining -= 1;
                self.queue.pop().expect("empty queue")
            }
        }
    }

    fn open(&mut self, _: WordOpenMode, _: usize, _: &CardStatistics) {}

    fn init(&mut self, _: &mut DeckData) {
        self.initialized = true;
    }
}
impl WorstWordsSRSModule {
    pub(crate) fn new() -> WorstWordsSRSModule {
        WorstWordsSRSModule {
            initialized: false,
            pool_size: 15,
            rounds_count: 2,
            rounds_remaining: 0,
            pool: vec![],
            queue: vec![],
        }
    }
    fn fill_pool(&mut self, set: &DeckData) {
        let mut sorted = set
            .set
            .clone()
            .into_iter()
            .zip(0..set.set.len())
            .collect::<Vec<_>>();
        sorted.sort_by_key(|c| c.0.score);
        let mut worst = sorted[0..self.pool_size]
            .iter()
            .map(|(_, index)| *index)
            .collect::<Vec<usize>>();
        worst.shuffle(&mut rand::rng());
        self.pool = worst;
    }
}

#[derive(Clone, Copy)]
pub enum WordOpenMode {
    Easy,
    Ok,
    Hard,
    None,
}