rml-core 0.1.0

use rand::Rng;
use std::collections::HashMap;
use std::fs::File;
use std::io::{Read, Write};

pub const ALLOWED_CHARS: &str =
    "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789 .,!?:;-";
const VOCAB_SIZE: usize = ALLOWED_CHARS.len(); // e.g. 70

lazy_static::lazy_static! {
    static ref CHAR_TO_INDEX: HashMap<char, usize> = {
        let mut map = HashMap::new();
        for (i, c) in ALLOWED_CHARS.chars().enumerate() {
            map.insert(c, i);
        }
        map
    };

    static ref INDEX_TO_CHAR: Vec<char> = {
        ALLOWED_CHARS.chars().collect()
    };
}

const HIDDEN: usize = 128; // Significantly larger hidden layer for more complex patterns
const LR: f32 = 0.005; // Lower learning rate for more stable convergence
pub const CONTEXT_SIZE: usize = 4; // Larger context for better text coherence

#[derive(Clone)]
pub struct NGramModel {
    // Weights for the input -> hidden layer
    w1: [[f32; HIDDEN]; VOCAB_SIZE * CONTEXT_SIZE],
    // Weights for the hidden -> output layer
    w2: [[f32; VOCAB_SIZE]; HIDDEN],
    // Mapping from n-grams to indices
    context_to_index: HashMap<Vec<usize>, usize>,
    next_context_index: usize,
}

impl NGramModel {
    pub fn new() -> Self {
        let mut rng = rand::thread_rng();

        // Initialization of weight matrices
        let mut w1 = [[0.0; HIDDEN]; VOCAB_SIZE * CONTEXT_SIZE];
        let mut w2 = [[0.0; VOCAB_SIZE]; HIDDEN];

        // Random initialization of weights
        for i in 0..VOCAB_SIZE * CONTEXT_SIZE {
            for j in 0..HIDDEN {
                w1[i][j] = rng.gen_range(-0.1..0.1); // Smaller initialization for more stable training
            }
        }

        for i in 0..HIDDEN {
            for j in 0..VOCAB_SIZE {
                w2[i][j] = rng.gen_range(-0.1..0.1);
            }
        }

        Self {
            w1,
            w2,
            context_to_index: HashMap::new(),
            next_context_index: 0,
        }
    }

    // Helper function to determine or create an index for a context
    fn get_or_create_context_index(&mut self, context: &[usize]) -> usize {
        let context_vec = context.to_vec();

        if let Some(&idx) = self.context_to_index.get(&context_vec) {
            return idx;
        }

        // If we have reached the maximum number of contexts, we use a fallback
        if self.next_context_index >= VOCAB_SIZE * CONTEXT_SIZE {
            // Simple hash-based fallback
            let mut hash = 0;
            for &c in context {
                hash = (hash * 31 + c) % (VOCAB_SIZE * CONTEXT_SIZE);
            }
            return hash;
        }

        let idx = self.next_context_index;
        self.context_to_index.insert(context_vec, idx);
        self.next_context_index += 1;
        idx
    }

    pub fn forward(&mut self, context: &[usize]) -> Vec<f32> {
        if context.len() != CONTEXT_SIZE {
            panic!("Context must contain exactly {} characters", CONTEXT_SIZE);
        }

        let input_idx = self.get_or_create_context_index(context);

        // Forward pass to hidden layer
        let mut hidden = [0.0; HIDDEN];
        for j in 0..HIDDEN {
            hidden[j] = self.w1[input_idx][j].tanh();
        }

        // Forward pass to output layer
        let mut output = vec![0.0; VOCAB_SIZE];
        for i in 0..VOCAB_SIZE {
            for j in 0..HIDDEN {
                output[i] += hidden[j] * self.w2[j][i];
            }
        }

        // Softmax normalization for probabilities
        softmax(&mut output);
        output
    }

    pub fn train(&mut self, context: &[usize], target: usize) {
        if context.len() != CONTEXT_SIZE {
            panic!("Context must contain exactly {} characters", CONTEXT_SIZE);
        }

        // Ensure the target is within the vocabulary
        let target_idx = target;

        let input_idx = self.get_or_create_context_index(context);

        // Forward pass, as in forward()
        let mut hidden = [0.0; HIDDEN];
        for j in 0..HIDDEN {
            hidden[j] = self.w1[input_idx][j].tanh();
        }

        let mut logits = [0.0; VOCAB_SIZE];
        for i in 0..VOCAB_SIZE {
            for j in 0..HIDDEN {
                logits[i] += hidden[j] * self.w2[j][i];
            }
        }

        // Calculation of error with softmax cross-entropy
        let mut probs = logits;
        softmax(&mut probs);
        probs[target_idx] -= 1.0; // Simple form of cross-entropy derivative

        // Backpropagation for W2 (Hidden -> Output)
        for j in 0..HIDDEN {
            for i in 0..VOCAB_SIZE {
                self.w2[j][i] -= LR * probs[i] * hidden[j];
            }
        }

        // Backpropagation for W1 (Input -> Hidden)
        for j in 0..HIDDEN {
            // Derivative of tanh activation: 1 - tanh²(x)
            let grad = (1.0 - hidden[j] * hidden[j]) * self.w2[j][target_idx];
            self.w1[input_idx][j] -= LR * grad;
        }
    }

