import Shared
import Foundation

private enum RBTreeColor {
    case red
    case black
}

private enum RotationDirection {
    case left
    case right
}

// MARK: - RBNode
public class RBTreeNode<T: Comparable>: Equatable {
    public typealias RBNode = RBTreeNode<T>

    fileprivate var color: RBTreeColor = .black
    fileprivate var key: T?
    var leftChild: RBNode?
    var rightChild: RBNode?
    fileprivate weak var parent: RBNode?

    public init(key: T?, leftChild: RBNode?, rightChild: RBNode?, parent: RBNode?) {
        self.key = key
        self.leftChild = leftChild
        self.rightChild = rightChild
        self.parent = parent

        self.leftChild?.parent = self
        self.rightChild?.parent = self
    }

    public convenience init(key: T?) {
        self.init(key: key, leftChild: RBNode(), rightChild: RBNode(), parent: RBNode())
    }

    // For initialising the nullLeaf
    public convenience init() {
        self.init(key: nil, leftChild: nil, rightChild: nil, parent: nil)
        self.color = .black
    }

    var isRoot: Bool {
        return parent == nil
    }

    var isLeaf: Bool {
        return rightChild == nil && leftChild == nil
    }

    var isNullLeaf: Bool {
        return key == nil && isLeaf && color == .black
    }

    var isLeftChild: Bool {
        return parent?.leftChild === self
    }

    var isRightChild: Bool {
        return parent?.rightChild === self
    }

    var grandparent: RBNode? {
        return parent?.parent
    }

    var sibling: RBNode? {
        if isLeftChild {
            return parent?.rightChild
        } else {
            return parent?.leftChild
        }
    }

    var uncle: RBNode? {
        return parent?.sibling
    }
}

// MARK: - RedBlackTree
public class RedBlackTree<T: Comparable> {
    public typealias RBNode = RBTreeNode<T>

    fileprivate(set) var root: RBNode
    fileprivate(set) var size = 0
    fileprivate let nullLeaf = RBNode()

    public init() {
        root = nullLeaf
    }
}

// MARK: - Equatable protocol
extension RBTreeNode {
    static public func ==<T>(lhs: RBTreeNode<T>, rhs: RBTreeNode<T>) -> Bool {
        return lhs.key == rhs.key
    }
}

// MARK: - Finding a nodes successor
extension RBTreeNode {
    /*
     * Returns the inorder successor node of a node
     * The successor is a node with the next larger key value of the current node
     */
    public func getSuccessor() -> RBNode? {
        // If node has right child: successor min of this right tree
        if let rightChild = self.rightChild {
            if !rightChild.isNullLeaf {
                return rightChild.minimum()
            }
        }
        // Else go upward until node left child
        var currentNode = self
        var parent = currentNode.parent
        while currentNode.isRightChild {
            if let parent = parent {
                currentNode = parent
            }
            parent = currentNode.parent
        }
        return parent
    }
}

// MARK: - Searching
extension RBTreeNode {
    /*
     * Returns the node with the minimum key of the current subtree
     */
    public func minimum() -> RBNode? {
        if let leftChild = leftChild {
            if !leftChild.isNullLeaf {
                return leftChild.minimum()
            }
            return self
        }
        return self
    }

    /*
     * Returns the node with the maximum key of the current subtree
     */
    public func maximum() -> RBNode? {
        if let rightChild = rightChild {
            if !rightChild.isNullLeaf {
                return rightChild.maximum()
            }
            return self
        }
        return self
    }
}

extension RedBlackTree {
    /*
     * Returns the node with the given key |input| if existing
     */
    public func search(input: T) -> RBNode? {
        return search(key: input, node: root)
    }

    /*
     * Returns the node with given |key| in subtree of |node|
     */
    fileprivate func search(key: T, node: RBNode?) -> RBNode? {
        // If node nil -> key not found
        guard let node = node else {
            return nil
        }
        // If node is nullLeaf == semantically same as if nil
        if !node.isNullLeaf {
            if let nodeKey = node.key {
                // Node found
                if key == nodeKey {
                    return node
                } else if key < nodeKey {
                    return search(key: key, node: node.leftChild)
                } else {
                    return search(key: key, node: node.rightChild)
                }
            }
        }
        return nil
    }
}

