import java.util.Iterator;
import java.util.LinkedList;

// This implementation is based on the BST class provided by the OpenDSA project:
//
// Distributed under the MIT License
//
// Copyright (c) 2011-2020 - Ville Karavirta and Cliff Shaffer
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

/**
 * Non-self-balancing Binary search tree based dictionary class.
 *
 * This code is a modified version of the OpenDSA Binary Search Tree implementation.
 *
 */
public class BSTDictionary<K extends Comparable<K>, V> implements Dictionary<K, V> {
  protected BSTNode<K, V> root; // Root of the BST
  private int nodeCount; // Number of nodes in the BST

  /** Constructor. */
  BSTDictionary() {
    root = null;
    nodeCount = 0;
  }

  /** Reinitialize tree. */
  public void clear() {
    root = null;
    nodeCount = 0;
  }

  /**
   * Insert a record into the tree. Keys can be anything, but they must be Comparable.
   *
   * @param k The key to insert.
   * @param v The value to insert.
   */
  @Override
  public void put(K k, V v) {
    root = insertHelp(root, k, v);
  }

  /**
   * Remove a record from the tree.
   *
   * @param key The key value of record to remove
   * @return The value removed, null if there is none.
   */
  @Override
  public V remove(K key) {
    V temp = findHelp(root, key); // First find it
    if (temp != null) {
      root = removeHelp(root, key); // Now remove it
      nodeCount--;
    }
    return temp;
  }

  /**
   * Return the record with key k, null if none exists.
   *
   * @param key The key to find
   * @return The value associated with the key, or null if not found
   */
  @Override
  public V get(K key) {
    return findHelp(root, key);
  }

  /**
   * Return the number of records in the dictionary.
   *
   * @return The number of records in the dictionary
   */
  public int size() {
    return nodeCount;
  }

  /**
   * Return a record that matches the key.
   *
   * @param rt The root of the subtree
   * @param key The key to find
   * @return The value associated with the key, or null if not found
   */
  private V findHelp(BSTNode<K, V> rt, K key) {
    if (rt == null) {
      return null;
    }
    if (rt.key().compareTo(key) > 0) {
      return findHelp(rt.left(), key);
    } else if (rt.key().compareTo(key) == 0) {
      return rt.value();
    } else {
      return findHelp(rt.right(), key);
    }
  }


  /**
   * Return the current subtree, modified to contain the new item.
   *
   * @param rt The root of the subtree
   * @param k The key to insert
   * @param v The value to insert
   * @return The modified subtree
   */
  private BSTNode<K, V> insertHelp(BSTNode<K, V> rt, K k, V v) {
    if (rt == null) {
      nodeCount++;
      return new BSTNode<K, V>(k, v);
    }
    if (rt.key().compareTo(k) > 0) {
      rt.setLeft(insertHelp(rt.left(), k, v));
    } else if (rt.key().compareTo(k) == 0) {
      rt.setValue(v);
    } else {
      rt.setRight(insertHelp(rt.right(), k, v));
    }
    return rt;
  }

  /**
   * Delete the maximum valued element in a subtree.
   *
   * @param rt The root of the subtree
   * @return The modified subtree with the maximum element removed
   */
  private BSTNode<K, V> deleteMax(BSTNode<K, V> rt) {
    if (rt.right() == null) {
      return rt.left();
    }
    rt.setRight(deleteMax(rt.right()));
    return rt;
  }


  /**
   * Get the maximum valued element in a subtree.
   *
   * @param rt The root of the subtree
   * @return The node with the maximum key
   */
  private BSTNode<K, V> getMax(BSTNode<K, V> rt) {
    if (rt.right() == null) {
      return rt;
    }
    return getMax(rt.right());
  }

  /**
   * Remove a node with the given key.
   *
   * @param rt The root of the subtree
   * @param key The key to remove
   * @return The tree with the node removed
   */
  private BSTNode<K, V> removeHelp(BSTNode<K, V> rt, K key) {
    if (rt == null) {
      return null;
    }
    if (rt.key().compareTo(key) > 0) {
      rt.setLeft(removeHelp(rt.left(), key));
    } else if (rt.key().compareTo(key) < 0) {
      rt.setRight(removeHelp(rt.right(), key));
    } else { // Found it
      if (rt.left() == null) {
        return rt.right();
      } else if (rt.right() == null) {
        return rt.left();
      } else { // Two children
        BSTNode<K, V> temp = getMax(rt.left());
        rt.setKey(temp.key());
        rt.setValue(temp.value());
        rt.setLeft(deleteMax(rt.left()));
      }
    }
    return rt;
  }

