/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package org.apache.cassandra.utils.btree;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collection;
import java.util.Comparator;
import java.util.List;

import javax.annotation.Nullable;

import accord.utils.AsymmetricComparator;
import accord.utils.Invariants;
import accord.utils.QuadFunction;
import accord.utils.SortedArrays;

import org.apache.cassandra.utils.caching.TinyThreadLocalPool;

import io.netty.util.concurrent.FastThreadLocal;

import static accord.utils.SortedArrays.Search.CEIL;
import static accord.utils.SortedArrays.Search.FLOOR;
import static org.apache.cassandra.utils.btree.BTree.BRANCH_FACTOR;
import static org.apache.cassandra.utils.btree.BTree.getBranchKeyEnd;
import static org.apache.cassandra.utils.btree.BTree.getChildCount;
import static org.apache.cassandra.utils.btree.BTree.getChildStart;
import static org.apache.cassandra.utils.btree.BTree.getKeyEnd;
import static org.apache.cassandra.utils.btree.BTree.getLeafKeyEnd;
import static org.apache.cassandra.utils.btree.BTree.isEmpty;
import static org.apache.cassandra.utils.btree.BTree.isLeaf;
import static org.apache.cassandra.utils.btree.BTree.size;
import static org.apache.cassandra.utils.btree.BTree.sizeMap;

/**
 * A very simple extension to BTree to provide an Augmented Interval BTree.
 *
 * TODO (desired): there are a number of obvious performance improvements that should be pursued at the earliest suitable opportunity:
 *      - IntervalMaxIndex should include an index over any direct child leaves to bound worst case to O(lgN + m) rather than O(lgN + m.B)
 *      - IntervalMaxIndex should include the immediately preceding self-key in any max key, so we can avoid the O(N) loop and simply pre-visit the self-key
 *      - IntervalMaxIndex should omit entirely keys/nodes that do not extend past the following key
 *      - Updater/Subtraction should re-use parts of IntervalMaxIndex where appropriate
 */
public class IntervalBTree
{
    /**
     * Defines additional SORT comparators that must be symmetric for building/maintaining an interval tree
     */
    public interface IntervalComparators<I> extends WithIntervalComparators<I, I>
    {
        Comparator<I> totalOrder();
        Comparator<I> endWithEndSorter();
    }

    /**
     * Defines search comparators for finding an element in an IntervalBTree.
     *
     * A search for an element K will potentially include any interval I
     * for which any comparator yields compare(K, I) == 0.
     *
     * These comparators DO NOT need to be symmetric in either type OR comparison result,
     * and should knowingly impose the inclusive/exclusive criteria of the interval
     * by returning non-zero answers as appropriate.
     *
     * Specifically, compare(a, a) should not return 0 if the relevant bound is exclusive.
     */
    public interface WithIntervalComparators<K, I>
    {
        AsymmetricComparator<K, I> startWithStartSeeker();
        AsymmetricComparator<K, I> startWithEndSeeker();
        AsymmetricComparator<K, I> endWithStartSeeker();
    }

    public static class InclusiveEndHelper
    {
        public static int keyStartWithStart(int c) { return equalsMeansAfter(c); }
        public static int keyStartWithEnd(int c) { return c; }
        public static int keyEndWithStart(int c) { return equalsMeansBefore(c); }
        public static int startWithStart(int c) { return equalsMeansAfter(c); }
        public static int startWithEnd(int c) { return equalsMeansAfter(c); }
        public static int endWithStart(int c) { return equalsMeansBefore(c); }
        private static int equalsMeansAfter(int c) { return c == 0 ? 1 : c; }
        private static int equalsMeansBefore(int c) { return c == 0 ? -1 : c; }
    }

