/*
 * 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.dht.tokenallocator;

import java.util.*;

import com.google.common.collect.HashMultimap;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.Maps;
import com.google.common.collect.Multimap;

import org.apache.cassandra.dht.IPartitioner;
import org.apache.cassandra.dht.Token;

/**
 * A Replication Aware allocator for tokens, that attempts to ensure an even distribution of ownership across
 * the known cluster for the provided replication strategy.
 *
 * A unit is shorthand for a "unit of ownership" which translates roughly to a node, or a disk on the node,
 * a CPU on the node, or some other relevant unit of ownership. These units should be the lowest rung over which
 * ownership needs to be evenly distributed. At the moment only nodes as a whole are treated as units, but that
 * will change with the introduction of token ranges per disk.
 */
class ReplicationAwareTokenAllocator<Unit> extends TokenAllocatorBase<Unit>
{
    final Multimap<Unit, Token> unitToTokens;
    final int replicas;

    ReplicationAwareTokenAllocator(NavigableMap<Token, Unit> sortedTokens, ReplicationStrategy<Unit> strategy, IPartitioner partitioner)
    {
        super(sortedTokens, strategy, partitioner);
        unitToTokens = HashMultimap.create();
        for (Map.Entry<Token, Unit> en : sortedTokens.entrySet())
            unitToTokens.put(en.getValue(), en.getKey());
        this.replicas = strategy.replicas();
    }

    public int getReplicas()
    {
        return replicas;
    }

    public Collection<Token> addUnit(Unit newUnit, int numTokens)
    {
        assert !unitToTokens.containsKey(newUnit);

        if (unitCount() < replicas)
            // Allocation does not matter; everything replicates everywhere.
            return generateRandomTokens(newUnit, numTokens);
        if (numTokens > sortedTokens.size())
            // Some of the heuristics below can't deal with this case. Use random for now, later allocations can fix any problems this may cause.
            return generateRandomTokens(newUnit, numTokens);

        // ============= construct our initial token ring state =============

        double optTokenOwnership = optimalTokenOwnership(numTokens);
        Map<Object, GroupInfo> groups = Maps.newHashMap();
        Map<Unit, UnitInfo<Unit>> unitInfos = createUnitInfos(groups);
        if (groups.size() < replicas)
        {
            // We need at least replicas groups to do allocation correctly. If there aren't enough, 
            // use random allocation.
            // This part of the code should only be reached via the RATATest. StrategyAdapter should disallow
            // token allocation in this case as the algorithm is not able to cover the behavior of NetworkTopologyStrategy.
            return generateRandomTokens(newUnit, numTokens);
        }

        // initialise our new unit's state (with an idealised ownership)
        // strategy must already know about this unit
        UnitInfo<Unit> newUnitInfo = new UnitInfo<>(newUnit, numTokens * optTokenOwnership, groups, strategy);

        // build the current token ring state
        TokenInfo<Unit> tokens = createTokenInfos(unitInfos, newUnitInfo.group);
        newUnitInfo.tokenCount = numTokens;

        // ============= construct and rank our candidate token allocations =============

        // walk the token ring, constructing the set of candidates in ring order
        // as the midpoints between all existing tokens
        CandidateInfo<Unit> candidates = createCandidates(tokens, newUnitInfo, optTokenOwnership);

        // Evaluate the expected improvements from all candidates and form a priority queue.
        PriorityQueue<Weighted<CandidateInfo<Unit>>> improvements = new PriorityQueue<>(sortedTokens.size());
        CandidateInfo<Unit> candidate = candidates;
        do
        {
            double impr = evaluateImprovement(candidate, optTokenOwnership, 1.0 / numTokens);
            improvements.add(new Weighted<>(impr, candidate));
            candidate = candidate.next;
        } while (candidate != candidates);

        // ============= iteratively take the best candidate, and re-rank =============

        CandidateInfo<Unit> bestToken = improvements.remove().value;
        for (int vn = 1; ; ++vn)
        {
            candidates = bestToken.removeFrom(candidates);
            confirmCandidate(bestToken);

            if (vn == numTokens)
                break;

            while (true)
            {
                // Get the next candidate in the queue. Its improvement may have changed (esp. if multiple tokens
                // were good suggestions because they could improve the same problem)-- evaluate it again to check
                // if it is still a good candidate.
                bestToken = improvements.remove().value;
                double impr = evaluateImprovement(bestToken, optTokenOwnership, (vn + 1.0) / numTokens);
                Weighted<CandidateInfo<Unit>> next = improvements.peek();

                // If it is better than the next in the queue, it is good enough. This is a heuristic that doesn't
                // get the best results, but works well enough and on average cuts search time by a factor of O(vnodes).
                if (next == null || impr >= next.weight)
                    break;
                improvements.add(new Weighted<>(impr, bestToken));
            }
        }

        return ImmutableList.copyOf(unitToTokens.get(newUnit));
    }

