midi-keys/src/typing.rs

use itertools::Itertools;
use std::collections::HashMap;

#[derive(Debug, Clone)]
pub enum Decoded {
    Keystroke(Keystroke),
    Partial(Vec<u8>, Vec<(Vec<u8>, Keystroke)>),
    Invalid(Vec<u8>),
}

#[derive(Debug, Clone)]
pub struct Mapping {
    sequences: Vec<(Vec<u8>, Keystroke)>,
    key_to_seq: HashMap<Keystroke, Vec<u8>>,
}

impl Mapping {
    pub fn new(mut sequences: Vec<(Vec<u8>, Keystroke)>) -> Mapping {
        sequences.sort();
        let key_to_seq = sequences.iter().cloned().map(|(a, b)| (b, a)).collect();

        Mapping {
            sequences,
            key_to_seq,
        }
    }

    pub fn encode(&self, key: Keystroke) -> Option<Vec<u8>> {
        self.key_to_seq.get(&key).cloned()
    }

    pub fn decode(&self, notes: &Vec<u8>) -> Decoded {
        if let Some(key) = self.decode_exact(notes) {
            return Decoded::Keystroke(key.clone());
        } else if let Some(keys) = self.prefix_range(notes) {
            return Decoded::Partial(notes.clone(), keys);
        } else {
            return Decoded::Invalid(notes.clone());
        }
    }

    pub fn decode_exact(&self, notes: &Vec<u8>) -> Option<Keystroke> {
        let index = self
            .sequences
            .binary_search_by(|(n, _key)| n.cmp(&notes))
            .ok()?;
        let (_notes, key) = self.sequences.get(index)?;
        Some(key.clone())
    }

    pub fn prefix_range(&self, notes: &Vec<u8>) -> Option<Vec<(Vec<u8>, Keystroke)>> {
        // we only want to match prefixes with exactly this length of notes
        let valid_seqs: Vec<_> = self
            .sequences
            .iter()
            .filter(|(n, _key)| n.starts_with(&notes))
            .cloned()
            .collect();

        if valid_seqs.is_empty() {
            None
        } else {
            Some(valid_seqs)
        }
    }

    pub fn collated_remaining(&self, notes: &Vec<u8>) -> Option<Vec<(u8, Vec<Keystroke>)>> {
        let next_sequences = self.prefix_range(notes)?;

        let seqs = self
            .sequences
            .iter()
            .filter(|(n, _key)| n.starts_with(&notes))
            .map(|(n, key)| (*n.get(notes.len()).unwrap(), key.clone()));

        let mut grouped_seqs = Vec::new();

        for (note, keys) in &seqs.chunk_by(|(n, _key)| *n) {
            let all_keys: Vec<_> = keys.map(|(n,k)| k.clone()).collect();
            grouped_seqs.push((note, all_keys));
        }

        //    grouped.map(|(n, keys)| (*n, keys.collect()));
        //     .collect();

        Some(grouped_seqs)
    }

    pub fn print_sequences(&self) {
        for (seq, key) in self.sequences.iter() {
            for deg in seq {
                print!("{} ", c_scale_note_name(*deg));
            }
            println!("=> {key:?}");
        }
    }
}

pub fn c_scale_note_name(degree: u8) -> String {
    match degree {
        0 => "C",
        1 => "D",
        2 => "E",
        3 => "F",
        4 => "G",
        5 => "A",
        6 => "B",
        _ => "?",
    }
    .to_string()
}

pub fn standard_melodic_mapping() -> Mapping {
    let modspecial = [
        Keystroke::Modifier(Modifier::Shift),
        Keystroke::Special(Special::CapsLock),
        Keystroke::Special(Special::Backspace),
    ];

    let alpha = 'a'..='z';
    let num = '0'..='9';
    let punct = ".,! \n".chars();
    let chars = alpha.chain(num).chain(punct).map(|c| Keystroke::Char(c));

    let keys = modspecial.iter().cloned().chain(chars);

    let sequences: Vec<_> = keys
        .enumerate()
        .map(|(k, key)| (base_values(k as u32, 7, 2), key))
        .collect();

    Mapping::new(sequences)
}

#[derive(Debug, Clone)]
pub struct TypingState {
    mapping: Mapping,
    current: Vec<u8>,
}

