simdphrase 0.1.1

use std::{cmp::Reverse, collections::HashSet, path::Path};

use crate::{
    db::{Document, DB, MAX_WINDOW_LEN},
    decreasing_window_iter::DecreasingWindows,
    error::DbError,
    roaringish::MAX_VALUE,
    utils::{normalize, tokenize},
    RoaringishPacked, Searcher,
};
use fxhash::FxHashMap;
use gxhash::{HashMap as GxHashMap, HashMapExt};
use heed::RwTxn;
use hyperloglogplus::{HyperLogLog, HyperLogLogPlus};

/// Specifies how the common tokens are treated during indexing.
#[derive(Debug)]
pub enum CommonTokens {
    /// Fixed list specified by the user.
    List(HashSet<String>),
    /// Top `n` most frequent tokens.
    FixedNum(u32),
    /// Percentage of the top `n` most frequent tokens.
    Percentage(f64),
}

/// Batch of documents to be indexed.
#[derive(Debug)]
struct Batch<D: Document> {
    /// Monotonically increasing batch id.
    batch_id: u32,

    /// Used to estimate the number of distinct tokens.
    hllp_tokens: HyperLogLogPlus<Box<str>, gxhash::GxBuildHasher>,

    /// Monotonically increasing token id (cleared after each batch).
    next_token_id: u32,

    /// Maps tokens to token ids (cleared after each batch).
    token_to_token_id: GxHashMap<Box<str>, u32>,

    /// Maps token ids to their roaringish packed data (cleared after each batch).
    ///
    /// This should be in sync with `token_id_to_token`
    token_id_to_roaringish_packed: Vec<RoaringishPacked>,
    /// Maps token ids to tokens (cleared after each batch).
    ///
    /// This should be in sync with `token_id_to_roaringish_packed`
    token_id_to_token: Vec<Box<str>>,

    // this 3 containers are in sync
    /// Document ids in the batch (cleared after each batch).
    ///
    /// This should be in sync with `documents` and `tokenized_docs`.
    doc_ids: Vec<u32>,
    /// Documents in the batch (cleared after each batch).
    ///
    /// This should be in sync with `doc_ids` and `tokenized_docs`.
    documents: Vec<D>,

    /// Tokenized representation of the documents in the batch (cleared after each batch).
    /// This representation is done by storing the token id.
    ///
    /// This should be in sync with `doc_ids` and `documents`.
    tokenized_docs: Vec<Vec<u32>>,
}

impl<D: Document> Batch<D> {
    /// Constructs a new batch.
    fn new() -> Self {
        Self {
            batch_id: 0,
            // This can't fail
            hllp_tokens: HyperLogLogPlus::new(18, gxhash::GxBuildHasher::default()).unwrap(),
            next_token_id: 0,
            token_to_token_id: GxHashMap::new(),
            token_id_to_roaringish_packed: Vec::new(),
            token_id_to_token: Vec::new(),
            doc_ids: Vec::new(),
            documents: Vec::new(),
            tokenized_docs: Vec::new(),
        }
    }

    /// Get an overestimated number of distinct tokens.
    fn estimate_number_of_distinct_tokens(&mut self) -> u64 {
        (self.hllp_tokens.count() * 1.015f64) as u64
    }

    /// Clears the batch.
    ///
    /// This should be called after each flush and before the start of a new batch.
    fn clear(&mut self) {
        self.next_token_id = 0;
        self.token_to_token_id.clear();
        self.token_id_to_roaringish_packed.clear();
        self.token_id_to_token.clear();

        self.doc_ids.clear();
        self.documents.clear();
        self.tokenized_docs.clear();
    }

    /// Adds a document to the batch and starts the indexing process.
    ///
    /// `count_freq` is used to count the frequency of each token. This should
    /// only be used in the first batch, allowing us to generate the common tokens.
    fn push(&mut self, doc_id: u32, content: &str, doc: D, count_freq: impl FnMut(&str)) {
        let tokenized_doc = self.index_doc(content, doc_id, count_freq);
        self.doc_ids.push(doc_id);
        self.documents.push(doc);
        self.tokenized_docs.push(tokenized_doc);
    }

    /// Get the token id for the input `token`. If the token is not present in the
    /// in the batch then it's added and a new token id is generated by incrementing
    /// the current value.
    fn get_token_id(
        token: &str,
        hllp_tokens: &mut HyperLogLogPlus<Box<str>, gxhash::GxBuildHasher>,
        token_to_token_id: &mut GxHashMap<Box<str>, u32>,
        token_id_to_token: &mut Vec<Box<str>>,
        token_id_to_roaringish_packed: &mut Vec<RoaringishPacked>,
        next_token_id: &mut u32,
    ) -> u32 {
        hllp_tokens.insert(token);

        let (_, token_id) = token_to_token_id
            .raw_entry_mut()
            .from_key(token)
            .or_insert_with(|| {
                let current_token_id = *next_token_id;
                *next_token_id += 1;
                (token.to_string().into_boxed_str(), current_token_id)
            });

        if *token_id as usize >= token_id_to_token.len() {
            token_id_to_token.push(token.to_string().into_boxed_str());
            token_id_to_roaringish_packed.push(RoaringishPacked::default());
        }

        *token_id
    }