    /**
     * Apply the accumulation function over all intersecting intervals in the tree
     */
    public static <Find, Interval, P1, P2, V> V accumulate(Object[] btree, WithIntervalComparators<Find, Interval> comparators, Find find, QuadFunction<P1, P2, Interval, V, V> function, P1 p1, P2 p2, V accumulate)
    {
        if (isLeaf(btree))
        {
            AsymmetricComparator<Find, Interval> startWithEnd = comparators.startWithEndSeeker();
            AsymmetricComparator<Find, Interval> endWithStart = comparators.endWithStartSeeker();
            int keyEnd = getLeafKeyEnd(btree);
            for (int i = 0; i < keyEnd; ++i)
            {
                Interval interval = (Interval) btree[i];
                if (endWithStart.compare(find, interval) >= 0 && startWithEnd.compare(find, interval) <= 0)
                    accumulate = function.apply(p1, p2, interval, accumulate);
            }
        }
        else
        {
            // start/end represent those keys/children we are guaranteed to intersect (and so may visit without any further comparisons)
            int startKey = SortedArrays.binarySearch(btree, 0, getKeyEnd(btree), find, (AsymmetricComparator)comparators.startWithStartSeeker(), CEIL);
            int startChild;

            if (startKey >= 0) startChild = startKey + 1;
            else startChild = (startKey = -1 - startKey);
            int endKey = SortedArrays.binarySearch(btree, startKey, getKeyEnd(btree), find, (AsymmetricComparator)comparators.endWithStartSeeker(), FLOOR);
            int endChild;
            if (endKey >= 0) { if (endKey < getKeyEnd(btree)) ++endKey; endChild = 1 + endKey; }
            else endChild = 1 + (endKey = -1 - endKey);

            {   // descend anything with a max that overlaps us that we won't already visit
                if (startChild > 0 || startKey > 0)
                    accumulate = accumulateMaxOnly(startChild + (startChild == endChild ? 1 : 0), startKey, btree, comparators, find, function, p1, p2, accumulate);
            }

            int childOffset = getChildStart(btree);
            if (startChild == startKey && startChild < endChild)
                accumulate = accumulate((Object[]) btree[childOffset + startChild++], comparators, find, function, p1, p2, accumulate);

            if (startKey < startChild && startKey < endKey)
                accumulate = function.apply(p1, p2, (Interval) btree[startKey++], accumulate);

            while (startChild < endChild - 1)
            {
                Invariants.require(startKey == startChild);
                accumulate = accumulate((Object[]) btree[childOffset + startChild++], function, p1, p2, accumulate);
                accumulate = function.apply(p1, p2, (Interval) btree[startKey++], accumulate);
            }
            if (startKey < endKey)
                accumulate = function.apply(p1, p2, (Interval) btree[startKey], accumulate);
            if (startChild < endChild)
                accumulate = accumulate((Object[]) btree[childOffset + startChild], comparators, find, function, p1, p2, accumulate);
        }
        return accumulate;
    }

    public static <R, P1, P2, V> V accumulate(Object[] btree, QuadFunction<P1, P2, R, V, V> function, P1 p1, P2 p2, V accumulate)
    {
        if (isLeaf(btree))
        {
            for (int i = 0, maxi = getLeafKeyEnd(btree); i < maxi; ++i)
                accumulate = function.apply(p1, p2, (R) btree[i], accumulate);
        }
        else
        {
            int keyEnd = getBranchKeyEnd(btree);
            for (int i = 0; i < keyEnd; ++i)
            {
                accumulate = accumulate((Object[]) btree[keyEnd + i], function, p1, p2, accumulate);
                accumulate = function.apply(p1, p2, (R) btree[i], accumulate);
            }
            accumulate = accumulate((Object[]) btree[2 * keyEnd], function, p1, p2, accumulate);
        }
        return accumulate;
    }

