import scala collection immutable HashMap class MinimizedTrie leaves H

import scala.collection.immutable.HashMap

class MinimizedTrie[T](leaves: HashMap[String, MinimizedTrie[T]], endPoint: Boolean, childrenSize: Int, data: Option[T]) {
  def this(endPoint: Boolean, data: Option[T]) = this(HashMap[String, MinimizedTrie[T]](), endPoint, 0, data)

  private def leavesNumber() = leaves.size

  private def getSingleLeaf() = leaves.last

  private def countWords(leaves: HashMap[String, MinimizedTrie[T]]) = leaves.foldLeft(0)(_ + _._2.size())

  private def findMaxPrefix(word: String): String = {
    for (edge <- leaves.keys) {
      val prefix = (word, edge).zipped.takeWhile(Function.tupled(_ == _)).unzip._1.mkString
      if (!prefix.isEmpty) {
        return prefix
      }
    }
    ""
  }

  def size(): Int = {
    if (endPoint) {
      return childrenSize + 1
    } else {
      return childrenSize
    }
  }

  def add(word: String, value: T): MinimizedTrie[T] = {
    if (word.isEmpty) {
      return new MinimizedTrie[T](leaves, true, childrenSize, Option[T](value))
    }

    val maxPrefix = findMaxPrefix(word)
    if (maxPrefix.isEmpty) {
      val newLeaf = word -> new MinimizedTrie[T](true, Option[T](value))
      return new MinimizedTrie[T](leaves + newLeaf, endPoint, childrenSize + 1, data)
    }

    val keyVal = leaves.find(elem => elem._1.startsWith(maxPrefix)).get
    val edge = keyVal._1
    val trie = keyVal._2

    if (edge == maxPrefix) {
      val newLeaf = edge -> trie.add(word.substring(edge.length), value)
      val newLeaves = leaves + newLeaf
      val size = countWords(newLeaves)
      return new MinimizedTrie[T](newLeaves, endPoint, size, data)
    } else {
      val end = edge.substring(maxPrefix.length)
      val middleVertex = new MinimizedTrie[T](HashMap(end -> trie), false, trie.size(), None)
      val newLeaf = maxPrefix -> middleVertex.add(word.substring(maxPrefix.length), value)

      val newLeaves = (leaves - edge) + newLeaf
      val size = countWords(newLeaves)
      return new MinimizedTrie[T](newLeaves, endPoint, size, data)
    }
  }

  def remove(word: String): MinimizedTrie[T] = {
    //this.print()
    //println()
    if (word.isEmpty) {
      return new MinimizedTrie(leaves, false, childrenSize, None)
    }

    val maxPrefix = findMaxPrefix(word)
    if (maxPrefix.isEmpty) {
      return this
    }

    val keyVal = leaves.find(elem => elem._1.startsWith(maxPrefix)).get
    val edge = keyVal._1
    val trie = keyVal._2

    if (edge == maxPrefix) {
      val newLeaf = edge -> trie.remove(word.substring(edge.length))

      if (newLeaf._2.size() == 0) {
        val newLeaves = leaves - edge
        val size = countWords(newLeaves)
        return new MinimizedTrie[T](newLeaves, endPoint, size, data)
      } else {
        newLeaf._2.leavesNumber() match {
          case 1 =>
            val nextLeaf = newLeaf._2.getSingleLeaf()
            if (newLeaf._2.size() == nextLeaf._2.size()) {
              val leaf = (edge + nextLeaf._1) -> nextLeaf._2
              val newLeaves = (leaves - edge) + leaf
              val size = countWords(newLeaves)
              return new MinimizedTrie[T](newLeaves, endPoint, size, data)
            } else {
              val newLeaves = leaves + newLeaf
              val size = countWords(newLeaves)
              return new MinimizedTrie[T](newLeaves, endPoint, size, data)
            }
          case _ =>
            val newLeaves = leaves + newLeaf
            val size = countWords(newLeaves)
            return new MinimizedTrie[T](newLeaves, endPoint, size, data)
        }
      }
    } else {
      return this
    }
  }

  def get(word: String): Option[T] = {
    if (word.isEmpty) {
      if (endPoint) {
        return data
      } else {
        return None
      }
    }

    val maxPrefix = findMaxPrefix(word)
    if (maxPrefix.isEmpty) {
      return None
    }

    val keyVal = leaves.find(elem => elem._1.startsWith(maxPrefix)).get
    if (keyVal._1 != maxPrefix) {
      return None
    }
    keyVal._2.get(word.substring(maxPrefix.length))
  }  

  def contains(word: String): Boolean = {
    return get(word) != None
  }

  def print(tabs: Int = 0) {
    if (endPoint) println(" " * tabs + "<" + data.get + ">")
    for (elem <- leaves) {
      println(" " * tabs + elem._1)
      elem._2.print(tabs + elem._1.length)
    }
  }

}

var t = new MinimizedTrie[Int](false, None)

t = t.add("aba", 0)
t = t.add("abab", 1)
t = t.add("a", 2)
t.print()
println()
t = t.remove("abab")
t.print()
println("aba: " + t.get("aba"))