    private Collection<Token> generateRandomTokens(Unit newUnit, int numTokens)
    {
        Set<Token> tokens = new HashSet<>(numTokens);
        while (tokens.size() < numTokens)
        {
            Token token = partitioner.getRandomToken();
            if (!sortedTokens.containsKey(token))
            {
                tokens.add(token);
                sortedTokens.put(token, newUnit);
                unitToTokens.put(newUnit, token);
            }
        }
        return tokens;
    }

    /**
     * Construct the token ring as a CircularList of TokenInfo,
     * and populate the ownership of the UnitInfo's provided
     */
    private TokenInfo<Unit> createTokenInfos(Map<Unit, UnitInfo<Unit>> units, GroupInfo newUnitGroup)
    {
        // build the circular list
        TokenInfo<Unit> prev = null;
        TokenInfo<Unit> first = null;
        for (Map.Entry<Token, Unit> en : sortedTokens.entrySet())
        {
            Token t = en.getKey();
            UnitInfo<Unit> ni = units.get(en.getValue());
            TokenInfo<Unit> ti = new TokenInfo<>(t, ni);
            first = ti.insertAfter(first, prev);
            prev = ti;
        }

        TokenInfo<Unit> curr = first;
        do
        {
            populateTokenInfoAndAdjustUnit(curr, newUnitGroup);
            curr = curr.next;
        } while (curr != first);

        return first;
    }

    private CandidateInfo<Unit> createCandidates(TokenInfo<Unit> tokens, UnitInfo<Unit> newUnitInfo, double initialTokenOwnership)
    {
        TokenInfo<Unit> curr = tokens;
        CandidateInfo<Unit> first = null;
        CandidateInfo<Unit> prev = null;
        do
        {
            CandidateInfo<Unit> candidate = new CandidateInfo<Unit>(partitioner.midpoint(curr.prev.token, curr.token), curr, newUnitInfo);
            first = candidate.insertAfter(first, prev);

            candidate.replicatedOwnership = initialTokenOwnership;
            populateCandidate(candidate);

            prev = candidate;
            curr = curr.next;
        } while (curr != tokens);
        prev.next = first;
        return first;
    }

    private void populateCandidate(CandidateInfo<Unit> candidate)
    {
        // Only finding replication start would do.
        populateTokenInfo(candidate, candidate.owningUnit.group);
    }

