package algs33;
import stdlib.*;
import java.util.Iterator;
import java.util.NoSuchElementException;
import algs13.Stack;
/* ***********************************************************************
 *  Compilation:  javac RandomizedBST.java
 *  Execution:    java RandomizedBST
 *
 *  Symbol table (map) implemented with a randomized BST.
 *
 *
 *************************************************************************/

public class XRandomizedBST<K extends Comparable<? super K>, V> implements Iterable<K> {

  private Node<K,V> root;   // root of the BST

  // BST helper node data type
  private static class Node<K,V> {
    public K key;          // key
    public V val;          // associated data
    public Node<K,V> left, right;   // left and right subtrees
    public int N;              // node count of descendents

    public Node(K key, V val) {
      this.key = key;
      this.val = val;
      this.N   = 1;
    }
  }


  /* ***********************************************************************
   *  BST search
   *************************************************************************/

  public boolean contains(K key) {
    return (get(key) != null);
  }

  // return value associated with the given key
  // if no such value, return null
  // if multiple such values, return first one on path from root
  public V get(K key) {
    return get(root, key);
  }

  private V get(Node<K,V> x, K key) {
    if (x == null) return null;
    int cmp = key.compareTo(x.key);
    if      (cmp == 0) return x.val;
    else if (cmp  < 0) return get(x.left,  key);
    else               return get(x.right, key);
  }


  /* ***********************************************************************
   *  randomized insertion
   *************************************************************************/
  public void put(K key, V val) {
    root = put(root, key, val);
  }

  // make new node the root with uniform probability
  private Node<K,V> put(Node<K,V> x, K key, V val) {
    if (x == null) return new Node<>(key, val);
    int cmp = key.compareTo(x.key);
    if (cmp == 0) { x.val = val; return x; }
    if (StdRandom.bernoulli(1.0 / (size(x) + 1.0))) return putRoot(x, key, val);
    if (cmp < 0) x.left  = put(x.left,  key, val);
    else         x.right = put(x.right, key, val);
    // (x.N)++;
    fix(x);
    return x;
  }


  private Node<K,V> putRoot(Node<K,V> x, K key, V val) {
    if (x == null) return new Node<>(key, val);
    int cmp = key.compareTo(x.key);
    if      (cmp == 0) { x.val = val; return x; }
    else if (cmp  < 0) { x.left  = putRoot(x.left,  key, val); x = rotR(x); }
    else               { x.right = putRoot(x.right, key, val); x = rotL(x); }
    return x;
  }



  /* ***********************************************************************
   *  deletion
   *************************************************************************/
  private Node<K,V> joinLR(Node<K,V> a, Node<K,V> b) {
    if (a == null) return b;
    if (b == null) return a;

    if (StdRandom.bernoulli((double) size(a) / (size(a) + size(b))))  {
      a.right = joinLR(a.right, b);
      fix(a);
      return a;
    }
    else {
      b.left = joinLR(a, b.left);
      fix(b);
      return b;
    }
  }

  private Node<K,V> remove(Node<K,V> x, K key) {
    if (x == null) return null;
    int cmp = key.compareTo(x.key);
    if      (cmp == 0) x = joinLR(x.left, x.right);
    else if (cmp  < 0) x.left  = remove(x.left,  key);
    else               x.right = remove(x.right, key);
    fix(x);
    return x;
  }

  // remove and return value associated with given key; if no such key, return null
  public V remove(K key) {
    V val = get(key);
    root = remove(root, key);
    return val;
  }

  /* ***********************************************************************
   *  Selection
   *************************************************************************/

  // return the kth largest key
  public K select(int k) { Node<K,V> x = select(root, k); return x.key; }
  private Node<K,V> select(Node<K,V> x, int k) {
    if (x == null) return null;
    int t = size(x.left);
    if      (t > k) return select(x.left,  k);
    else if (t < k) return select(x.right, k-t-1);
    else            return x;
  }



  // return the smallest key
  public K min() {
    K key = null;
    for (Node<K,V> x = root; x != null; x = x.left)
      key = x.key;
    return key;
  }

  // return the largest key
  public K max() {
    K key = null;
    for (Node<K,V> x = root; x != null; x = x.right)
      key = x.key;
    return key;
  }

  // return the smallest key >= query key; if no such key return null
  public K ceil(K key) {
    Node<K,V> best = ceil(root, key, null);
    if (best == null) return null;
    return best.key;
  }
  private Node<K,V> ceil(Node<K,V> x, K key, Node<K,V> best) {
    if      (x == null)        return best;
    else if (eq(key, x.key))   return x;
    else if (less(key, x.key)) return ceil(x.left,  key, x);
    else                       return ceil(x.right, key, best);
  }

  // return the smallest key >= query key; if no such key return null
  public K ceil2(K key) {
    Node<K,V> best = null;
    Node<K,V> x = root;
    while (x != null) {
      int cmp = key.compareTo(x.key);
      if      (cmp < 0) { best = x; x = x.left; }
      else if (cmp > 0) { x = x.right;          }
      else              return x.key;
    }
    if (best == null) return null;
    return best.key;
  }


  /* *********************************************************************
   *  Iterate using inorder traversal using a stack.
   *  Iterating through N elements takes O(N) time.
   ***********************************************************************/
  public Iterator<K> iterator() { return new BSTIterator(root); }

  // an iterator
  private class BSTIterator implements Iterator<K> {
    private Stack<Node<K,V>> stack = new Stack<>();

