package cyclops.data.base;


import com.oath.cyclops.matching.Deconstruct.Deconstruct2;
import cyclops.companion.Comparators;
import cyclops.control.Option;
import cyclops.data.ImmutableList;
import cyclops.data.LazySeq;
import cyclops.data.Seq;
import cyclops.reactive.ReactiveSeq;
import lombok.AllArgsConstructor;
import cyclops.data.tuple.Tuple;
import cyclops.data.tuple.Tuple2;
import lombok.EqualsAndHashCode;

import java.io.Serializable;
import java.util.Arrays;
import java.util.Objects;
import java.util.function.Supplier;


@AllArgsConstructor
public final class HAMT<K, V> implements Serializable {

    private static final long serialVersionUID = 1L;
    static final int BITS_IN_INDEX = 5;
    static final int SIZE = (int) StrictMath.pow(2, BITS_IN_INDEX);
    static final int MIN_INDEX = 0;
    static final int MAX_INDEX = SIZE - 1;
    static final int MASK = (1 << BITS_IN_INDEX) - 1;

    public static <K, V> Node<K, V> empty() {
        return EmptyNode.Instance;
    }

    public interface Node<K, V> extends Serializable {


        default Node<K, V> put(K key, V value) {
            return plus(0, key.hashCode(), key, value);
        }

        default Option<V> get(K key) {
            return get(0, key.hashCode(), key);
        }

        default V getOrElse(K key, V alt) {
            return getOrElse(0, key.hashCode(), key, alt);
        }

        default boolean containsKey(K key) {
            return get(key).isPresent();
        }

        default Node<K, V> minus(K key) {
            return minus(0, key.hashCode(), key);
        }

        public Node<K, V> plus(int bitShiftDepth, int hash, K key, V value);

        public Option<V> get(int bitShiftDepth, int hash, K key);

        public V getOrElse(int bitShiftDepth, int hash, K key, V alt);

        public V getOrElseGet(int bitShiftDepth, int hash, K key, Supplier<? extends V> alt);

        public Node<K, V> minus(int bitShiftDepth, int hash, K key);

        int size();

        LazySeq<Tuple2<K, V>> lazyList();

        ReactiveSeq<Tuple2<K, V>> stream();
        default ReactiveSeq<Tuple2<K, V>> streamNaturalOrder(){
            return stream();
        }
    }


    public static final class EmptyNode<K, V> implements Node<K, V> {
        private static final long serialVersionUID = 1L;
        static final EmptyNode Instance = new EmptyNode();

        @Override
        public Node<K, V> plus(int bitShiftDepth, int hash, K key, V value) {
            return new ValueNode<>(hash, key, value);
        }

        @Override
        public Option<V> get(int bitShiftDepth, int hash, K key) {
            return Option.none();
        }

        @Override
        public V getOrElse(int bitShiftDepth, int hash, K key, V alt) {
            return alt;
        }

        @Override
        public V getOrElseGet(int bitShiftDepth, int hash, K key, Supplier<? extends V> alt) {
            return alt.get();
        }

        @Override
        public Node<K, V> minus(int bitShiftDepth, int hash, K key) {
            return this;
        }

        @Override
        public int size() {
            return 0;
        }

        @Override
        public LazySeq<Tuple2<K, V>> lazyList() {
            return LazySeq.empty();
        }

        @Override
        public ReactiveSeq<Tuple2<K, V>> stream() {
            return ReactiveSeq.empty();
        }

        public String toString() {
            return "[]";
        }

        @Override
        public int hashCode() {
            return 1;
        }

        @Override
        public boolean equals(Object obj) {
            return this == obj;
        }
    }

    @AllArgsConstructor
    @EqualsAndHashCode
    public static final class ValueNode<K, V> implements Node<K, V>, Deconstruct2<K, V> {
        private static final long serialVersionUID = 1L;
        private final int hash;
        public final K key;
        public final V value;

        @Override
        public Node<K, V> plus(int bitShiftDepth, int hash, K key, V value) {
            ValueNode<K, V> newNode = new ValueNode<>(hash, key, value);
            return isMatch(hash, key) ? newNode : merge(bitShiftDepth, newNode);
        }

        private Node<K, V> merge(int bitShiftDepth, ValueNode<K, V> that) {
            //hash each merge into a collision node if hashes are the same, otherwise store in new location under a BitsetNode
            if (hash == that.hash)
                return new CollisionNode<>(hash, Seq.of(Tuple.tuple(key, value), that.unapply()));
            //create new BitsetNode
            int mask1 = BitsetNode.mask(hash, bitShiftDepth);
            int mask2 = BitsetNode.mask(that.hash, bitShiftDepth);
            int posThis = BitsetNode.bitpos(mask1);
            int posThat = BitsetNode.bitpos(mask2);
            int newBitset = posThis | posThat;
            if (mask1 == mask2) { //collision
                Node<K, V> merged = merge(bitShiftDepth + BITS_IN_INDEX, that);
                return new BitsetNode<>(newBitset, 2, new Node[]{merged});
            }
            Node<K, V>[] ordered = mask1 < mask2 ? new Node[]{this, that} : new Node[]{that, this};
            return new BitsetNode<>(newBitset, 2, ordered);
        }