    /**
     * Incorporates the selected candidate into the ring, adjusting ownership information and calculated token
     * information.
     */
    private void confirmCandidate(CandidateInfo<Unit> candidate)
    {
        // This process is less efficient than it could be (loops through each vnode's replication span instead
        // of recalculating replicationStart, replicationThreshold from existing data + new token data in an O(1)
        // case analysis similar to evaluateImprovement). This is fine as the method does not dominate processing
        // time.

        // Put the accepted candidate in the token list.
        UnitInfo<Unit> newUnit = candidate.owningUnit;
        Token newToken = candidate.token;
        sortedTokens.put(newToken, newUnit.unit);
        unitToTokens.put(newUnit.unit, newToken);

        TokenInfo<Unit> prev = candidate.prevInRing();
        TokenInfo<Unit> newTokenInfo = new TokenInfo<>(newToken, newUnit);
        newTokenInfo.replicatedOwnership = candidate.replicatedOwnership;
        newTokenInfo.insertAfter(prev, prev);   // List is not empty so this won't need to change head of list.

        // Update data for candidate.
        populateTokenInfoAndAdjustUnit(newTokenInfo, newUnit.group);

        ReplicationVisitor replicationVisitor = new ReplicationVisitor();
        assert newTokenInfo.next == candidate.split;
        for (TokenInfo<Unit> curr = newTokenInfo.next; !replicationVisitor.visitedAll(); curr = curr.next)
        {
            // update the candidate between curr and next
            candidate = candidate.next;
            populateCandidate(candidate);

            if (!replicationVisitor.add(curr.owningUnit.group))
                continue;    // If we've already seen this group, the token cannot be affected.

            populateTokenInfoAndAdjustUnit(curr, newUnit.group);
        }

        replicationVisitor.clean();
    }

    /**
     * Calculates the {@code replicationStart} of a token, as well as {@code replicationThreshold} which is chosen in a way
     * that permits {@code findUpdatedReplicationStart} to quickly identify changes in ownership.
     */
    private Token populateTokenInfo(BaseTokenInfo<Unit, ?> token, GroupInfo newUnitGroup)
    {
        GroupInfo tokenGroup = token.owningUnit.group;
        PopulateVisitor visitor = new PopulateVisitor();

        // Replication start = the end of a token from the RF'th different group seen before the token.
        Token replicationStart;
        // The end of a token from the RF-1'th different group seen before the token.
        Token replicationThreshold = token.token;
        GroupInfo currGroup;
        for (TokenInfo<Unit> curr = token.prevInRing(); ; curr = curr.prev)
        {
            replicationStart = curr.token;
            currGroup = curr.owningUnit.group;
            if (!visitor.add(currGroup))
                continue; // Group is already seen.
            if (visitor.visitedAll())
                break;

            replicationThreshold = replicationStart;
            // Another instance of the same group precedes us in the replication range of the ring,
            // so this is where our replication range begins
            if (currGroup == tokenGroup)
                break;
        }
        if (newUnitGroup == tokenGroup)
            // new token is always a boundary (as long as it's closer than replicationStart)
            replicationThreshold = token.token;
        else if (newUnitGroup != currGroup && visitor.seen(newUnitGroup))
            // already has new group in replication span before last seen. cannot be affected
            replicationThreshold = replicationStart;
        visitor.clean();

        token.replicationThreshold = replicationThreshold;
        token.replicationStart = replicationStart;
        return replicationStart;
    }

    private void populateTokenInfoAndAdjustUnit(TokenInfo<Unit> populate, GroupInfo newUnitGroup)
    {
        Token replicationStart = populateTokenInfo(populate, newUnitGroup);
        double newOwnership = replicationStart.size(populate.token);
        double oldOwnership = populate.replicatedOwnership;
        populate.replicatedOwnership = newOwnership;
        populate.owningUnit.ownership += newOwnership - oldOwnership;
    }