  /**
   * Return a string representation of the dictionary in the format {key1=value1, key2=value2, ...}.
   * Keys are listed in sorted order (in-order traversal of the BST).
   *
   * @return A string representation of the dictionary
   */
  @Override
  public String toString() {
    StringBuilder sb = new StringBuilder();
    sb.append("{");
    boolean first = true;
    for (K key : this) {
      if (!first) {
        sb.append(", ");
      }
      sb.append(key);
      sb.append("=");
      sb.append(get(key));
      first = false;
    }
    sb.append("}");
    return sb.toString();
  }



  @Override
  public Iterator<K> iterator() {
    // The simplest way to to create an iterator is to perform a full traversal at the time the
    // iterator is constructed, storing all of the keys in a list. The `next` method can then just
    // return the next available element in the list. This is not a great implementation. It
    // requires O(n) space and O(n) time, even if we only end up iterating through 1 item.
    //
    // A *better* implementation would keep track of where we are in an in-order traversal between
    // calls to `next`. This can be accomplished using a stack to simulate the call stack that would
    // exist during an in-order traversal. You are free to do this the easy way, but I
    // encourage you to think about the smart way if you have time.
    //
    // You do not need to implement remove.
    return new BSTIterator();
  }

  /**
   * Stack-based iterator for BST that performs an in-order traversal.
   */
  private class BSTIterator implements Iterator<K> {
    private LinkedList<BSTNode<K, V>> stack;

    public BSTIterator() {
      stack = new LinkedList<>();
      pushLeft(root);
    }

    /** Push all left descendants onto the stack. */
    private void pushLeft(BSTNode<K, V> node) {
      while (node != null) {
        stack.push(node);
        node = node.left();
      }
    }

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

    @Override
    public K next() {
      if (!hasNext()) {
        throw new java.util.NoSuchElementException();
      }
      BSTNode<K, V> node = stack.pop();
      pushLeft(node.right());
      return node.key();
    }
  }



  private class BSTNode<K extends Comparable<K>, V> {
    private K key; // Key for this node
    private V value; // Key for this node
    private BSTNode<K, V> left; // Pointer to left child
    private BSTNode<K, V> right; // Pointer to right child

    // Constructors
    BSTNode() {
      left = right = null;
    }

    BSTNode(K key, V value) {
      left = right = null;
      this.key = key;
      this.value = value;
    }

    BSTNode(K key, V value, BSTNode<K, V> l, BSTNode<K, V> r) {
      left = l;
      right = r;
      this.key = key;
      this.value = value;
    }

    // Get and set the key
    public K key() {
      return key;
    }

    public void setKey(K k) {
      key = k;
    }

    // Get and set the value
    public V value() {
      return value;
    }

    public void setValue(V v) {
      value = v;
    }

    // Get and set the left child
    public BSTNode<K, V> left() {
      return left;
    }

    public void setLeft(BSTNode<K, V> p) {
      left = p;
    }

    // Get and set the right child
    public BSTNode<K, V> right() {
      return right;
    }

    public void setRight(BSTNode<K, V> p) {
      right = p;
    }

    // return TRUE if a leaf node, FALSE otherwise
    public boolean isLeaf() {
      return (left == null) && (right == null);
    }

    // This implementation is based on the BST class provided by the OpenDSA project:
    //
    // Distributed under the MIT License
    //
    // Copyright (c) 2011-2020 - Ville Karavirta and Cliff Shaffer
    //
    // Permission is hereby granted, free of charge, to any person obtaining a copy
    // of this software and associated documentation files (the "Software"), to deal
    // in the Software without restriction, including without limitation the rights
    // to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
    // copies of the Software, and to permit persons to whom the Software is
    // furnished to do so, subject to the following conditions:
    //
    // The above copyright notice and this permission notice shall be included in
    // all copies or substantial portions of the Software.
    //
    // THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
    // IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
    // FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
    // AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
    // LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
    // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
    // THE SOFTWARE.

  }

}