/*
 * 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 accord.local;

import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.function.Function;

import javax.annotation.Nullable;

import accord.primitives.Range;
import accord.primitives.Ranges;
import accord.utils.Invariants;

public interface ShardDistributor
{
    // TODO (expected, topology): this is overly simplistic: need to supply existing distribution, and support
    //                            gradual local redistribution to keep number of shards eventually the same
    List<Ranges> split(Ranges ranges);

    /**
     * Return a single subSpit from a range given the total number of splits the range should be divided into and
     * the index of the split to be returned.
     *
     * If the range is not sufficiently divisible, some time slots will return null, and the others unitary portions of the range
     */
    @Nullable Range splitRange(Range range, int from, int to, int totalSplits);

    Ranges selectFirstSubRanges(Range range, Ranges subRanges, int totalSplits);

    int numberOfSplitsPossible(Range range);

    class EvenSplit<T> implements ShardDistributor
    {
        public interface Splitter<T>
        {
            default boolean splittable(Range range, int numSplits) { return true; }
            T sizeOf(Range range);
            Range subRange(Range range, T start, T end);

            T zero();
            T valueOf(int v);
            T add(T a, T b);
            T subtract(T a, T b);
            T divide(T a, int i);
            T divide(T a, T b);
            T multiply(T a, int i);
            int min(T v, int i);
            int compare(T a, T b);
        }

        final int numberOfShards;
        final Function<Ranges, ? extends Splitter<T>> splitter;

        public EvenSplit(int numberOfShards, Function<Ranges, ? extends Splitter<T>> splitter)
        {
            this.numberOfShards = numberOfShards;
            this.splitter = splitter;
        }

        @Override
        public int numberOfSplitsPossible(Range range)
        {
            Splitter<T> splitter = this.splitter.apply(Ranges.of(range));
            return splitter.min(splitter.sizeOf(range), Integer.MAX_VALUE);
        }

        @Override
        public Range splitRange(Range range, int from, int to, int numSplits)
        {
            Invariants.requireArgument(from <= to);
            Invariants.requireArgument(to <= numSplits);
            Splitter<T> splitter = this.splitter.apply(Ranges.single(range));
            if (!splitter.splittable(range, numSplits))
                return range;

            T size = splitter.sizeOf(range);
            T splitSize = splitter.divide(size, numSplits);
            T remainder = splitter.subtract(size, splitter.multiply(splitSize, numSplits));
            T splitPlusOneRate = remainder.equals(splitter.zero()) ? null : splitter.divide(splitter.add(remainder, splitter.valueOf(numSplits-1)), remainder);
            T splitBegin = splitter.multiply(splitSize, from);
            if (splitPlusOneRate != null)
                splitBegin = splitter.add(splitBegin, splitter.divide(splitter.valueOf(from), splitPlusOneRate));

            T splitEnd;
            if (to == numSplits)
            {
                splitEnd = size;
            }
            else
            {
                splitEnd = splitter.multiply(splitSize, to);
                if (splitPlusOneRate != null)
                    splitEnd = splitter.add(splitEnd, splitter.divide(splitter.valueOf(to), splitPlusOneRate));
            }

            if (splitBegin.equals(splitEnd))
                return null;

            return splitter.subRange(range, splitBegin, splitEnd);
        }

        @Override
        public Ranges selectFirstSubRanges(Range range, Ranges subRanges, int numSplits)
        {
            Splitter<T> splitter = this.splitter.apply(Ranges.of(range));
            T totalSize = splitter.sizeOf(range);
            T splitSize = splitter.divide(totalSize, numSplits);

            T runningTotal = splitter.zero();
            int i = 0;
            while (i < subRanges.size())
            {
                Range next = subRanges.get(i);
                T sizeOfNext = splitter.sizeOf(next);
                T totalWithNext = splitter.add(sizeOfNext, runningTotal);
                if (splitter.compare(totalWithNext, splitSize) >= 0)
                    break;
                runningTotal = totalWithNext;
                ++i;
            }

            if (i == subRanges.size())
                return subRanges;

            Range next = subRanges.get(i);
            next = splitter.subRange(next, splitter.zero(), splitter.subtract(splitSize, runningTotal));
            if (i == 0)
                return Ranges.of(next);

            Ranges ranges = subRanges.slice(0, i);
            return ranges.with(Ranges.of(next));
        }

        @Override
        public List<Ranges> split(Ranges ranges)
        {
            if (ranges.isEmpty())
                return Collections.emptyList();

            Splitter<T> splitter = this.splitter.apply(ranges);
            T totalSize = splitter.zero();
            for (int i = 0 ; i < ranges.size() ; ++i)
                totalSize = splitter.add(totalSize, splitter.sizeOf(ranges.get(i)));

            if (splitter.compare(totalSize, splitter.zero()) <= 0)
                throw new IllegalStateException();

            int numberOfSplits = splitter.min(totalSize, this.numberOfShards);
            T shardSize = splitter.divide(totalSize, numberOfSplits);
            List<Range> buffer = new ArrayList<>(ranges.size());
            List<Ranges> result = new ArrayList<>(numberOfSplits);
            int ri = 0;
            T rOffset = splitter.zero();
            T rSize = splitter.sizeOf(ranges.get(0));
            while (result.size() < numberOfSplits)
            {
                T required = result.size() < numberOfSplits - 1 ? shardSize : splitter.subtract(totalSize, splitter.multiply(shardSize, (numberOfSplits - 1)));
                while (true)
                {
                    if (splitter.compare(splitter.subtract(rSize, rOffset), required) >= 0)
                    {
                        buffer.add(splitter.subRange(ranges.get(ri), rOffset, splitter.add(rOffset, required)));
                        result.add(Ranges.ofSortedAndDeoverlapped(buffer.toArray(new Range[0])));
                        buffer.clear();
                        rOffset = splitter.add(rOffset, required);
                        if (splitter.compare(rOffset, rSize) >= 0 && ++ri < ranges.size())
                        {
                            Invariants.require(splitter.compare(rOffset, rSize) == 0);
                            rOffset = splitter.zero();
                            rSize = splitter.sizeOf(ranges.get(ri));
                        }
                        break;
                    }
                    else
                    {
                        buffer.add(splitter.subRange(ranges.get(ri), rOffset, rSize));
                        required = splitter.subtract(required, splitter.subtract(rSize, rOffset));
                        rOffset = splitter.zero();
                        rSize = splitter.sizeOf(ranges.get(++ri));
                    }
                }
            }
            return result;
        }
    }
}
