/**
 * This work is based on https://deegen1.github.io/wbtree/index.html
 *
 * WBTreeSetBase.py - v1.01
 * Copyright 2022 Alec Dee - MIT license - SPDX: MIT
 * deegen1.github.io - akdee144@gmail.com
 *
 * This class implements reasonable interface for weight balanced tree, but it is not java-opinionated
 * For example, it does not care about Iterators, sub-sets, submaps, spliterators etc
 *
 * We use this structure inside SortedSet and SortedMap classes.
 * WBTreeSet and WBTreeMap implement all java idioms, when core algorithmic stuff is stored here.
 *
 * Such decomposition makes few things somewhat less effective (e.g. map contains double-wrapped entries) but
 * it allows us to synchronize with python original source code.
 */
package net.ranides.assira.collection.maps;

import net.ranides.assira.generic.CompareUtils;

import java.io.Serializable;
import java.util.Comparator;
import java.util.function.Consumer;

class WBTreeMapBase<K,V> implements Serializable {

    private static final long serialVersionUID = 2L;

    public enum FindMode {
        EQ,
        EQG,
        LT,
        LE,
        GT,
        GE
    }

    private final Comparator<? super K> cmp;
    private final boolean unique;
    private Node<K,V> root;

    public WBTreeMapBase() {
        this(CompareUtils.comparator(), true);
    }

    public WBTreeMapBase(Comparator<? super K> cmp, boolean unique) {
        this.cmp = cmp;
        this.unique = unique;
        this.root = null;
    }

    public void clear() {
        this.root = null;
    }

    public int size() {
        return weight(root);
    }

    public Node<K,V> at(int i) {
        Node<K,V> node = root;
        int weight = size();
        if(i<0) {
            i += weight;
        }
        if(i<0 || i>=weight) {
            return null;
        }
        while (true) {
            Node<K,V> left = node.left;
            int lw = weight(left);
            if(i > lw) {
                i -= lw +1;
                node = node.right;
            }
            else if(i == lw) {
                break;
            } else {
                node = left;
            }
        }
        return node;
    }

    public void forEach(Consumer<K> consumer) {
        forEachNode(node -> consumer.accept(node.key));
    }

    void forEachNode(Consumer<Node<K,V>> consumer) {
        Node<K,V> node = first();
        while(node != null) {
            consumer.accept(node);
            node = next(node);
        }
    }

    public Node<K,V> first() {
        Node<K,V> node = root;
        Node<K,V> ret = null;

        while(node != null) {
            ret = node;
            node = node.left;
        }

        return ret;
    }

    public Node<K,V> last() {
        Node<K,V> node = root;
        Node<K,V> ret = null;

        while (node != null) {
            ret = node;
            node = node.right;
        }

        return ret;
    }

    public Node<K,V> find(K value) {
        return find(value, FindMode.EQ);
    }

    public Node<K,V> find(K value, FindMode mode) {
        Node<K,V> node = root;
        Node<K,V> ret = null;

        boolean lset = mode == FindMode.LT || mode == FindMode.LE;
        boolean rset = mode == FindMode.GT || mode == FindMode.GE;
        boolean eset = mode == FindMode.EQ || mode == FindMode.EQG || mode == FindMode.LE || mode == FindMode.GE;
        boolean eright = mode == FindMode.EQG || mode == FindMode.LE || mode == FindMode.GT;

        while(node != null) {
            int c = cmp.compare(node.key, value);
            if(c < 0) {
                if(lset) {
                    ret = node;
                }
                node = node.right;
            }
            else if(c > 0) {
                if(rset) {
                    ret = node;
                }
                node = node.left;
            }
            else {
                if(eset) {
                    ret = node;
                    if(unique) {
                        break;
                    }
                }
                node = eright ? node.right : node.left;
            }
        }

        return ret;
    }

    public Node<K,V> add(K key, V value) {
        Node<K,V> prev = null;
        Node<K,V> node = root;
        int c = 0;
        while(node != null) {
            c = cmp.compare(node.key, key);
            if(c == 0 && this.unique) {
                node.key = key;
                node.value = value;
                return node;
            }
            prev = node;
            if(c <= 0) {
                node = node.right;
            } else {
                node = node.left;
            }
        }
        node = new Node<>(key, value);
        node.parent = prev;
        if(prev == null) {
            this.root = node;
        } else {
            if(c <=0) {
                prev.right = node;
            } else {
                prev.left = node;
            }
            rebalance(prev);
        }
        return node;
    }

    public Node<K,V> remove(K value) {
        Node<K,V> node = find(value);
        if(node != null) {
            removeNode(node);
        }
        return node;
    }