// MARK: - Finding maximum and minimum value
extension RedBlackTree {
    /*
     * Returns the minimum key value of the whole tree
     */
    public func minValue() -> T? {
        guard let minNode = root.minimum() else {
            return nil
        }
        return minNode.key
    }

    /*
     * Returns the maximum key value of the whole tree
     */
    public func maxValue() -> T? {
        guard let maxNode = root.maximum() else {
            return nil
        }
        return maxNode.key
    }
}

// MARK: - Inserting new nodes
extension RedBlackTree {
    /*
     * Insert a node with key |key| into the tree
     * 1. Perform normal insert operation as in a binary search tree
     * 2. Fix red-black properties
     * Runntime: O(log n)
     */
    public func insert(key: T) {
        if root.isNullLeaf {
            root = RBNode(key: key)
        } else {
            insert(input: RBNode(key: key), node: root)
        }

        size += 1
    }

    /*
     * Nearly identical insert operation as in a binary search tree
     * Differences: All nil pointers are replaced by the nullLeaf, we color the inserted node red,
     * after inserting we call insertFixup to maintain the red-black properties
     */
    private func insert(input: RBNode, node: RBNode) {
        guard let inputKey = input.key, let nodeKey = node.key else {
            return
        }
        if inputKey < nodeKey {
            guard let child = node.leftChild else {
                addAsLeftChild(child: input, parent: node)
                return
            }
            if child.isNullLeaf {
                addAsLeftChild(child: input, parent: node)
            } else {
                insert(input: input, node: child)
            }
        } else {
            guard let child = node.rightChild else {
                addAsRightChild(child: input, parent: node)
                return
            }
            if child.isNullLeaf {
                addAsRightChild(child: input, parent: node)
            } else {
                insert(input: input, node: child)
            }
        }
    }

    private func addAsLeftChild(child: RBNode, parent: RBNode) {
        parent.leftChild = child
        child.parent = parent
        child.color = .red
        insertFixup(node: child)
    }

    private func addAsRightChild(child: RBNode, parent: RBNode) {
        parent.rightChild = child
        child.parent = parent
        child.color = .red
        insertFixup(node: child)
    }

    /*
     * Fixes possible violations of the red-black property after insertion
     * Only violation of red-black properties occurs at inserted node |z| and z.parent
     * We have 3 distinct cases: case 1, 2a and 2b
     * - case 1: may repeat, but only h/2 steps, where h is the height of the tree
     * - case 2a -> case 2b -> red-black tree
     * - case 2b -> red-black tree
     */
    private func insertFixup(node z: RBNode) {
        if !z.isNullLeaf {
            guard let parentZ = z.parent else {
                return
            }
            // If both |z| and his parent are red -> violation of red-black property -> need to fix it
            if parentZ.color == .red {
                guard let uncle = z.uncle else {
                    return
                }
                // Case 1: Uncle red -> recolor + move z
                if uncle.color == .red {
                    parentZ.color = .black
                    uncle.color = .black
                    if let grandparentZ = parentZ.parent {
                        grandparentZ.color = .red
                        // Move z to grandparent and check again
                        insertFixup(node: grandparentZ)
                    }
                }
                // Case 2: Uncle black
                else {
                    var zNew = z
                    // Case 2.a: z right child -> rotate
                    if parentZ.isLeftChild && z.isRightChild {
                        zNew = parentZ
                        leftRotate(node: zNew)
                    } else if parentZ.isRightChild && z.isLeftChild {
                        zNew = parentZ
                        rightRotate(node: zNew)
                    }
                    // Case 2.b: z left child -> recolor + rotate
                    zNew.parent?.color = .black
                    if let grandparentZnew = zNew.grandparent {
                        grandparentZnew.color = .red
                        if z.isLeftChild {
                            rightRotate(node: grandparentZnew)
                        } else {
                            leftRotate(node: grandparentZnew)
                        }
                        // We have a valid red-black-tree
                    }
                }
            }
        }
        root.color = .black
    }
}