impl TypingState {
    pub fn new(mapping: Mapping) -> TypingState {
        TypingState {
            mapping,
            current: vec![],
        }
    }

    /// Add the value into the typing state.
    pub fn enter(&mut self, value: u8) {
        self.current.push(value);
    }

    /// Attempts to decode the current value stored. It will reset if it's either
    /// invalid or valid, but will retain the current buffer if it's partial.
    pub fn emit(&mut self) -> Decoded {
        let decoded = self.mapping.decode(&self.current);

        if !matches!(decoded, Decoded::Partial(..)) {
            self.current.clear();
        }

        decoded
    }

    pub fn current_notes(&self) -> Vec<u8> {
        self.current.clone()
    }

    pub fn mapping(&self) -> &Mapping {
        &self.mapping
    }
}

/// Chunks a u8 into a sequence of values in a given base.
pub fn base_values(mut x: u32, base: u32, len: usize) -> Vec<u8> {
    let mut vals = vec![];
    for _ in 0..len {
        let digit = x % (base as u32);
        vals.push(digit as u8);
        x = x / base;
    }
    vals.reverse();

    vals
}

/// Determines the scale degree of a given MIDI note relative to the provided
/// scale root.
pub fn midi_to_scale_degree(root: u8, note: u8) -> Option<u8> {
    let semitones = (note as i32 - root as i32).rem_euclid(12);
    Some(match semitones {
        0 => 0,
        2 => 1,
        4 => 2,
        5 => 3,
        7 => 4,
        9 => 5,
        11 => 6,
        _ => return None,
    })
}

#[derive(Debug, PartialEq, Eq, Hash, Clone, PartialOrd, Ord)]
pub enum Keystroke {
    Char(char),
    Modifier(Modifier),
    Special(Special),
}

#[derive(Debug, PartialEq, Eq, Hash, Clone, PartialOrd, Ord)]
pub enum Modifier {
    Shift,
}

#[derive(Debug, PartialEq, Eq, Hash, Clone, PartialOrd, Ord)]
pub enum Special {
    CapsLock,
    Backspace,
}

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

    #[test]
    fn encodes_lowercase() {
        let mapping = standard_melodic_mapping();
        assert_eq!(mapping.encode(Keystroke::Char('a')), Some(vec![0, 3]));
        assert_eq!(mapping.encode(Keystroke::Char('b')), Some(vec![0, 4]));
        assert_eq!(mapping.encode(Keystroke::Char('z')), Some(vec![4, 0]));
    }

    #[test]
    fn encodes_numbers() {
        let mapping = standard_melodic_mapping();
        assert_eq!(mapping.encode(Keystroke::Char('0')), Some(vec![4, 1]));
        assert_eq!(mapping.encode(Keystroke::Char('1')), Some(vec![4, 2]));
        assert_eq!(mapping.encode(Keystroke::Char('9')), Some(vec![5, 3]));
    }

    #[test]
    fn converts_to_scale_degree() {
        assert_eq!(midi_to_scale_degree(60, 58), None);
        assert_eq!(midi_to_scale_degree(60, 59), Some(6));
        assert_eq!(midi_to_scale_degree(60, 60), Some(0));
        assert_eq!(midi_to_scale_degree(60, 61), None);
        assert_eq!(midi_to_scale_degree(60, 62), Some(1));
        assert_eq!(midi_to_scale_degree(60, 63), None);
        assert_eq!(midi_to_scale_degree(60, 64), Some(2));
        assert_eq!(midi_to_scale_degree(60, 65), Some(3));
        assert_eq!(midi_to_scale_degree(60, 66), None);
        assert_eq!(midi_to_scale_degree(60, 67), Some(4));
        assert_eq!(midi_to_scale_degree(60, 68), None);
        assert_eq!(midi_to_scale_degree(60, 69), Some(5));
        assert_eq!(midi_to_scale_degree(60, 70), None);
        assert_eq!(midi_to_scale_degree(60, 71), Some(6));
        assert_eq!(midi_to_scale_degree(60, 72), Some(0));
        assert_eq!(midi_to_scale_degree(60, 73), None);
        assert_eq!(midi_to_scale_degree(60, 74), Some(1));
    }
}