    /**
     * Evaluates the improvement in variance for both units and individual tokens when candidate is inserted into the
     * ring.
     */
    private double evaluateImprovement(CandidateInfo<Unit> candidate, double optTokenOwnership, double newUnitMult)
    {
        double tokenChange = 0;

        UnitInfo<Unit> candidateUnit = candidate.owningUnit;
        Token candidateEnd = candidate.token;

        // Form a chain of units affected by the insertion to be able to qualify change of unit ownership.
        // A unit may be affected more than once.
        UnitAdjustmentTracker<Unit> unitTracker = new UnitAdjustmentTracker<>(candidateUnit);

        // Reflect change in ownership of the splitting token (candidate).
        tokenChange += applyOwnershipAdjustment(candidate, candidateUnit, candidate.replicationStart, candidateEnd, optTokenOwnership, unitTracker);

        // Loop through all vnodes that replicate candidate or split and update their ownership.
        ReplicationVisitor replicationVisitor = new ReplicationVisitor();
        for (TokenInfo<Unit> curr = candidate.split; !replicationVisitor.visitedAll(); curr = curr.next)
        {
            UnitInfo<Unit> currUnit = curr.owningUnit;

            if (!replicationVisitor.add(currUnit.group))
                continue;    // If this group is already seen, the token cannot be affected.

            Token replicationEnd = curr.token;
            Token replicationStart = findUpdatedReplicationStart(curr, candidate);
            tokenChange += applyOwnershipAdjustment(curr, currUnit, replicationStart, replicationEnd, optTokenOwnership, unitTracker);
        }
        replicationVisitor.clean();

        double nodeChange = unitTracker.calculateUnitChange(newUnitMult, optTokenOwnership);
        return -(tokenChange + nodeChange);
    }

    /**
     * Returns the start of the replication span for the token {@code curr} when {@code candidate} is inserted into the
     * ring.
     */
    private Token findUpdatedReplicationStart(TokenInfo<Unit> curr, CandidateInfo<Unit> candidate)
    {
        return furtherStartToken(curr.replicationThreshold, candidate.token, curr.token);
    }

    /**
     * Applies the ownership adjustment for the given element, updating tracked unit ownership and returning the change
     * of variance.
     */
    private double applyOwnershipAdjustment(BaseTokenInfo<Unit, ?> curr, UnitInfo<Unit> currUnit,
            Token replicationStart, Token replicationEnd,
            double optTokenOwnership, UnitAdjustmentTracker<Unit> unitTracker)
    {
        double oldOwnership = curr.replicatedOwnership;
        double newOwnership = replicationStart.size(replicationEnd);
        double tokenCount = currUnit.tokenCount;
        assert tokenCount > 0;
        unitTracker.add(currUnit, newOwnership - oldOwnership);
        return (sq(newOwnership - optTokenOwnership) - sq(oldOwnership - optTokenOwnership)) / sq(tokenCount);
    }

    /**
     * Tracker for unit ownership changes. The changes are tracked by a chain of UnitInfos where the adjustedOwnership
     * field is being updated as we see changes in token ownership.
     *
     * The chain ends with an element that points to itself; this element must be specified as argument to the
     * constructor as well as be the first unit with which 'add' is called; when calculating the variance change
     * a separate multiplier is applied to it (used to permit more freedom in choosing the first tokens of a unit).
     */
    private static class UnitAdjustmentTracker<Unit>
    {
        UnitInfo<Unit> unitsChain;

        UnitAdjustmentTracker(UnitInfo<Unit> newUnit)
        {
            unitsChain = newUnit;
        }

        void add(UnitInfo<Unit> currUnit, double diff)
        {
            if (currUnit.prevUsed == null)
            {
                assert unitsChain.prevUsed != null || currUnit == unitsChain;

                currUnit.adjustedOwnership = currUnit.ownership + diff;
                currUnit.prevUsed = unitsChain;
                unitsChain = currUnit;
            }
            else
            {
                currUnit.adjustedOwnership += diff;
            }
        }

        double calculateUnitChange(double newUnitMult, double optTokenOwnership)
        {
            double unitChange = 0;
            UnitInfo<Unit> unitsChain = this.unitsChain;
            // Now loop through the units chain and add the unit-level changes. Also clear the groups' seen marks.
            while (true)
            {
                double newOwnership = unitsChain.adjustedOwnership;
                double oldOwnership = unitsChain.ownership;
                double tokenCount = unitsChain.tokenCount;
                double diff = (sq(newOwnership / tokenCount - optTokenOwnership) - sq(oldOwnership / tokenCount - optTokenOwnership));
                UnitInfo<Unit> prev = unitsChain.prevUsed;
                unitsChain.prevUsed = null;
                if (unitsChain != prev)
                    unitChange += diff;
                else
                {
                    unitChange += diff * newUnitMult;
                    break;
                }
                unitsChain = prev;
            }
            this.unitsChain = unitsChain;
            return unitChange;
        }
    }


