starfuzz

//// A module providing tokenization capabilities.
//// Tokenization splits raw text into smaller components such as characters,
//// words, or n-grams before similarity metrics or vector transformations.

import gleam/string
import gleam/list
import starfuzz/internal/ngram

/// Represents tokenization errors.
pub type TokenizeError {
  /// Returned when the n-gram size `n` is invalid (e.g. less than or equal to 0).
  InvalidN(Int)
}

/// Splits the string into its constituent Unicode grapheme characters.
pub fn characters(text: String) -> List(String) {
  string.to_graphemes(text)
}

/// Splits the string into words, using whitespace characters as boundaries.
pub fn words(text: String) -> List(String) {
  let chars = string.to_graphemes(text)
  let #(acc, current_word) = list.fold(chars, #([], []), fn(state, g) {
    let #(word_list, current) = state
    case is_whitespace(g) {
      True -> {
        case current {
          [] -> #(word_list, [])
          _ -> #([string.join(list.reverse(current), ""), ..word_list], [])
        }
      }
      False -> #(word_list, [g, ..current])
    }
  })
  let final_list = case current_word {
    [] -> acc
    _ -> [string.join(list.reverse(current_word), ""), ..acc]
  }
  list.reverse(final_list)
}

fn is_whitespace(char: String) -> Bool {
  case char {
    " " | "\t" | "\n" | "\r" -> True
    _ -> False
  }
}

/// Generates character n-grams of size `n` from the input text.
pub fn character_ngrams(text: String, n: Int) -> Result(List(String), TokenizeError) {
  let chars = characters(text)
  case ngram.ngrams(chars, n) {
    Ok(windows) -> Ok(list.map(windows, fn(w) { string.join(w, "") }))
    Error(Nil) -> Error(InvalidN(n))
  }
}

/// Generates word n-grams of size `n` from the input text.
pub fn word_ngrams(text: String, n: Int) -> Result(List(List(String)), TokenizeError) {
  let wrds = words(text)
  case ngram.ngrams(wrds, n) {
    Ok(windows) -> Ok(windows)
    Error(Nil) -> Error(InvalidN(n))
  }
}