MarkovChain.scala

import shapeless._
import shapeless.nat._
import shapeless.ops.nat.GT._
import shapeless.ops.nat.GTEq._
import shapeless.ops.nat._
import shapeless.syntax.sized._

import scala.collection.mutable

/**
  * Concurrent mutable Markov's Chain
  */
final class MarkovChain[Order <: Nat, T] private (implicit n: ToInt[Order])  {

  import MarkovChain.SizedVectorOps

  type Probs = Map[T, Int]
  type State = Sized[Vector[T], Order]
  
  def apply(state: State): Probs = this.synchronized {
    dict(state)
  }

  def add[S <: Nat : ToInt](sentence: Sized[Vector[T], S])
      (implicit ev: S > Order): this.type = this.synchronized {
    sentence.sliding[Succ[Order]] foreach { xs =>
      val value = xs.last
      val state = xs.take[Order]
      val probs = dict.getOrElse(state, Map.empty)
      val newProb = probs.getOrElse(value, 0) + 1
      dict.put(state, probs.updated(value, newProb))
    }
    this
  }

  private val dict = mutable.Map
    .empty[State, Probs]
    .withDefaultValue(Map.empty)
}

object MarkovChain {
  def firstOrder[T] = new MarkovChain[_1, T]
  def secondOrder[T] = new MarkovChain[_2, T]
  def thirdOrder[T] = new MarkovChain[_3, T]
  def nOrder[N <: Nat: ToInt, T] = new MarkovChain[N, T]

  implicit final class SizedVectorOps[S <: Nat : ToInt, T](xs: Sized[Vector[T], S]) {
    def sliding[M <: Nat : ToInt](implicit ev: S >= M): Iterator[Sized[Vector[T], M]] = {
      xs.unsized.sliding(toInt[M]).map(_.ensureSized[M])
    }
  }
}