    /// Indexes `content`s for this `doc_id`.
    ///
    /// `count_freq` is used to count the frequency of each token. This should
    /// only be used in the first batch, allowing us to generate the common tokens.
    fn index_doc(
        &mut self,
        content: &str,
        doc_id: u32,
        mut count_freq: impl FnMut(&str),
    ) -> Vec<u32> {
        let mut tokenized_doc = Vec::new();
        let mut token_id_to_positions: FxHashMap<u32, Vec<u32>> = FxHashMap::new();
        let content = normalize(content);
        for (pos, token) in tokenize(&content).enumerate().take(MAX_VALUE as usize) {
            let token_id = Self::get_token_id(
                token,
                &mut self.hllp_tokens,
                &mut self.token_to_token_id,
                &mut self.token_id_to_token,
                &mut self.token_id_to_roaringish_packed,
                &mut self.next_token_id,
            );

            count_freq(token);

            token_id_to_positions
                .entry(token_id)
                .or_default()
                .push(pos as u32);
            tokenized_doc.push(token_id);
        }

        for (token_id, positions) in token_id_to_positions.iter() {
            self.token_id_to_roaringish_packed[*token_id as usize].push(doc_id, positions);
        }
        tokenized_doc
    }

    /// Flushes the batch.
    fn flush(
        &mut self,
        db: &DB<D>,
        rwtxn: &mut RwTxn,
        common_tokens: &HashSet<Box<str>>,
        mmap_size: &mut usize,
    ) -> Result<(), DbError> {
        log::info!("Flushing batch");
        let b = std::time::Instant::now();

        // Nothing to do if the batch is empty.
        if self.doc_ids.is_empty() {
            log::debug!("Empty batch, nothing to flush");
            return Ok(());
        }

        self.merge_common_tokens(common_tokens);

        db.write_token_to_roaringish_packed(
            &self.token_to_token_id,
            &self.token_id_to_roaringish_packed,
            mmap_size,
            self.batch_id,
        )?;
        db.write_doc_id_to_document(rwtxn, &self.doc_ids, &self.documents)?;

        self.batch_id += 1;
        self.clear();
        log::info!("Flush took {:?}", b.elapsed());
        Ok(())
    }

    /// Merges the tokens for all of the documents in the batch.
    /// This will create new tokens and consequently new token ids.
    ///
    /// The generation is done by merging up to [MAX_WINDOW_LEN] tokens at a time.
    /// We are only allowed to merge:
    /// * Common tokens with other common tokens.
    /// * Rare tokens with a common token.
    /// * Common tokens with a rare token.
    ///
    /// So if it's impossible for the generated token to have two rare tokens.
    /// Also if rare tokens can only be in the first or last position, for example:
    ///
    /// # Examples
    /// ```
    /// c c
    /// c c c
    /// r c
    /// r c c
    /// c r
    /// c c r
    /// ```
    ///
    /// This will generate all possible combinations of the merging process.
    fn merge_common_tokens(&mut self, common_tokens: &HashSet<Box<str>>) {
        log::debug!("Merging common tokens");
        if common_tokens.is_empty() {
            return;
        }

        let b = std::time::Instant::now();
        for (tokenized_doc, doc_id) in self.tokenized_docs.iter().zip(self.doc_ids.iter()) {
            let mut token_id_to_positions: FxHashMap<u32, Vec<u32>> = FxHashMap::new();
            let it = DecreasingWindows::new(tokenized_doc, MAX_WINDOW_LEN);
            for (pos, token_ids) in it.enumerate() {
                let token_id = token_ids[0];
                let token = &self.token_id_to_token[token_id as usize];
                let is_first_token_rare = !common_tokens.contains(token);

                for i in 1..token_ids.len() {
                    let token_id = token_ids[i];
                    let token = &self.token_id_to_token[token_id as usize];
                    let is_token_rare = !common_tokens.contains(token);
                    if is_first_token_rare && is_token_rare {
                        break;
                    }
                    let token: String = token_ids[..i + 1]
                        .iter()
                        .map(|token_id| self.token_id_to_token[*token_id as usize].as_ref())
                        .intersperse(" ")
                        .collect();
                    let token_id = Self::get_token_id(
                        &token,
                        &mut self.hllp_tokens,
                        &mut self.token_to_token_id,
                        &mut self.token_id_to_token,
                        &mut self.token_id_to_roaringish_packed,
                        &mut self.next_token_id,
                    );
                    token_id_to_positions
                        .entry(token_id)
                        .or_default()
                        .push(pos as u32);
                    if is_token_rare {
                        break;
                    }
                }
            }

            for (token_id, positions) in token_id_to_positions.iter() {
                self.token_id_to_roaringish_packed[*token_id as usize].push(*doc_id, positions);
            }
        }
        log::debug!("Merge took {:?}", b.elapsed());
    }
}

/// Responsible for indexing documents.
pub struct Indexer {
    batch_size: Option<u32>,
    common_tokens: Option<CommonTokens>,
}