    public BSTIterator(Node<K,V> x) {
      while (x != null) {
        stack.push(x);
        x = x.left;
      }
    }

    public boolean hasNext()  { return !stack.isEmpty();                    }

    // it's optional and we don't want to support it
    public void remove()      { throw new UnsupportedOperationException();  }

    public K next() {
      if (!hasNext()) throw new NoSuchElementException();
      Node<K,V> x = stack.pop();
      K key = x.key;
      x = x.right;
      while (x != null) {
        stack.push(x);
        x = x.left;
      }
      return key;
    }
  }




  /* ***********************************************************************
   *  Utility functions.
   *************************************************************************/

  // return number of nodes in subtree rooted at x
  public int size() { return size(root); }
  private int size(Node<K,V> x) {
    if (x == null) return 0;
    else           return x.N;
  }

  // height of tree (empty tree height = 0)
  public int height() { return height(root); }
  private int height(Node<K,V> x) {
    if (x == null) return 0;
    return 1 + Math.max(height(x.left), height(x.right));
  }


  /* ***********************************************************************
   *  helper BST functions
   *************************************************************************/

  // fix subtree count field
  private void fix(Node<K,V> x) {
    if (x == null) return;
    x.N = 1 + size(x.left) + size(x.right);
  }

  // right rotate
  private Node<K,V> rotR(Node<K,V> h) {
    Node<K,V> x = h.left;
    h.left = x.right;
    x.right = h;
    fix(h);
    fix(x);
    return x;
  }

  // left rotate
  private Node<K,V> rotL(Node<K,V> h) {
    Node<K,V> x = h.right;
    h.right = x.left;
    x.left = h;
    fix(h);
    fix(x);
    return x;
  }


  /* ***********************************************************************
   *  Debugging functions that test the integrity of the tree
   *************************************************************************/

  // check integrity of subtree count fields
  public boolean check() { return checkCount() && isBST(); }

  // check integrity of count fields
  private boolean checkCount() { return checkCount(root); }
  private boolean checkCount(Node<K,V> x) {
    if (x == null) return true;
    return checkCount(x.left) && checkCount(x.right) && (x.N == 1 + size(x.left) + size(x.right));
  }


  // does this tree satisfy the BST property?
  private boolean isBST() { return isBST(root, min(), max()); }

  // are all the values in the BST rooted at x between min and max, and recursively?
  private boolean isBST(Node<K,V> x, K min, K max) {
    if (x == null) return true;
    if (less(x.key, min) || less(max, x.key)) return false;
    return isBST(x.left, min, x.key) && isBST(x.right, x.key, max);
  }



  /* ***********************************************************************
   *  helper comparison functions
   *************************************************************************/

  private boolean less(K k1, K k2) {
    return k1.compareTo(k2) < 0;
  }

  private boolean eq(K k1, K k2) {
    return k1.compareTo(k2) == 0;
  }



  /* ***********************************************************************
   *  test client
   *************************************************************************/
  public static void main(String[] args) {
    XRandomizedBST<String, String> st = new XRandomizedBST<>();

    // insert some key-value pairs
    st.put("www.cs.princeton.edu",   "128.112.136.11");
    st.put("www.cs.princeton.edu",   "128.112.136.35");    // overwrite old value
    st.put("www.princeton.edu",      "128.112.130.211");
    st.put("www.math.princeton.edu", "128.112.18.11");
    st.put("www.yale.edu",           "130.132.51.8");
    st.put("www.amazon.com",         "207.171.163.90");
    st.put("www.simpsons.com",       "209.123.16.34");
    st.put("www.stanford.edu",       "171.67.16.120");
    st.put("www.google.com",         "64.233.161.99");
    st.put("www.ibm.com",            "129.42.16.99");
    st.put("www.apple.com",          "17.254.0.91");
    st.put("www.slashdot.com",       "66.35.250.150");
    st.put("www.whitehouse.gov",     "204.153.49.136");
    st.put("www.espn.com",           "199.181.132.250");
    st.put("www.snopes.com",         "66.165.133.65");
    st.put("www.movies.com",         "199.181.132.250");
    st.put("www.cnn.com",            "64.236.16.20");
    st.put("www.iitb.ac.in",         "202.68.145.210");


    StdOut.println(st.get("www.cs.princeton.edu"));
    StdOut.println(st.get("www.harvardsucks.com"));
    StdOut.println(st.get("www.simpsons.com"));
    StdOut.println();

    StdOut.println("integrity check: " + st.check());
    StdOut.println();

    StdOut.println("ceil(www.simpsonr.com) = " + st.ceil("www.simpsonr.com"));
    StdOut.println("ceil(www.simpsons.com) = " + st.ceil("www.simpsons.com"));
    StdOut.println("ceil(www.simpsont.com) = " + st.ceil("www.simpsont.com"));

    StdOut.println("ceil(www.simpsonr.com) = " + st.ceil2("www.simpsonr.com"));
    StdOut.println("ceil(www.simpsons.com) = " + st.ceil2("www.simpsons.com"));
    StdOut.println("ceil(www.simpsont.com) = " + st.ceil2("www.simpsont.com"));
    StdOut.println();

    for (int i = 0; i < st.size(); i++) {
      StdOut.println(i + "th: key  " + st.select(i));
    }
    StdOut.println();

    StdOut.println("min key: " + st.min());
    StdOut.println("max key: " + st.max());
    StdOut.println("size:    " + st.size());
    StdOut.println("height:  " + st.height());
    StdOut.println();
  }

}