/* -----------------------------------------------------------------------------
Copyright 2019-2021 Kevin P. Barry

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
----------------------------------------------------------------------------- */

// Author: Kevin P. Barry [ta0kira@gmail.com]

define SearchTree {
  @value optional Node<#k,#v> root

  new () {
    return #self{ empty }
  }

  set (k,v) {
    root <- Node<#k,#v>.insert(root,k,v)
    return self
  }

  remove (k) {
    root <- Node<#k,#v>.delete(root,k)
    return self
  }

  get (k) {
    return Node<#k,#v>.find(root,k)
  }

  defaultOrder () {
    return ForwardTreeOrder:create<?,?>(root)
  }

  reverseOrder () {
    return ReverseTreeOrder:create<?,?>(root)
  }
}

define ValidatedTree {
  @value optional Node<#k,#v> root

  new () {
    return #self{ empty }
  }

  set (k,v) {
    root <- Node<#k,#v>.insert(root,k,v)
    \ Node<#k,#v>.validate(root)
    return self
  }

  remove (k) {
    root <- Node<#k,#v>.delete(root,k)
    \ Node<#k,#v>.validate(root)
    return self
  }

  get (k) {
    return Node<#k,#v>.find(root,k)
  }
}

concrete Node<#k,#v> {
  #k defines LessThan<#k>

  @type insert (optional Node<#k,#v>,#k,#v) -> (optional Node<#k,#v>)
  @type delete (optional Node<#k,#v>,#k) -> (optional Node<#k,#v>)
  @type find (optional Node<#k,#v>,#k) -> (optional #v)
  @type validate (optional Node<#k,#v>) -> ()

  @value getLower () -> (optional Node<#k,#v>)
  @value getHigher () -> (optional Node<#k,#v>)
  @value getKey () -> (#k)
  @value getValue () -> (#v)
}

define Node {
  @value Int height
  @value #k key
  @value #v value
  @value optional Node<#k,#v> lower
  @value optional Node<#k,#v> higher

  insert (node,k,v) {
    if (!present(node)) {
      return #self{ 1, k, v, empty, empty }
    }
    Node<#k,#v> node2 <- require(node)
    if (k `#k.lessThan` node2.getKey()) {
      \ node2.setLower(insert(node2.getLower(),k,v))
      return rebalance(node2)
    } elif (node2.getKey() `#k.lessThan` k) {
      \ node2.setHigher(insert(node2.getHigher(),k,v))
      return rebalance(node2)
    } else {
      \ node2.setValue(v)
      return node2
    }
  }

  delete (node,k) {
    if (!present(node)) {
      return empty
    }
    Node<#k,#v> node2 <- require(node)
    if (k `#k.lessThan` node2.getKey()) {
      \ node2.setLower(delete(node2.getLower(),k))
      return rebalance(node2)
    } elif (node2.getKey() `#k.lessThan` k) {
      \ node2.setHigher(delete(node2.getHigher(),k))
      return rebalance(node2)
    } else {
      return rebalance(removeNode(node2))
    }
  }

  find (node,k) {
    if (present(node)) {
      scoped {
        Node<#k,#v> node2 <- require(node)
      } in if (k `#k.lessThan` node2.getKey()) {
        return find(node2.getLower(),k)
      } elif (node2.getKey() `#k.lessThan` k) {
        return find(node2.getHigher(),k)
      } else {
        return node2.getValue()
      }
    } else {
      return empty
    }
  }

  validate (node) {
    if (present(node)) {
      \ require(node).validateOrder()
      \ require(node).validateBalance()
    }
  }

  getLower () { return lower }

  getHigher () { return higher }

  getKey () { return key }

  getValue () { return value }

  @value validateOrder () -> ()
  validateOrder () {
    if (present(lower)) {
      if (!(require(lower).getKey() `#k.lessThan` getKey())) {
        fail("bad lower order")
      }
      \ require(lower).validateOrder()
    }
    if (present(higher)) {
      if (!(getKey() `#k.lessThan` require(higher).getKey())) {
        fail("bad higher order")
      }
      \ require(higher).validateOrder()
    }
  }

  @value validateBalance () -> ()
  validateBalance () {
    scoped {
      Int balance <- getBalance()
    } in if (balance > 1 || balance < -1) {
      fail("out of balance: " + balance.formatted())
    }
    if (present(lower)) {
      \ require(lower).validateBalance()
    }
    if (present(higher)) {
      \ require(higher).validateBalance()
    }
  }

  @value updateHeight () -> ()
  updateHeight () {
    scoped {
      Int l <- 0
      Int h <- 0
      if (present(lower)) {
        l <- require(lower).getHeight()
      }
      if (present(higher)) {
        h <- require(higher).getHeight()
      }
    } in if (l > h) {
      height <- l + 1
    } else {
      height <- h + 1
    }
  }

  @value getBalance () -> (Int)
  getBalance () {
    scoped {
      Int l <- 0
      Int h <- 0
      if (present(lower)) {
        l <- require(lower).getHeight()
      }
      if (present(higher)) {
        h <- require(higher).getHeight()
      }
    } in return h - l
  }

  @type rebalance (optional Node<#k,#v>) -> (optional Node<#k,#v>)
  rebalance (node) {
    if (!present(node)) {
      return empty
    }
    Node<#k,#v> node2 <- require(node)
    \ node2.updateHeight()
    scoped {
      Int balance <- node2.getBalance()
    } in if (balance > 1) {
      return pivotLower(node2)
    } elif (balance < -1) {
      return pivotHigher(node2)
    } else {
      return node2
    }
  }

  @type pivotHigher (Node<#k,#v>) -> (Node<#k,#v>)
  pivotHigher (node) (newNode) {
    if (require(node.getLower()).getBalance() > 0) {
      \ node.setLower(pivotLower(require(node.getLower())))
    }
    newNode <- require(node.getLower())
    \ node.setLower(newNode.getHigher())
    \ node.updateHeight()
    \ newNode.setHigher(node)
    \ newNode.updateHeight()
  }

  @type pivotLower (Node<#k,#v>) -> (Node<#k,#v>)
  pivotLower (node) (newNode) {
    if (require(node.getHigher()).getBalance() < 0) {
      \ node.setHigher(pivotHigher(require(node.getHigher())))
    }
    newNode <- require(node.getHigher())
    \ node.setHigher(newNode.getLower())
    \ node.updateHeight()
    \ newNode.setLower(node)
    \ newNode.updateHeight()
  }

  @type removeNode (Node<#k,#v>) -> (optional Node<#k,#v>)
  removeNode (node) (newNode) {
    if (node.getBalance() < 0) {
      optional Node<#k,#v> temp, newNode <- removeHighest(node.getLower())
      \ node.setLower(temp)
    } else {
      optional Node<#k,#v> temp, newNode <- removeLowest(node.getHigher())
      \ node.setHigher(temp)
    }
    if (present(newNode)) {
      \ swapChildren(node,require(newNode))
      \ require(newNode).updateHeight()
    }
  }

  @type removeHighest (optional Node<#k,#v>) -> (optional Node<#k,#v>,optional Node<#k,#v>)
  removeHighest (node) (newNode,removed) {
    if (!present(node)) {
      return empty, empty
    }
    Node<#k,#v> node2 <- require(node)
    if (present(node2.getHigher())) {
      optional Node<#k,#v> temp, removed <- removeHighest(node2.getHigher())
      \ node2.setHigher(temp)
      newNode <- rebalance(node2)
    } else {
      newNode <- node2.getLower()
      \ node2.setLower(empty)
      removed <- node
    }
  }

  @type removeLowest (optional Node<#k,#v>) -> (optional Node<#k,#v>,optional Node<#k,#v>)
  removeLowest (node) (newNode,removed) {
    if (!present(node)) {
      return empty, empty
    }
    Node<#k,#v> node2 <- require(node)
    if (present(node2.getLower())) {
      optional Node<#k,#v> temp, removed <- removeLowest(node2.getLower())
      \ node2.setLower(temp)
      newNode <- rebalance(node2)
    } else {
      newNode <- node2.getHigher()
      \ node2.setHigher(empty)
      removed <- node
    }
  }

  @type swapChildren (Node<#k,#v>,Node<#k,#v>) -> ()
  swapChildren (l,r) {
    scoped {
      optional Node<#k,#v> temp <- l.getLower()
      \ l.setLower(r.getLower())
    } in \ r.setLower(temp)
    scoped {
      optional Node<#k,#v> temp <- l.getHigher()
      \ l.setHigher(r.getHigher())
    } in \ r.setHigher(temp)
  }

  @value getHeight () -> (Int)
  getHeight () { return height }

  @value setValue (#v) -> ()
  setValue (v) { value <- v }

  @value setLower (optional Node<#k,#v>) -> ()
  setLower (l) { lower <- l }

  @value setHigher (optional Node<#k,#v>) -> ()
  setHigher (h) { higher <- h }
}

concrete TreeKeyValue<#k,#v> {
  refines KeyValue<#k,#v>

  @category create<#k,#v> (#k,#v) -> (TreeKeyValue<#k,#v>)
}

define TreeKeyValue {
  @value #k key
  @value #v value

  create (key,value) {
    return TreeKeyValue<#k,#v>{ key, value }
  }

  getKey () {
    return key
  }

  getValue () {
    return value
  }
}

concrete ForwardTreeOrder<#k,#v> {
  refines Order<KeyValue<#k,#v>>
  #k defines LessThan<#k>

  @category create<#k,#v>
    #k defines LessThan<#k>
  (optional Node<#k,#v>) -> (optional ForwardTreeOrder<#k,#v>)
}

define ForwardTreeOrder {
  $ReadOnly[node,prev]$

  @value Node<#k,#v> node
  @value optional ForwardTreeOrder<#k,#v> prev

  create (node) {
    optional Node<#k,#v> node2 <- node
    optional ForwardTreeOrder<#k,#v> current <- empty
    while (present(node2)) {
      current <- ForwardTreeOrder<#k,#v>{ require(node2), current }
      node2 <- require(node2).getLower()
    }
    return current
  }

  next () {
    // Algorithm:
    // 1. Pop self from the stack.
    // 2. Traverse lower to the bottom starting from the higher child of self.
    optional Node<#k,#v> node2 <- node.getHigher()
    optional ForwardTreeOrder<#k,#v> current <- prev
    while (present(node2)) {
      current <- ForwardTreeOrder<#k,#v>{ require(node2), current }
      node2 <- require(node2).getLower()
    }
    return current
  }

  get () {
    return TreeKeyValue:create<?,?>(node.getKey(),node.getValue())
  }
}

concrete ReverseTreeOrder<#k,#v> {
  refines Order<KeyValue<#k,#v>>
  #k defines LessThan<#k>

  @category create<#k,#v>
    #k defines LessThan<#k>
  (optional Node<#k,#v>) -> (optional ReverseTreeOrder<#k,#v>)
}

define ReverseTreeOrder {
  $ReadOnly[node,prev]$

  @value Node<#k,#v> node
  @value optional ReverseTreeOrder<#k,#v> prev

  create (node) {
    optional Node<#k,#v> node2 <- node
    optional ReverseTreeOrder<#k,#v> current <- empty
    while (present(node2)) {
      current <- ReverseTreeOrder<#k,#v>{ require(node2), current }
      node2 <- require(node2).getHigher()
    }
    return current
  }

  next () {
    // Algorithm:
    // 1. Pop self from the stack.
    // 2. Traverse higher to the bottom starting from the lower child of self.
    optional Node<#k,#v> node2 <- node.getLower()
    optional ReverseTreeOrder<#k,#v> current <- prev
    while (present(node2)) {
      current <- ReverseTreeOrder<#k,#v>{ require(node2), current }
      node2 <- require(node2).getHigher()
    }
    return current
  }

  get () {
    return TreeKeyValue:create<?,?>(node.getKey(),node.getValue())
  }
}