        @Override
        public Option<V> get(int bitShiftDepth, int hash, K key) {
            return isMatch(hash, key) ? Option.of(value) : Option.none();
        }

        @Override
        public V getOrElse(int bitShiftDepth, int hash, K key, V alt) {
            return isMatch(hash, key) ? value : alt;
        }

        @Override
        public V getOrElseGet(int bitShiftDepth, int hash, K key, Supplier<? extends V> alt) {
            return isMatch(hash, key) ? value : alt.get();
        }

        private boolean isMatch(int hash, K key) {
            return this.hash == hash && Objects.equals(this.key, key);
        }

        @Override
        public Node<K, V> minus(int bitShiftDepth, int hash, K key) {
            return isMatch(hash, key) ? EmptyNode.Instance : this;
        }

        public int hash() {
            return hash;
        }

        @Override
        public int size() {
            return 1;
        }

        @Override
        public LazySeq<Tuple2<K, V>> lazyList() {
            return LazySeq.of(unapply());
        }

        @Override
        public ReactiveSeq<Tuple2<K, V>> stream() {
            return ReactiveSeq.of(Tuple.tuple(key, value));
        }

        @Override
        public Tuple2<K, V> unapply() {
            return Tuple.tuple(key, value);
        }

        public String toString() {
            return "[h:" + hash + ",k:" + key + ",v:" + value + "]";
        }
    }

    @EqualsAndHashCode
    public static final class CollisionNode<K, V> implements Node<K, V> {

        private final int hash;
        private final int size;
        private final ImmutableList<Tuple2<K, V>> bucket;

        public CollisionNode(int hash, ImmutableList<Tuple2<K, V>> bucket) {
            this.hash = hash;
            this.size = bucket.size();
            this.bucket = bucket;
        }

        private static final long serialVersionUID = 1L;

        @Override
        public Node<K, V> plus(int bitShiftDepth, int hash, K key, V value) {
            ImmutableList<Tuple2<K, V>> filtered = bucket.filter(t -> !Objects.equals(key, t._1()));

            if (this.hash == hash) {
                return filtered.size() == 0 ? new ValueNode<>(hash, key, value) : new CollisionNode<>(hash, filtered.prepend(Tuple.tuple(key, value)));
            }
            return merge(bitShiftDepth, hash, new ValueNode<>(hash, key, value));
        }

        private Node<K, V> merge(int bitShiftDepth, int thatHash, Node<K, V> that) {
            //hash each merge into a collision node if hashes are the same, otherwise store in new location under a BitsetNode
            if (hash == thatHash)
                return new CollisionNode<>(hash, bucket.prependAll(that.lazyList()));
            //create new BitsetNode
            int mask1 = BitsetNode.mask(hash, bitShiftDepth);
            int mask2 = BitsetNode.mask(thatHash, bitShiftDepth);
            int posThis = BitsetNode.bitpos(mask1);
            int posThat = BitsetNode.bitpos(mask2);
            int newBitset = posThis | posThat;
            if (mask1 == mask2) { //collision
                Node<K, V> merged = merge(bitShiftDepth + BITS_IN_INDEX, thatHash, that);
                return new BitsetNode<>(newBitset, 2, new Node[]{merged});
            }
            Node<K, V>[] ordered = mask1 < mask2 ? new Node[]{this, that} : new Node[]{that, this};
            return new BitsetNode<>(newBitset, 2, ordered);
        }

        @Override
        public Option<V> get(int bitShiftDepth, int hash, K key) {
            if (this.hash == hash) {
                return bucket.stream().filter(t -> Objects.equals(key, t._1())).takeOne().map(Tuple2::_2);
            }
            return Option.none();
        }

        @Override
        public V getOrElse(int bitShiftDepth, int hash, K key, V alt) {
            return get(bitShiftDepth, hash, key).orElse(alt);
        }

        @Override
        public V getOrElseGet(int bitShiftDepth, int hash, K key, Supplier<? extends V> alt) {
            return get(bitShiftDepth, hash, key).orElseGet(alt);
        }

        @Override
        public Node<K, V> minus(int bitShiftDepth, int hash, K key) {
            if (this.hash == hash) {
                return new CollisionNode<>(hash, bucket.filter(t -> !Objects.equals(key, t._1())));
            }
            return this;
        }
        @Override
        public ReactiveSeq<Tuple2<K, V>> streamNaturalOrder() {
            return stream().sorted(Comparators.naturalOrderIdentityComparator());
        }

        @Override
        public int size() {
            return bucket.size();
        }

        @Override
        public LazySeq<Tuple2<K, V>> lazyList() {
            return bucket.lazySeq();
        }

        @Override
        public ReactiveSeq<Tuple2<K, V>> stream() {
            return bucket.stream();
        }

        public String toString() {
            return "[COLLISION : h:" + hash + "," + bucket.toString() + "]";
        }
    }

    @AllArgsConstructor
    @EqualsAndHashCode
    public static final class BitsetNode<K, V> implements Node<K, V> {
        public final int bitset;
        private final int size;
        private final Node<K, V>[] nodes;
        private static final long serialVersionUID = 1L;