impl Indexer {
    /// Creates a new indexer.
    ///
    /// * If `batch_size` is [None] then the indexer will index all the documents in a single batch.
    /// * If `common_tokens` is [None] then merging will happen.
    pub fn new(batch_size: Option<u32>, common_tokens: Option<CommonTokens>) -> Self {
        Self {
            batch_size,
            common_tokens,
        }
    }

    /// Generates the list of common tokens to be used
    /// in the merging phase
    fn generate_common_tokens(
        &self,
        token_to_freq: &GxHashMap<Box<str>, u32>,
    ) -> HashSet<Box<str>> {
        let Some(common_tokens) = &self.common_tokens else {
            return HashSet::new();
        };
        match common_tokens {
            CommonTokens::List(tokens) => tokens
                .into_iter()
                .map(|t| t.to_string().clone().into_boxed_str())
                .collect(),
            CommonTokens::FixedNum(max) => {
                let max = (*max as usize).min(token_to_freq.len());
                let mut token_to_freq: Vec<_> = token_to_freq.iter().collect();
                token_to_freq.sort_unstable_by_key(|(_, freq)| Reverse(*freq));
                token_to_freq[0..max]
                    .iter()
                    .map(|(token, _)| (*token).clone())
                    .collect()
            }
            CommonTokens::Percentage(p) => {
                let max = (token_to_freq.len() as f64 * *p) as usize;
                let mut token_to_freq: Vec<_> = token_to_freq.iter().collect();
                token_to_freq.sort_unstable_by_key(|(_, freq)| Reverse(*freq));
                token_to_freq[0..max]
                    .iter()
                    .map(|(token, _)| (*token).clone())
                    .collect()
            }
        }
    }

    /// Indexes an iterator of documents.
    ///
    /// This iterator should essentially return a tuple `(&str, D)`, where
    /// `D` is the form of the document that will be serialized and stored in the database.
    ///
    /// So the content of the document (`&str`) can be different from the stored version (`D`).
    ///
    /// The type `D` is anything that can be serialized by [rkyv].
    /// 
    /// This returns a [Searcher] object and the number of indexed documents.
    pub fn index<S, D, I, P>(&self, docs: I, path: P, db_size: usize) -> Result<(Searcher<D>, u32), DbError>
    where
        S: AsRef<str>,
        I: IntoIterator<Item = (S, D)>,
        D: Document,
        P: AsRef<Path>,
    {
        let path = path.as_ref();
        let db = DB::truncate(path, db_size)?;
        let mut rwtxn = db.env.write_txn()?;

        let mut batch = Batch::new();

        let batch_size = self.batch_size.unwrap_or(u32::MAX);
        let mut it = docs.into_iter();

        let mut token_to_freq = GxHashMap::new();
        let mut next_doc_id = 0;
        let mut mmap_size = 0;

        log::info!("Starting first batch");
        // Index the first batch to generate the common tokens
        let b = std::time::Instant::now();
        for (content, doc) in it.by_ref() {
            let doc_id = next_doc_id;
            next_doc_id += 1;

            batch.push(doc_id, content.as_ref(), doc, |token| {
                let (_, freq) = token_to_freq
                    .raw_entry_mut()
                    .from_key(token)
                    .or_insert_with(|| (token.to_owned().into_boxed_str(), 0));
                *freq += 1;
            });

            if next_doc_id % batch_size == 0 {
                break;
            }
        }
        log::info!("First batch took {:?}", b.elapsed());

        let common_tokens = self.generate_common_tokens(&token_to_freq);
        drop(token_to_freq);

        batch.flush(&db, &mut rwtxn, &common_tokens, &mut mmap_size)?;

        // Index the rest of the documents
        log::info!("Starting new batch");
        let mut b = std::time::Instant::now();
        for (content, doc) in it {
            let doc_id = next_doc_id;
            next_doc_id += 1;

            batch.push(doc_id, content.as_ref(), doc, |_| {});

            if next_doc_id % batch_size == 0 {
                log::info!("Batch took {:?}", b.elapsed());
                b = std::time::Instant::now();
                batch.flush(&db, &mut rwtxn, &common_tokens, &mut mmap_size)?;
                log::info!("Starting new batch");
            }
        }

        // Flush the last batch
        batch.flush(&db, &mut rwtxn, &common_tokens, &mut mmap_size)?;

        let number_of_distinct_tokens = batch.estimate_number_of_distinct_tokens();
        log::debug!(
            "Approximation for the number of distinct tokens: {}",
            number_of_distinct_tokens
        );

        // Write to db
        db.write_common_tokens(&mut rwtxn, &common_tokens)?;
        db.generate_mmap_file(
            number_of_distinct_tokens,
            mmap_size,
            batch.batch_id,
            &mut rwtxn,
        )?;

        let b = std::time::Instant::now();
        log::info!("Commiting");
        rwtxn.commit()?;
        log::info!("Commit took {:?}", b.elapsed());

        let searcher = Searcher::new(path)?;
        Ok((searcher, next_doc_id))
    }
}