    public void removeNode(Node<K,V> node) {
        Node<K,V> p = node.parent;
        Node<K,V> l = node.left;
        Node<K,V> r = node.right;
        Node<K,V> bal = null;
        Node<K,V> next = null;

        if(r == null) {
            // case 1
            bal = p;
            next = l;
            l = null;
        }
        else if(r.left != null) {
            // case 3
            next = r;
            while(next.left != null) {
                bal = next;
                next = next.left;
            }
            Node<K,V> c = next.right;
            bal.left = c;
            if(c != null) {
                c.parent = bal;
            }
        }
        else {
            // case 2
            bal = r;
            next = r;
            r = null;
        }
        // replace node with next
        if(p == null) {
            root = next;
        }
        else if(p.left == node) {
            p.left = next;
        }
        else {
            p.right = next;
        }

        if(next != null) {
            next.parent = p;
            if(l != null) {
                next.left = l;
                l.parent = next;
            }
            if(r != null) {
                next.right = r;
                r.parent = next;
            }
        }

        rebalance(bal);
    }

    void rebalance(Node<K,V> next) {
        while (next != null) {
            Node<K,V> node = next;
            Node<K,V> orig = next;
            next = node.parent;

            Node<K,V> l = node.left;
            Node<K,V> r = node.right;
            int lw = weight(l);
            int rw = weight(r);

            if(rw * 5 + 2< lw * 2) {
                if(weight(l.left) * 5 < lw *2) {
                    node.left = rotateLeft(l);
                }
                node = rotateRight(node);
            }
            else if (lw * 5 + 2 < rw * 2) {
                if( weight(r.right) * 5 < rw*2) {
                    node.right = rotateRight(r);
                }
                node = rotateLeft(node);
            }
            else {
                node.weight = lw + rw + 1;
                continue;
            }
            if(next == null) {
                this.root = node;
            }
            else if(next.left == orig) {
                next.left = node;
            }
            else {
                next.right = node;
            }
        }
    }

    public static int weight(Node<?,?> node) {
        return node != null ? node.weight : 0;
    }

    public static <K,V> Node<K,V> next(Node<K,V> node) {
        Node<K,V> child = node.right;
        if(child != null) {
            while(child != null) {
                node = child;
                child = child.left;
            }
        } else {
            while(node != null && node.right == child) {
                child = node;
                node = node.parent;
            }
        }
        return node;
    }

    public static <K,V> Node<K,V> prev(Node<K,V> node) {
        Node<K,V> child = node.left;
        if(child != null) {
            while(child != null) {
                node = child;
                child = child.right;
            }
        } else {
            while(node != null && node.left == child) {
                child = node;
                node = node.parent;
            }
        }
        return node;
    }

    public static <K,V> Node<K,V> rotateLeft(Node<K,V> a) {
        Node<K,V> b=a.right;
        Node<K,V> r=b.left;
        b.parent=a.parent;
        b.left=a;
        a.parent=b;
        a.right=r;
        if (r != null) {
            r.parent = a;
        }
        recomputeWeight(a);
        recomputeWeight(b);
        return b;
    }

    public static <K,V> Node<K,V> rotateRight(Node<K,V> a) {
        Node<K,V> b = a.left;
        Node<K,V> l = b.right;
        b.parent = a.parent;
        b.right = a;
        a.parent = b;
        a.left = l;
        if (l != null) {
            l.parent = a;
        }
        recomputeWeight(a);
        recomputeWeight(b);
        return b;
    }

    public static void recomputeWeight(Node<?,?> node) {
        node.weight = 1 + weight(node.left) + weight(node.right);
    }

    public static int indexOf(Node<?,?> node) {
        if(node == null) {
            return -1;
        }
        int i = -1;
        Node<?,?> prev = node.right;
        while (node != null) {
            if (node.right == prev) {
                i += weight(node.left) + 1;
            }
            prev = node;
            node = node.parent;
        }
        return i;
    }

    public static class Node<K, V> extends AMap.AEntry<K, V> implements Serializable {

        private static final long serialVersionUID = 2L;

        public int weight;
        public Node<K,V> parent;
        public Node<K,V> left;
        public Node<K,V> right;
        public K key;
        public V value;

        public Node(K key, V value) {
            this.weight = 1;
            this.parent = null;
            this.left = null;
            this.right = null;
            this.key = key;
            this.value = value;
        }

        @Override
        public K getKey() {
            return key;
        }

        @Override
        public V getValue() {
            return value;
        }

        @Override
        public V setValue(V value) {
            V prev = this.value;
            this.value = value;
            return prev;
        }
    }




}
