package trees;

import java.io.*;
import java.util.*;

public class BinaryTree implements Iterable {

    /**
     * Note that for each recursive method in Node, there's a method in BTree.
     */
    public class Node extends BaseNode {

        Node(char c) {
            this.c = c;
        }

        /**
         * In a tree, everything is done recursively. First we add...
         */
        boolean add(char c) {
            if (c > this.c) {
                if (right == null) {
                    System.out.printf("+++ data = %c is added as a rightChild of %s%n", c, this);
                    right = new Node(c);
                } else {
                    System.out.printf("    following right child of %s%n", c, this);
                    right.add(c);
                }
            } else if (c < this.c) {
                if (left == null) {
                    System.out.printf("+++ data = %c is added as a leftChild of %s%n", c, this);
                    left = new Node(c);
                } else {
                    System.out.printf("    following left child of %s%n", c, this);
                    left.add(c);
                }
            }
            return true;
        }

        @Override
        void toString(StringBuilder sb) {
            sb.append('(');
            if (left == null) {
                sb.append('-');
            } else {
                left.toString(sb);
            }
            sb.append(' ');
            sb.append(c);
            sb.append(' ');
            if (right == null) {
                sb.append('-');
            } else {
                right.toString(sb);
            }
            sb.append(')');
        }
        /**
         * Print out a textual representation of the tree.
         */
        public String toString() {
            StringBuilder sb = new StringBuilder();
            toString(sb);
            return sb.toString();
        }

        /**
         * Determine the height of the tree.
         */
        int height() {
            int hr = 0, hl = 0;
            if (right != null) {
                hr = right.height();
            }
            if (left != null) {
                hl = left.height();
            }
            if (hr > hl) {
                return 1 + hr;
            } else {
                return 1 + hl;
            }
        }

        void draw(PrintWriter pw) {
            pw.printf("<table border=1 cellpadding=5 cellspacing=0>%n");
            if (c == ' ') {
                pw.printf("<tr><td class='top' colspan=2>[Space]</td></tr>%n");
            } else {
                pw.printf("<tr><td class='top' colspan=2>%c</td></tr>%n", c);
            }

            pw.printf("<tr><td class='left'>");
            if (left != null) {
                ((Node) left).draw(pw);
            } else {
                pw.printf("&empty;");
            }

            pw.printf("</td><td class='right'>");
            if (right != null) {
                ((Node) right).draw(pw);
            } else {
                pw.printf("&empty;");
            }

            pw.printf("</td></tr></table>%n");
        }

        Node findSmallest() {
            if (left != null) {
                return ((Node) left).findSmallest();
            }
            return this;
        }

        BaseNode remove(char c) {
            if (this.c == c) {
                // No children, removal is easy
                if (left == null && right == null) {
                    return null;
                }
                // left is null, removal is easy
                if (left == null) {
                    return right;
                }
                // right is null, removal is easy
                if (right == null) {
                    return left;
                }
                // Removal is harder. Find the smallest. It will
                // be the new head of this subtree.
                BaseNode smallest = ((Node) right).findSmallest();
                //          2
                //         /  \
                //        1    \
                //              5
                //             / \
                //            3   6
                //             \
                //              4
                //          3
                //         /  \
                //        1    \
                //              5
                //             / \
                //            4   6
                //
                // Remove the smallest from the right tree...
                BaseNode removedRight = ((Node) right).remove(smallest.c); // Node 3 removed from 5
                // Make the right tree the right child of the smallest...
                smallest.right = removedRight; // 3.right = 5
                smallest.left = left; // 3.left = 1
                return smallest;
            } else if (c < this.c && left != null) {
                left = ((Node) left).remove(c);
            } else if (c > this.c && right != null) {
                right = ((Node) right).remove(c);
            }
            return this;
        }
    }
    Node head;

    public int height() {
        if (head == null) {
            return 0;
        } else {
            return head.height();
        }
    }

    public int nodeCount() {
        if (head == null) {
            return 0;
        } else {
            return head.nodeCount();
        }
    }

    public boolean add(char c) {
        System.out.printf("Adding data = %c%n", c);
        if (head == null) {
            head = new Node(c);
        } else {
            head.add(c);
        }
        System.out.printf("New tree = %s%n%n", head);
        return true;
    }