    // Save the model to a file
    pub fn save(&self, filename: &str) -> std::io::Result<()> {
        let mut file = File::create(filename)?;

        // Save the dimensions
        file.write_all(&(VOCAB_SIZE as u32).to_le_bytes())?;
        file.write_all(&(HIDDEN as u32).to_le_bytes())?;
        file.write_all(&(CONTEXT_SIZE as u32).to_le_bytes())?;

        // Save W1
        for row in &self.w1 {
            for &value in row {
                file.write_all(&value.to_le_bytes())?;
            }
        }

        // Save W2
        for row in &self.w2 {
            for &value in row {
                file.write_all(&value.to_le_bytes())?;
            }
        }

        // Save the number of contexts
        file.write_all(&(self.context_to_index.len() as u32).to_le_bytes())?;

        // Save the context map
        for (context, &index) in &self.context_to_index {
            // Save the context length
            file.write_all(&(context.len() as u32).to_le_bytes())?;

            // Save the context elements
            for &c in context {
                file.write_all(&(c as u32).to_le_bytes())?;
            }

            // Save the index
            file.write_all(&(index as u32).to_le_bytes())?;
        }

        // Save the next context index
        file.write_all(&(self.next_context_index as u32).to_le_bytes())?;

        Ok(())
    }

    // Load the model from a file
    pub fn load(filename: &str) -> std::io::Result<Self> {
        let mut file = File::open(filename)?;

        // Read the dimensions
        let mut buffer = [0; 4];

        file.read_exact(&mut buffer)?;
        let vocab_size = u32::from_le_bytes(buffer) as usize;

        file.read_exact(&mut buffer)?;
        let hidden_size = u32::from_le_bytes(buffer) as usize;

        file.read_exact(&mut buffer)?;
        let context_size = u32::from_le_bytes(buffer) as usize;

        // Check if the dimensions match
        if vocab_size != VOCAB_SIZE || hidden_size != HIDDEN || context_size != CONTEXT_SIZE {
            return Err(std::io::Error::new(
                std::io::ErrorKind::InvalidData,
                format!(
                    "Dimensions do not match: expected {}x{}x{}, found {}x{}x{}",
                    VOCAB_SIZE, HIDDEN, CONTEXT_SIZE, vocab_size, hidden_size, context_size
                ),
            ));
        }

        // Create a new model
        let mut model = Self::new();

        // Read W1
        for i in 0..VOCAB_SIZE * CONTEXT_SIZE {
            for j in 0..HIDDEN {
                file.read_exact(&mut buffer)?;
                model.w1[i][j] = f32::from_le_bytes(buffer);
            }
        }

        // Read W2
        for i in 0..HIDDEN {
            for j in 0..VOCAB_SIZE {
                file.read_exact(&mut buffer)?;
                model.w2[i][j] = f32::from_le_bytes(buffer);
            }
        }

        // Read the number of contexts
        file.read_exact(&mut buffer)?;
        let num_contexts = u32::from_le_bytes(buffer) as usize;

        // Read the context map
        for _ in 0..num_contexts {
            // Read the context length
            file.read_exact(&mut buffer)?;
            let context_len = u32::from_le_bytes(buffer) as usize;

            // Read the context elements
            let mut context = Vec::with_capacity(context_len);
            for _ in 0..context_len {
                file.read_exact(&mut buffer)?;
                context.push(u32::from_le_bytes(buffer) as usize);
            }

            // Read the index
            file.read_exact(&mut buffer)?;
            let index = u32::from_le_bytes(buffer) as usize;

            // Add the context to the map
            model.context_to_index.insert(context, index);
        }

        // Read the next context index
        file.read_exact(&mut buffer)?;
        model.next_context_index = u32::from_le_bytes(buffer) as usize;

        Ok(model)
    }
}

// Helper function: Softmax for normalizing output probabilities
pub fn softmax(x: &mut [f32]) {
    let max = x.iter().copied().fold(f32::NEG_INFINITY, f32::max);
    let sum: f32 = x
        .iter_mut()
        .map(|v| {
            *v = (*v - max).exp();
            *v
        })
        .sum();
    for v in x.iter_mut() {
        *v /= sum;
    }
}

// Helper function: Sampling from probabilities with temperature
pub fn sample(probs: &[f32]) -> usize {
    // Lower temperature = more conservative decisions (more typical characters)
    // Higher temperature = more randomness and creativity
    const TEMPERATURE: f32 = 0.3; // Even lower temperature for more conservative selection

    // Copy the probabilities and apply temperature
    let mut adjusted_probs = Vec::with_capacity(probs.len());
    for &p in probs {
        adjusted_probs.push(p.powf(1.0 / TEMPERATURE));
    }

    // Renormalize
    let sum: f32 = adjusted_probs.iter().sum();
    for p in &mut adjusted_probs {
        *p /= sum;
    }

    // Perform sampling
    let mut cumulative_sum = 0.0;
    let r: f32 = rand::random();
    for (i, &p) in adjusted_probs.iter().enumerate() {
        cumulative_sum += p;
        if r < cumulative_sum {
            return i;
        }
    }

    VOCAB_SIZE - 1
}

// Helper function to prepare a text for training
pub fn prepare_training_data(text: &str) -> Vec<(Vec<usize>, usize)> {
    // Only use letters, numbers, spaces and some punctuation marks
    let allowed_chars: Vec<char> = ALLOWED_CHARS.chars().collect();

    // Filter the characters
    let filtered_text: String = text.chars().filter(|c| allowed_chars.contains(c)).collect();

    // Convert to indices
    let chars: Vec<usize> = filtered_text.chars().map(|c| char_to_index(c)).collect();

    let mut data = Vec::new();

    // Fill our training data with contexts and target characters
    for i in CONTEXT_SIZE..chars.len() {
        let context = chars[i - CONTEXT_SIZE..i].to_vec();
        let target = chars[i];
        data.push((context, target));
    }

    data
}

// Helper function for converting text to indices and back
pub fn char_to_index(c: char) -> usize {
    *CHAR_TO_INDEX.get(&c).unwrap_or(&0)
}

pub fn index_to_char(idx: usize) -> char {
    INDEX_TO_CHAR[idx % VOCAB_SIZE]
}