        @Override
        public Node<K, V> plus(int bitShiftDepth, int hash, K key, V value) {
            int bitPos = bitpos(hash, bitShiftDepth);
            int arrayPos = index(bitPos);
            Node<K, V> node = (absent(bitPos) ? EmptyNode.Instance : nodes[arrayPos]).plus(bitShiftDepth + BITS_IN_INDEX, hash, key, value);
            if (absent(bitPos)) {
                int addedBit = bitset | bitPos;
                Node<K, V>[] addedNodes = new Node[nodes.length + 1];
                System.arraycopy(nodes, 0, addedNodes, 0, arrayPos);
                addedNodes[arrayPos] = node;
                System.arraycopy(nodes, arrayPos, addedNodes, arrayPos + 1, nodes.length - arrayPos);
                return new BitsetNode<>(addedBit, size(addedNodes), addedNodes);
            } else {
                Node<K, V>[] updatedNodes = Arrays.copyOf(nodes, nodes.length);
                updatedNodes[arrayPos] = node;
                return new BitsetNode<>(bitset, size(updatedNodes), updatedNodes);
            }

        }

        static int size(Node[] n) {
            int res = 0;
            for (Node next : n) {
                res = res + next.size();
            }
            return res;
        }

        @Override
        public Option<V> get(int bitShiftDepth, int hash, K key) {
            int pos = bitpos(hash, bitShiftDepth);
            return absent(pos) ? Option.none() : find(bitShiftDepth, pos, hash, key);
        }

        @Override
        public V getOrElse(int bitShiftDepth, int hash, K key, V alt) {
            int pos = bitpos(hash, bitShiftDepth);
            return absent(pos) ? alt : find(bitShiftDepth, pos, hash, key, alt);
        }

        @Override
        public V getOrElseGet(int bitShiftDepth, int hash, K key, Supplier<? extends V> alt) {
            int pos = bitpos(hash, bitShiftDepth);
            return absent(pos) ? alt.get() : findOrGet(bitShiftDepth, pos, hash, key, alt);
        }

        public boolean absent(int pos) {
            return (bitset & pos) == 0;
        }

        private V findOrGet(int shift, int pos, int hash, K key, Supplier<? extends V> alt) {
            return nodes[index(pos)].getOrElseGet(shift + BITS_IN_INDEX, hash, key, alt);
        }

        private V find(int shift, int pos, int hash, K key, V alt) {
            return nodes[index(pos)].getOrElse(shift + BITS_IN_INDEX, hash, key, alt);
        }

        private Option<V> find(int shift, int pos, int hash, K key) {
            return nodes[index(pos)].get(shift + BITS_IN_INDEX, hash, key);
        }

        private Node<K, V> findNode(int pos) {
            return nodes[index(pos)];
        }

        @Override
        public Node<K, V> minus(int bitShiftDepth, int hash, K key) {
            if (nodes.length == 0)
                return this;
            int bitPos = bitpos(hash, bitShiftDepth);
            int arrayPos = index(bitPos);
            boolean absent = absent(bitPos);
            if (absent)
                return this;
            Node<K, V> node = nodes[arrayPos].minus(bitShiftDepth + BITS_IN_INDEX, hash, key);

            if (node instanceof EmptyNode) {
                int removedBit = bitset & ~bitPos;

                Node<K, V>[] removedNodes = new Node[nodes.length - 1];
                System.arraycopy(nodes, 0, removedNodes, 0, arrayPos);
                System.arraycopy(nodes, arrayPos + 1, removedNodes, arrayPos, nodes.length - arrayPos - 1);
                if (removedNodes.length == 1) {
                    return removedNodes[0];
                }
                return new BitsetNode<>(removedBit, size(removedNodes), removedNodes);
            } else {
                int removedBit = bitset & ~bitPos;


                Node<K, V>[] removedNodes = new Node[nodes.length];
                System.arraycopy(nodes, 0, removedNodes, 0, nodes.length);


                removedNodes[arrayPos] = node;
                return new BitsetNode<>(bitset, size(removedNodes), removedNodes);
            }
        }

        @Override
        public int size() {
            return size;
        }

        @Override
        public LazySeq<Tuple2<K, V>> lazyList() {
            return LazySeq.fromStream(stream());
        }

        @Override
        public ReactiveSeq<Tuple2<K, V>> stream() {
            return ReactiveSeq.of(nodes).flatMap(n -> n.stream());
        }

        static int bitpos(int hash, int shift) {
            return 1 << mask(hash, shift);
        }

        static int bitpos(int mask) {
            return 1 << mask;
        }

        static int mask(int hash, int shift) {
            return (hash >>> shift) & (SIZE - 1);
        }

        int index(int bit) {
            return Integer.bitCount(bitset & (bit - 1));
        }

        @Override
        public String toString() {
            StringBuilder s = new StringBuilder("{b:" + Integer.toBinaryString(bitset) + ",s:" + size);
            for (Node<K, V> next : nodes) {
                s.append("," + next.toString());
            }
            return s.append("}").toString();

        }
    }

}