    private static <Find, Interval, P1, P2, V> V accumulateMaxOnly(int ifChildBefore, int ifKeyBefore, Object[] btree, WithIntervalComparators<Find, Interval> comparators, Find find, QuadFunction<P1, P2, Interval, V, V> function, P1 p1, P2 p2, V accumulate)
    {
        AsymmetricComparator<Find, Interval> startWithEnd = comparators.startWithEndSeeker();
        AsymmetricComparator<Find, Interval> endWithStart = comparators.endWithStartSeeker();
        if (isLeaf(btree))
        {
            Invariants.require(ifChildBefore == Integer.MAX_VALUE);
            for (int i = 0, maxi = getLeafKeyEnd(btree); i < maxi; ++i)
            {
                Interval interval = (Interval) btree[i];
                if (startWithEnd.compare(find, interval) <= 0 && endWithStart.compare(find, interval) >= 0)
                    accumulate = function.apply(p1, p2, interval, accumulate);
            }
        }
        else
        {
            int keyEnd = getChildStart(btree);
            ifKeyBefore = Math.min(ifKeyBefore, keyEnd);

            IntervalMaxIndex intervalMaxIndex = getIntervalMaxIndex(btree);
            int descendMaxStart = SortedArrays.binarySearch((Interval[]) intervalMaxIndex.sortedByEnd, 0, intervalMaxIndex.sortedByEnd.length, find, comparators.startWithEndSeeker(), CEIL);
            if (descendMaxStart < 0)
                descendMaxStart = -1 - descendMaxStart;
            int[] sortedByEndIndex = intervalMaxIndex.sortedByEndIndex;
            for (int i = descendMaxStart; i < sortedByEndIndex.length; ++i)
            {
                int index = sortedByEndIndex[i];
                if (index < ifKeyBefore)
                {
                    Interval interval = (Interval) btree[index];
                    if (startWithEnd.compare(find, interval) <= 0 && endWithStart.compare(find, interval) >= 0)
                        accumulate = function.apply(p1, p2, interval, accumulate);
                }
                if (index < ifChildBefore)
                    accumulate = accumulateMaxOnly(Integer.MAX_VALUE, Integer.MAX_VALUE, (Object[]) btree[keyEnd + index], comparators, find, function, p1, p2, accumulate);
            }
        }
        return accumulate;
    }

    static class IntervalIndexAdapter implements Comparator<IntervalIndexAdapter.SortEntry>
    {
        static class SortEntry
        {
            Object sort;
            int index;
        }

        Object[] unpackedIndex;
        SortEntry[] sort = new SortEntry[BRANCH_FACTOR];
        Comparator endSorter, totalOrder;

        void setComparators(IntervalComparators comparators)
        {
            this.endSorter = comparators.endWithEndSorter();
            this.totalOrder = comparators.totalOrder();
        }

        void reset()
        {
            this.endSorter = this.totalOrder = null;
            for (int i = 0 ; i < sort.length && sort[i] != null ; i++)
                sort[i].sort = null;
        }