    /**
     * Helper class for marking/unmarking visited a chain of groups
     */
    private abstract class GroupVisitor
    {
        GroupInfo groupChain = GroupInfo.TERMINATOR;
        int seen = 0;

        abstract GroupInfo prevSeen(GroupInfo group);
        abstract void setPrevSeen(GroupInfo group, GroupInfo prevSeen);

        // true iff this is the first time we've visited this group
        boolean add(GroupInfo group)
        {
            if (prevSeen(group) != null)
                return false;
            ++seen;
            setPrevSeen(group, groupChain);
            groupChain = group;
            return true;
        }

        boolean visitedAll()
        {
            return seen >= replicas;
        }

        boolean seen(GroupInfo group)
        {
            return prevSeen(group) != null;
        }

        // Clean group seen markers.
        void clean()
        {
            GroupInfo groupChain = this.groupChain;
            while (groupChain != GroupInfo.TERMINATOR)
            {
                GroupInfo prev = prevSeen(groupChain);
                setPrevSeen(groupChain, null);
                groupChain = prev;
            }
            this.groupChain = GroupInfo.TERMINATOR;
        }
    }

    private class ReplicationVisitor extends GroupVisitor
    {
        GroupInfo prevSeen(GroupInfo group)
        {
            return group.prevSeen;
        }

        void setPrevSeen(GroupInfo group, GroupInfo prevSeen)
        {
            group.prevSeen = prevSeen;
        }
    }

    private class PopulateVisitor extends GroupVisitor
    {
        GroupInfo prevSeen(GroupInfo group)
        {
            return group.prevPopulate;
        }

        void setPrevSeen(GroupInfo group, GroupInfo prevSeen)
        {
            group.prevPopulate = prevSeen;
        }
    }

    private double optimalTokenOwnership(int tokensToAdd)
    {
        return 1.0 * replicas / (sortedTokens.size() + tokensToAdd);
    }

    /**
     * Selects from {@code t1}, {@code t2} the token that forms a bigger range with {@code towards} as the upper bound,
     * taking into account wrapping.
     * Unlike Token.size(), equality is taken to mean "same as" rather than covering the whole range.
     */
    private static Token furtherStartToken(Token t1, Token t2, Token towards)
    {
        if (t1.equals(towards))
            return t2;
        if (t2.equals(towards))
            return t1;

        return t1.size(towards) > t2.size(towards) ? t1 : t2;
    }

    private static double sq(double d)
    {
        return d * d;
    }


    /**
     * For testing, remove the given unit preserving correct state of the allocator.
     */
    void removeUnit(Unit n)
    {
        Collection<Token> tokens = unitToTokens.removeAll(n);
        sortedTokens.keySet().removeAll(tokens);
    }

    public int unitCount()
    {
        return unitToTokens.asMap().size();
    }

    public String toString()
    {
        return getClass().getSimpleName();
    }

    /**
     * TokenInfo about candidate new tokens/vnodes.
     */
    private static class CandidateInfo<Unit> extends BaseTokenInfo<Unit, CandidateInfo<Unit>>
    {
        // directly preceding token in the current token ring
        final TokenInfo<Unit> split;

        public CandidateInfo(Token token, TokenInfo<Unit> split, UnitInfo<Unit> owningUnit)
        {
            super(token, owningUnit);
            this.split = split;
        }

        TokenInfo<Unit> prevInRing()
        {
            return split.prev;
        }
    }

    static void dumpTokens(String lead, BaseTokenInfo<?, ?> tokens)
    {
        BaseTokenInfo<?, ?> token = tokens;
        do
        {
            System.out.format("%s%s: rs %s rt %s size %.2e%n", lead, token, token.replicationStart, token.replicationThreshold, token.replicatedOwnership);
            token = token.next;
        } while (token != null && token != tokens);
    }
}