unitoken 0.1.3

use std::collections::{BTreeMap, HashMap, hash_map};

use crate::{MyError, MyResult, spec::WordDisplay, traits::CanStrToWord};

use super::*;

pub trait ToWord<C> {
  fn to_word(self) -> Word<C>;
}

impl<C> ToWord<C> for Vec<C> {
  fn to_word(self) -> Word<C> {
    Arc::from(self.into_boxed_slice())
  }
}

impl<C: Clone> ToWord<C> for &[C] {
  fn to_word(self) -> Word<C> {
    Arc::from(self.to_owned().into_boxed_slice())
  }
}

impl ToWord<u8> for &str {
  fn to_word(self) -> Word<u8> {
    Arc::from(self.as_bytes().to_owned().into_boxed_slice())
  }
}

impl ToWord<Character> for &str {
  fn to_word(self) -> Word<Character> {
    let chars = self.chars().map(|ch| Character::Unicode(ch)).collect::<Vec<_>>();
    Arc::from(chars.into_boxed_slice())
  }
}

impl ToWord<u8> for u8 {
  fn to_word(self) -> Word<u8> {
    Arc::from(vec![self].into_boxed_slice())
  }
}

impl ToWord<Character> for u8 {
  fn to_word(self) -> Word<Character> {
    Arc::from(vec![Character::Byte(self)].into_boxed_slice())
  }
}

impl ToWord<Character> for char {
  fn to_word(self) -> Word<Character> {
    Arc::from(vec![Character::Unicode(self)].into_boxed_slice())
  }
}

pub trait WordDebugExt {
  fn debug_display(&self) -> String;
  fn to_string_lossy(&self) -> String;
}

impl WordDebugExt for Word<u8> {
  fn debug_display(&self) -> String {
    crate::spec::uni::UniSpec.word_display(self)
  }
  fn to_string_lossy(&self) -> String {
    String::from_utf8_lossy(self).to_string()
  }
}

impl WordDebugExt for Word<Character> {
  fn debug_display(&self) -> String {
    self
      .iter()
      .map(|c| match c {
        Character::Unicode(ch) => ch.to_string(),
        Character::Byte(b) => format!("\\x{:02x}", *b),
      })
      .collect()
  }
  fn to_string_lossy(&self) -> String {
    let mut buffer = Vec::new();
    self.iter().for_each(|c| match c {
      Character::Unicode(ch) => buffer.extend_from_slice(ch.to_string().as_bytes()),
      Character::Byte(b) => buffer.extend_from_slice(&[*b]),
    });
    String::from_utf8_lossy(&buffer).to_string()
  }
}

pub(crate) fn _merge<C, I>(words: &mut Vec<PreToken<C, I>>, merge: &Merge<C, I>, target_idx: I) -> BTreeMap<(I, I), MergeData>
where
  I: Ord + Copy,
{
  // all tp with target_idx MUST be positive, so that occurs_in should be added.
  // while tp without target_idx MUST be negative, and occurs_in should be removed.
  let mut changes = BTreeMap::<(I, I), MergeData>::new();
  for k in merge.data.occurs_in.iter().copied() {
    let w = &mut words[k as usize];
    // local freq tracks the frequency changes within this word.
    let mut local_freq = BTreeMap::<(I, I), Freq>::new();
    let w_idx = &w.idxs;
    let w_freq = w.freq;
    let mut new_idxs = Vec::with_capacity(w_idx.len());
    let mut i = 0;
    let mut last_tp: Option<(I, I)> = None;
    while i + 1 < w_idx.len() {
      let tp = (w_idx[i], w_idx[i + 1]);
      *local_freq.entry(tp).or_default() += 1;
      if tp == merge.tp {
        new_idxs.push(target_idx);
        i += 2;
        changes.entry(tp).or_default().freq -= w_freq;
        *local_freq.entry(tp).or_default() -= 1;
        // deal with left neighbor,
        // e.g. in "abcd", when merging "b" and "c",
        // old_tp = ("a", "b"), new_tp = ("a", "bc")
        if let Some(old_tp) = last_tp {
          let new_tp = (old_tp.0, target_idx);
          changes.entry(old_tp).or_default().freq -= w_freq;
          changes.entry(new_tp).or_default().freq += w_freq;
          *local_freq.entry(old_tp).or_default() -= 1;
          *local_freq.entry(new_tp).or_default() -= 1;
          // if i >= w_idx.len(), loop is end, and last_tp never reads
          // last_tp = Some(new_tp);
        }
        // deal with right neighbor, notice i+=2 above
        // e.g. in "abcd", when merging "b" and "c",
        // old_tp = ("c", "d"), new_tp = ("bc", "d")
        if i < w_idx.len() {
          let old_tp = (tp.1, w_idx[i]);
          let new_tp = (target_idx, old_tp.1);
          changes.entry(old_tp).or_default().freq -= w_freq;
          changes.entry(new_tp).or_default().freq += w_freq;
          // old_tp is not increased, so that it should not be decreased
          *local_freq.entry(old_tp).or_default() -= 0;
          // when combining "b" and "c" in "bcbc",
          // new_tp=("bc", "b") would be false positive occurs_in
          *local_freq.entry(new_tp).or_default() -= 1;
          last_tp = Some(new_tp);
        }
      } else {
        new_idxs.push(w_idx[i]);
        last_tp = Some(tp);
        i += 1;
      }
    }
    if i < w_idx.len() {
      new_idxs.push(w_idx[i]);
    }

    local_freq.iter().filter(|(_, i)| **i <= 0).for_each(|(tp, _)| {
      changes.entry(*tp).and_modify(|d| { d.occurs_in.insert(k as _); });
    });

    w.idxs = new_idxs;
  }
  changes
}

pub(crate) fn _vocab_get<C, I>(vocab: &BTreeMap<I, Word<C>>, idx: I) -> MyResult<Word<C>>
where
  C: CanStrToWord,
  I: IdxLike + HasChar<C>,
{
  vocab.get(&idx).cloned().or_else(|| idx.idx_to_word()).ok_or_else(|| MyError::OovIdx(idx.to_u64()))
}

pub(crate) fn _update_merge_map<C, I>(merge_map: &mut HashMap<(I, I), Merge<C, I>>, merge: &Merge<C, I>, changes: BTreeMap<(I, I), MergeData>, vocab: Option<&BTreeMap<I, Word<C>>>)
where
  I: IdxLike + HasChar<C>,
  C: CanStrToWord,
  Word<C>: WordDebugExt,
{
  for (tp, data) in changes {
    if tp == merge.tp {
      continue;
    }
    if data.freq == 0 {
      continue;
    }
    let entry = merge_map.entry(tp);
    let entry = match entry {
      hash_map::Entry::Occupied(e) => e.into_mut(),
      hash_map::Entry::Vacant(e) => {
        if let Some(vocab) = vocab {
          let content = (
            _vocab_get(vocab, tp.0).unwrap(),
            _vocab_get(vocab, tp.1).unwrap(),
          );
          e.insert(Merge::new(tp, content))
        } else {
          continue;
        }
      }
    };
    // println!("  Change {:?} {:?} {:?}: freq {} -> {}", tp, entry.content.0.display(), entry.content.1.display(), entry.data.freq, entry.data.freq + data.freq);
    entry.data.freq += data.freq;
    if data.freq > 0 {
      entry.data.occurs_in.extend(data.occurs_in);
    } else {
      data.occurs_in.iter().for_each(|doc_id| {
        entry.data.occurs_in.remove(doc_id);
      });
    }
  }
}