        public void override(Object[] branch, @Nullable List<Object[]> sourceNodes)
        {
            int childStart = getChildStart(branch), childCount = getChildCount(branch);
            for (int i = childCount - 1 ; i >= 0 && sort[i] == null ; --i)
                sort[i] = new SortEntry();

            if (sourceNodes != null && sourceNodes.isEmpty())
                sourceNodes = null;

            if (sourceNodes != null)
            {
                int maxUnpackedSize = 0;
                if (unpackedIndex == null)
                    unpackedIndex = new Object[BRANCH_FACTOR];

                int i = 0;
                for (Object[] sn : sourceNodes)
                {
                    IntervalMaxIndex index = getIntervalMaxIndex(sn);
                    int snStart = getChildStart(sn), snCount = getChildCount(sn);
                    for (int j = 0 ; j < index.sortedByEnd.length ; ++j)
                        unpackedIndex[index.sortedByEndIndex[j]] = index.sortedByEnd[j];

                    int j = 0;
                    while (i < childCount && j < snCount)
                    {
                        if (branch[childStart + i] == sn[snStart + j]) // child nodes are the same so we can copy max
                        {
                            sort[i].index = i;
                            if (j == snStart)
                            {
                                // in either case, the max we have represents the max of the child;
                                // must only compare with new branch key (if any) and take the branch if ends later
                                sort[i].sort = i == childStart || endSorter.compare(branch[i], unpackedIndex[j]) <= 0
                                               ? unpackedIndex[j] : branch[i];
                            }
                            else if (unpackedIndex[j] != sn[j])
                            {
                                // previous max was not the branch, so must have been within child
                                // like prior condition, must only compare with new branch key (if any) and take the branch if ends later
                                // however, we have a chance to save a comparison here if we also copied the key from the source node
                                sort[i].sort = i == childStart || branch[i] == sn[j] || endSorter.compare(branch[i], unpackedIndex[j]) <= 0
                                               ? unpackedIndex[j] : branch[i];
                            }
                            else if (i != childStart && endSorter.compare(sn[j], branch[i]) <= 0)
                            {
                                // previous max was previous branch key; new branch key has higher end so can simply use it
                                sort[i].sort = branch[i];
                            }
                            else
                            {
                                // otherwise we cannot infer anything directly
                                sort[i].sort = null;
                            }

                            ++i;
                            ++j;
                        }
                        else
                        {
                            if (i == childStart || j == snStart)
                                break; // if we've finished iteration of one or the other, break and continue with any further source nodes (if any)

                            int c = totalOrder.compare(branch[i], sn[j]);
                            if (c <= 0) sort[i++].sort = null;
                            if (c >= 0) ++j;
                        }
                    }

                    maxUnpackedSize = Math.max(maxUnpackedSize, snCount);
                }
                while (i < childCount)
                    sort[i++].sort = null;
                Arrays.fill(unpackedIndex, 0, maxUnpackedSize, null);
            }
            for (int i = 0; i < childCount; ++i)
            {
                if (sourceNodes != null && sort[i].sort != null)
                    continue;

                Object[] child = (Object[]) branch[i + childStart];
                Object max = maxChild(child);
                if (i < childStart && endSorter.compare(branch[i], max) > 0)
                    max = branch[i];
                sort[i].index = i;
                sort[i].sort = max;
            }
            Arrays.sort(sort, 0, childCount, this);
            Object[] sortedByEnd = new Object[childCount];
            int[] sortedByEndIndex = new int[childCount];
            for (int i = 0; i < childCount; ++i)
            {
                sortedByEnd[i] = sort[i].sort;
                sortedByEndIndex[i] = sort[i].index;
            }

            int[] sizeMap = sizeMap(branch);
            branch[branch.length - 1] = new IntervalMaxIndex(sortedByEnd, sortedByEndIndex, sizeMap);
        }

        private Object maxChild(Object[] child)
        {
            if (isLeaf(child))
            {
                Object max = child[0];
                for (int i = 1, end = getLeafKeyEnd(child); i < end; ++i)
                {
                    if (endSorter.compare(child[i], max) > 0)
                        max = child[i];
                }
                return max;
            }
            else
            {
                IntervalMaxIndex index = getIntervalMaxIndex(child);
                return index.sortedByEnd[index.sortedByEnd.length - 1];
            }
        }

        @Override
        public int compare(SortEntry o1, SortEntry o2)
        {
            return endSorter.compare(o1.sort, o2.sort);
        }
    }

    static class IntervalBranchBuilder extends BTree.BranchBuilder
    {
        List<Object[]> sourceNodes;
        IntervalIndexAdapter adapter;

        IntervalBranchBuilder(BTree.LeafOrBranchBuilder child)
        {
            super(child);
            sourceNodes = child.height == 1 ? new ArrayList<>() : null;
        }

        @Override
        BTree.BranchBuilder allocateParent()
        {
            return new IntervalBranchBuilder(this);
        }

        @Override
        void initParent()
        {
            super.initParent();
            ((IntervalBranchBuilder) parent).init(adapter);
        }