// MARK: - Deleting a node
extension RedBlackTree {
    /*
     * Delete a node with key |key| from the tree
     * 1. Perform standard delete operation as in a binary search tree
     * 2. Fix red-black properties
     * Runntime: O(log n)
     */
    public func delete(key: T) {
        if size == 1 {
            root = nullLeaf
            size -= 1
        } else if let node = search(key: key, node: root) {
            if !node.isNullLeaf {
                delete(node: node)
                size -= 1
            }
        }
    }

    /*
     * Nearly identical delete operation as in a binary search tree
     * Differences: All nil pointers are replaced by the nullLeaf,
     * after deleting we call insertFixup to maintain the red-black properties if the delted node was
     * black (as if it was red -> no violation of red-black properties)
     */
    private func delete(node z: RBNode) {
        var nodeY = RBNode()
        var nodeX = RBNode()
        if let leftChild = z.leftChild, let rightChild = z.rightChild {
            if leftChild.isNullLeaf || rightChild.isNullLeaf {
                nodeY = z
            } else {
                if let successor = z.getSuccessor() {
                    nodeY = successor
                }
            }
        }
        if let leftChild = nodeY.leftChild {
            if !leftChild.isNullLeaf {
                nodeX = leftChild
            } else if let rightChild = nodeY.rightChild {
                nodeX = rightChild
            }
        }
        nodeX.parent = nodeY.parent
        if let parentY = nodeY.parent {
            // Should never be the case, as parent of root = nil
            if parentY.isNullLeaf {
                root = nodeX
            } else {
                if nodeY.isLeftChild {
                    parentY.leftChild = nodeX
                } else {
                    parentY.rightChild = nodeX
                }
            }
        } else {
            root = nodeX
        }
        if nodeY != z {
            z.key = nodeY.key
        }
        // If sliced out node was red -> nothing to do as red-black-property holds
        // If it was black -> fix red-black-property
        if nodeY.color == .black {
            deleteFixup(node: nodeX)
        }
    }

    /*
     * Fixes possible violations of the red-black property after deletion
     * We have w distinct cases: only case 2 may repeat, but only h many steps, where h is the height
     * of the tree
     * - case 1 -> case 2 -> red-black tree
     *   case 1 -> case 3 -> case 4 -> red-black tree
     *   case 1 -> case 4 -> red-black tree
     * - case 3 -> case 4 -> red-black tree
     * - case 4 -> red-black tree
     */
    private func deleteFixup(node x: RBNode) {
        var xTmp = x
        if !x.isRoot && x.color == .black {
            guard var sibling = x.sibling else {
                return
            }
            // Case 1: Sibling of x is red
            if sibling.color == .red {
                // Recolor
                sibling.color = .black
                if let parentX = x.parent {
                    parentX.color = .red
                    // Rotation
                    if x.isLeftChild {
                        leftRotate(node: parentX)
                    } else {
                        rightRotate(node: parentX)
                    }
                    // Update sibling
                    if let sibl = x.sibling {
                        sibling = sibl
                    }
                }
            }
            // Case 2: Sibling is black with two black children
            if sibling.leftChild?.color == .black && sibling.rightChild?.color == .black {
                // Recolor
                sibling.color = .red
                // Move fake black unit upwards
                if let parentX = x.parent {
                    deleteFixup(node: parentX)
                }
                // We have a valid red-black-tree
            } else {
                // Case 3: a. Sibling black with one black child to the right
                if x.isLeftChild && sibling.rightChild?.color == .black {
                    // Recolor
                    sibling.leftChild?.color = .black
                    sibling.color = .red
                    // Rotate
                    rightRotate(node: sibling)
                    // Update sibling of x
                    if let sibl = x.sibling {
                        sibling = sibl
                    }
                }
                // Still case 3: b. One black child to the left
                else if x.isRightChild && sibling.leftChild?.color == .black {
                    // Recolor
                    sibling.rightChild?.color = .black
                    sibling.color = .red
                    // Rotate
                    leftRotate(node: sibling)
                    // Update sibling of x
                    if let sibl = x.sibling {
                        sibling = sibl
                    }
                }
                // Case 4: Sibling is black with red right child
                // Recolor
                if let parentX = x.parent {
                    sibling.color = parentX.color
                    parentX.color = .black
                    // a. x left and sibling with red right child
                    if x.isLeftChild {
                        sibling.rightChild?.color = .black
                        // Rotate
                        leftRotate(node: parentX)
                    }
                    // b. x right and sibling with red left child
                    else {
                        sibling.leftChild?.color = .black
                        //Rotate
                        rightRotate(node: parentX)
                    }
                    // We have a valid red-black-tree
                    xTmp = root
                }
            }
        }
        xTmp.color = .black
    }
}