    public void remove(char c) {
        if (head == null) {
            return;
        }
        head = (Node) ((Node) head).remove(c);
    }

    public String toString() {
        if (head == null) {
            return "()";
        }
        return head.toString();
    }

    /**
     * Draw the tree in HTML. It's easier to understand when viewed. Note: you
     * are not expected to learn how to write HTML for this course.
     */
    public void draw(String fname) throws IOException {
        FileWriter fw = new FileWriter(fname);
        PrintWriter pw = new PrintWriter(fw);
        pw.printf("<html>%n");
        pw.printf("<head>%n");
        pw.printf("<style>%n");
        pw.printf(".top { text-align: center; background-color: #FFDDFF; }");
        pw.printf(".left { text-align: center; background-color: #DDFFFF; width: 50%%; }");
        pw.printf(".right { text-align: center; background-color: #FFFFDD; width: 50%%; }");
        pw.printf("</style>%n");
        pw.printf("</head>%n");
        pw.printf("<body>%n");
        head.draw(pw);
        pw.printf("</body>%n");
        pw.printf("</html>%n");
        pw.close();
    }

    /**
     * Determine if the character c is contained within the BTree.
     */
    public boolean contains(char c) {
        if (head == null) {
            return false;
        } else {
            return head.contains(c);
        }
    }

    public void visit(BinaryTreeVisitor v) {
        if (head != null) {
            head.visit(v);
        }
    }

    public static BinaryTree test(String add, String rmv) {
        BinaryTree bt = new BinaryTree();
        Set<Character> sc = new HashSet<>();
        for (int i = 0; i < add.length(); i++) {
            char c = add.charAt(i);
            bt.add(c);
            sc.add(c);
            assert bt.nodeCount() == sc.size() : String.format("%d == %d: %c", bt.nodeCount(), sc.size(), c);
        }
        for (int i = 0; i < rmv.length(); i++) {
            char c = rmv.charAt(i);
            if (bt.contains(c)) {
                bt.remove(c);
            }
            if (sc.contains(c)) {
                sc.remove(c);
            }
            assert bt.nodeCount() == sc.size();
        }
        return bt;
    }

    public static void main(String[] args) throws Exception {
        test("apple", "ape");
        test("avacado", "down");
        test("window", "glow");
        BinaryTree bt = test("the_quick_brown_fox_jumped_over_the_lazy_dog", "");
        bt.draw("tree.html");
        System.out.printf("Tree height = %d%n", bt.height());
        System.out.printf("Tree node count = %d%n", bt.nodeCount());
        bt.visit(new BinaryTreeVisitor() {
            public void visit(BaseNode n) {
                System.out.println("c: " + n.c);
            }
        });
        int count = 0;
        for (Object ch : bt) {
            System.out.println("iter: " + ch);
            count++;
        }
        System.out.printf("%d == %d%n", count, bt.nodeCount());
        try (final var pw = new PrintWriter("viz.dot")) {
            pw.println("digraph G {");
            bt.visit(new BinaryTreeVisitor() {
                @Override
                public void visit(BaseNode n) {
                    if (n.left != null) {
                        pw.printf("%s -> %s;%n", n.c, n.left.c);
                    }
                    if (n.right != null) {
                        pw.printf("%s -> %s;%n", n.c, n.right.c);
                    }
                }
            });
            pw.println("}");
        }
    }

    public Iterator iterator() {
        return new Iterator() {
            // instance variables on the anonymous inner class
            Stack<BaseNode> stack = new Stack<>();
            BaseNode node = head;

            public boolean hasNext() {
                return node != null;
            }

            public Object next() {
                // Store the value we want to return
                Object ret = node.c;

                // save the right node on the
                // stack so we can look at it
                // later.
                if (node.right != null) {
                    stack.push(node.right);
                }

                // Move to the left node
                node = node.left;

                // If we've run out of left nodes
                // check to see if there's anything
                // on our stack.
                if (node == null) {
                    if (!stack.empty()) {
                        node = stack.pop();
                    }
                }

                return ret;
            }
        };
    }
}