        private IntervalBranchBuilder init(IntervalIndexAdapter adapter)
        {
            this.adapter = adapter;
            return this;
        }

        @Override
        int setDrainSizeMap(Object[] original, int keysInOriginal, Object[] branch, int keysInBranch)
        {
            int result = super.setDrainSizeMap(original, keysInOriginal, branch, keysInBranch);
            adapter.override(branch, sourceNodes);
            return result;
        }

        @Override
        void setRedistributedSizeMap(Object[] branch, int steal)
        {
            super.setRedistributedSizeMap(branch, steal);
            adapter.override(branch, sourceNodes);
        }

        @Override
        int setOverflowSizeMap(Object[] branch, int keys)
        {
            int result = super.setOverflowSizeMap(branch, keys);
            adapter.override(branch, sourceNodes);
            return result;
        }

        @Override
        void setSourceNode(Object[] sourceNode)
        {
            if (sourceNodes != null)
                sourceNodes.add(sourceNode);
            super.setSourceNode(sourceNode);
        }

        @Override
        void clearSourceNode()
        {
            super.clearSourceNode();
            if (sourceNodes != null)
                sourceNodes.clear();
        }
    }

    public static class FastIntervalTreeBuilder<V> extends BTree.FastBuilder<V>
    {
        private static final TinyThreadLocalPool<FastIntervalTreeBuilder<?>> POOL = new TinyThreadLocalPool<>();
        final IntervalIndexAdapter adapter = new IntervalIndexAdapter();

        @Override
        BTree.BranchBuilder allocateParent()
        {
            return new IntervalBranchBuilder(this);
        }

        @Override
        void initParent()
        {
            super.initParent();
            ((IntervalBranchBuilder) parent).init(adapter);
        }

        @Override
        public void close()
        {
            super.close();
            adapter.reset();
        }
    }

    static class IntervalUpdater<Compare, Existing extends Compare, Insert extends Compare> extends BTree.Updater<Compare, Existing, Insert>
    {
        static final TinyThreadLocalPool<IntervalUpdater> POOL = new TinyThreadLocalPool<>();
        private final IntervalIndexAdapter adapter = new IntervalIndexAdapter();

        static <Compare, Existing extends Compare, Insert extends Compare> IntervalUpdater<Compare, Existing, Insert> get(IntervalComparators<Compare> comparators)
        {
            TinyThreadLocalPool.TinyPool<IntervalUpdater> pool = POOL.get();
            IntervalUpdater<Compare, Existing, Insert> updater = pool.poll();
            if (updater == null)
                updater = new IntervalUpdater<>();
            updater.pool = pool;
            updater.adapter.setComparators(comparators);
            return updater;
        }

        @Override
        BTree.BranchBuilder allocateParent()
        {
            return new IntervalBranchBuilder(this);
        }

        @Override
        void initParent()
        {
            super.initParent();
            ((IntervalBranchBuilder) parent).init(adapter);
        }

        @Override
        public void close()
        {
            // TODO (required): validate this in IntervalBTreeTest
            super.close();
            adapter.reset();
        }
    }

    public static Object[] empty()
    {
        return BTree.empty();
    }

    public static Object[] singleton(Object value)
    {
        return BTree.singleton(value);
    }

    static class Subtraction<K, T extends K> extends BTree.AbstractSubtraction<K, T>
    {
        static final FastThreadLocal<Subtraction> SHARED = new FastThreadLocal<>();
        private final IntervalIndexAdapter adapter = new IntervalIndexAdapter();

        static <K, T extends K> Subtraction<K, T> get(IntervalComparators<K> comparators)
        {
            Subtraction subtraction = SHARED.get();
            if (subtraction == null)
                SHARED.set(subtraction = new Subtraction());
            subtraction.comparator = comparators.totalOrder();
            subtraction.adapter.setComparators(comparators);
            return subtraction;
        }

        @Override
        BTree.BranchBuilder allocateParent()
        {
            return new IntervalBranchBuilder(this);
        }