// MARK: - Rotation
extension RedBlackTree {
    /*
     * Left rotation around node x
     * Assumes that x.rightChild y is not a nullLeaf, rotates around the link from x to y,
     * makes y the new root of the subtree with x as y's left child and y's left child as x's right
     * child, where n = a node, [n] = a subtree
     *     |                |
     *     x                y
     *   /   \     ~>     /   \
     * [A]    y          x    [C]
     *       / \        / \
     *     [B] [C]    [A] [B]
     */
    fileprivate func leftRotate(node x: RBNode) {
        rotate(node: x, direction: .left)
    }

    /*
     * Right rotation around node y
     * Assumes that y.leftChild x is not a nullLeaf, rotates around the link from y to x,
     * makes x the new root of the subtree with y as x's right child and x's right child as y's left
     * child, where n = a node, [n] = a subtree
     *     |                |
     *     x                y
     *   /   \     <~     /   \
     * [A]    y          x    [C]
     *       / \        / \
     *     [B] [C]    [A] [B]
     */
    fileprivate func rightRotate(node x: RBNode) {
        rotate(node: x, direction: .right)
    }

    /*
     * Rotation around a node x
     * Is a local operation preserving the binary-search-tree property that only exchanges pointers.
     * Runntime: O(1)
     */
    private func rotate(node x: RBNode, direction: RotationDirection) {
        var nodeY: RBNode? = RBNode()

        // Set |nodeY| and turn |nodeY|'s left/right subtree into |x|'s right/left subtree
        switch direction {
        case .left:
            nodeY = x.rightChild
            x.rightChild = nodeY?.leftChild
            x.rightChild?.parent = x
        case .right:
            nodeY = x.leftChild
            x.leftChild = nodeY?.rightChild
            x.leftChild?.parent = x
        }

        // Link |x|'s parent to nodeY
        nodeY?.parent = x.parent
        if x.isRoot {
            if let node = nodeY {
                root = node
            }
        } else if x.isLeftChild {
            x.parent?.leftChild = nodeY
        } else if x.isRightChild {
            x.parent?.rightChild = nodeY
        }

        // Put |x| on |nodeY|'s left
        switch direction {
        case .left:
            nodeY?.leftChild = x
        case .right:
            nodeY?.rightChild = x
        }
        x.parent = nodeY
    }
}

// MARK: - Verify
extension RedBlackTree {
    /*
     * Verifies that the existing tree fulfills all red-black properties
     * Returns true if the tree is a valid red-black tree, false otherwise
     */
    public func verify() -> Bool {
        if root.isNullLeaf {
            print("The tree is empty")
            return true
        }
        return property2() && property4() && property5()
    }

    // Property 1: Every node is either red or black -> fullfilled through setting node.color of type
    // RBTreeColor

    // Property 2: The root is black
    private func property2() -> Bool {
        if root.color == .red {
            print("Property-Error: Root is red")
            return false
        }
        return true
    }

    // Property 3: Every nullLeaf is black -> fullfilled through initialising nullLeafs with color = black

    // Property 4: If a node is red, then both its children are black
    private func property4() -> Bool {
        return property4(node: root)
    }

    private func property4(node: RBNode) -> Bool {
        if node.isNullLeaf {
            return true
        }
        if let leftChild = node.leftChild, let rightChild = node.rightChild {
            if node.color == .red {
                if !leftChild.isNullLeaf && leftChild.color == .red {
                    print("Property-Error: Red node with key \(String(describing: node.key)) has red left child")
                    return false
                }
                if !rightChild.isNullLeaf && rightChild.color == .red {
                    print("Property-Error: Red node with key \(String(describing: node.key)) has red right child")
                    return false
                }
            }
            return property4(node: leftChild) && property4(node: rightChild)
        }
        return false
    }

    // Property 5: For each node, all paths from the node to descendant leaves contain the same number
    // of black nodes (same blackheight)
    private func property5() -> Bool {
        if property5(node: root) == -1 {
            return false
        } else {
            return true
        }
    }

    private func property5(node: RBNode) -> Int {
        if node.isNullLeaf {
            return 0
        }
        guard let leftChild = node.leftChild, let rightChild = node.rightChild else {
            return -1
        }
        let left = property5(node: leftChild)
        let right = property5(node: rightChild)

        if left == -1 || right == -1 {
            return -1
        } else if left == right {
            let addedHeight = node.color == .black ? 1 : 0
            return left + addedHeight
        } else {
            print("Property-Error: Black height violated at node with key \(String(describing: node.key))")
            return -1
        }
    }
}

// MARK: - Debugging
extension RBTreeNode: CustomDebugStringConvertible {
    public var debugDescription: String {
        var s = ""
        if isNullLeaf {
            s = "nullLeaf"
        } else {
            if let key = key {
                s = "key: \(key)"
            } else {
                s = "key: nil"
            }
            if let parent = parent {
                s += ", parent: \(String(describing: parent.key))"
            }
            if let left = leftChild {
                s += ", left = [" + left.debugDescription + "]"
            }
            if let right = rightChild {
                s += ", right = [" + right.debugDescription + "]"
            }
            s += ", color = \(color)"
        }
        return s
    }
}

extension RedBlackTree: CustomDebugStringConvertible {
    public var debugDescription: String {
        return root.debugDescription
    }
}

extension RBTreeNode: CustomStringConvertible {
    public var description: String {
        var s = ""
        if isNullLeaf {
            s += "nullLeaf"
        } else {
            if let left = leftChild {
                s += "(\(left.description)) <- "
            }
            if let key = key {
                s += "\(key)"
            } else {
                s += "nil"
            }
            s += ", \(color)"
            if let right = rightChild {
                s += " -> (\(right.description))"
            }
        }
        return s
    }
}

extension RedBlackTree: CustomStringConvertible {
    public var description: String {
        if root.isNullLeaf {
            return "[]"
        } else {
            return root.description
        }
    }
}

struct Node {
    let label: Character
    var inEdges: Set<Character> = []
    var outEdges: Set<Character> = []
}

struct Graph {
    var nodes: [Character: Node] = [:]
}

if let input: [String] = readInput(day: 7) {
    var graph = Graph()
    for line in input {
        let scanner = Scanner(string: line)
        scanner.scanString("Step", into: nil)
        var nodeString: NSString? = ""
        scanner.scanCharacters(from: .letters, into: &nodeString)
        let nodeA = String(nodeString!).first!
        scanner.scanString("must be finished before step", into: nil)
        scanner.scanCharacters(from: .letters, into: &nodeString)
        let nodeB = String(nodeString!).first!

        var firstNode = graph.nodes[nodeA] ?? Node(label: nodeA, inEdges: [], outEdges: [])
        var secondNode = graph.nodes[nodeB] ?? Node(label: nodeB, inEdges: [], outEdges: [])
        firstNode.outEdges.insert(nodeB)
        secondNode.inEdges.insert(nodeA)

        graph.nodes[nodeA] = firstNode
        graph.nodes[nodeB] = secondNode
    }

    let redBlackTree = RedBlackTree<Character>()
    // Initialize it with all nodes without in edges
    for (key, node) in graph.nodes {
        if node.inEdges.isEmpty {
            redBlackTree.insert(key: key)
        }
    }
    var output = ""

    while redBlackTree.size > 0 {
        // Pick the minimum node
        if let nodeKey = redBlackTree.minValue(), let node = graph.nodes[nodeKey] {
            // Add it to the output
            output.append(nodeKey)

            // Remove it from the RB tree
            redBlackTree.delete(key: nodeKey)

            // Check all next nodes
            for nextNodeKey in node.outEdges {
                if var nextNode = graph.nodes[nextNodeKey] {
                    // Remove the edge from its in edges
                    nextNode.inEdges.remove(nodeKey)
                    graph.nodes[nextNodeKey] = nextNode

                    // If the graph has no more in edges, consider it for the next round
                    if nextNode.inEdges.isEmpty {
                        redBlackTree.insert(key: nextNodeKey)
                    }
                }
            }
        }
    }

    print(output)
}