        @Override
        void initParent()
        {
            super.initParent();
            ((IntervalBranchBuilder) parent).init(adapter);
        }

        @Override
        void reset()
        {
            super.reset();
            adapter.reset();
        }
    }

    /**
     * Subtracts {@code insert} into {@code update}, applying {@code updateF} to each new item in {@code insert},
     * as well as any matched items in {@code update}.
     * <p>
     * Note that {@code UpdateFunction.noOp} is assumed to indicate a lack of interest in which value survives.
     */
    public static <Compare> Object[] subtract(Object[] toUpdate, Object[] subtract, IntervalComparators<Compare> comparators)
    {
        try (Subtraction subtraction = Subtraction.get(comparators))
        {
            return subtraction.subtract(toUpdate, subtract);
        }
    }

    /**
     * Inserts {@code insert} into {@code update}, applying {@code updateF} to each new item in {@code insert},
     * as well as any matched items in {@code update}.
     * <p>
     * Note that {@code UpdateFunction.noOp} is assumed to indicate a lack of interest in which value survives.
     */
    public static <Compare, Existing extends Compare, Insert extends Compare> Object[] update(Object[] existing,
                                                                                              Object[] insert,
                                                                                              IntervalComparators<Compare> comparators)
    {
        // perform some initial obvious optimisations
        if (isEmpty(insert))
            return existing; // do nothing if update is empty

        if (isEmpty(existing))
            return insert;

        if (isLeaf(insert))
        {
            // consider flipping the order of application, if update is much larger than insert and applying unary no-op
            int updateSize = size(existing);
            int insertSize = size(insert);
            int scale = Integer.numberOfLeadingZeros(updateSize) - Integer.numberOfLeadingZeros(insertSize);
            if (scale >= 4)
            {
                // i.e. at roughly 16x the size, or one tier deeper - very arbitrary, should pick more carefully
                // experimentally, at least at 64x the size the difference in performance is ~10x
                Object[] tmp = insert;
                insert = existing;
                existing = tmp;
            }
        }

        try (IntervalUpdater<Compare, Existing, Insert> updater = IntervalUpdater.get(comparators))
        {
            return updater.update(existing, insert, comparators.totalOrder(), (UpdateFunction) UpdateFunction.noOp);
        }
    }

    public static <V> Object[] build(Collection<V> build, IntervalComparators<V> comparators)
    {
        try (FastIntervalTreeBuilder<V> builder = IntervalBTree.fastBuilder(comparators))
        {
            for (V v : build)
                builder.add(v);
            return builder.build();
        }
    }

    /**
     * Build a tree of unknown size, in order.
     */
    public static <V> FastIntervalTreeBuilder<V> fastBuilder(IntervalComparators<V> comparators)
    {
        TinyThreadLocalPool.TinyPool<FastIntervalTreeBuilder<?>> pool = FastIntervalTreeBuilder.POOL.get();
        FastIntervalTreeBuilder<V> builder = (FastIntervalTreeBuilder<V>) pool.poll();
        if (builder == null)
            builder = new FastIntervalTreeBuilder<>();
        builder.pool = pool;
        builder.adapter.setComparators(comparators);
        return builder;
    }

    // TODO (desired): index leaf nodes
    static class IntervalMaxIndex
    {
        final Object[] sortedByEnd;
        final int[] sortedByEndIndex;
        final int[] sizeMap;

        IntervalMaxIndex(Object[] sortedByEnd, int[] sortedByEndIndex, int[] sizeMap)
        {
            this.sortedByEnd = sortedByEnd;
            this.sortedByEndIndex = sortedByEndIndex;
            this.sizeMap = sizeMap;
        }
    }

    /**
     * @return the size map for the branch node
     */
    static IntervalMaxIndex getIntervalMaxIndex(Object[] branchNode)
    {
        return (IntervalMaxIndex) branchNode[branchNode.length - 1];
    }
}
