package org.apache.helix.tools;

/*
 * 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.
 */

import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.Random;
import java.util.TreeMap;

import org.apache.helix.HelixException;
import org.apache.helix.ZNRecord;
import org.apache.helix.model.IdealState.IdealStateProperty;

/**
 * DefaultIdealStateCalculator tries to optimally allocate master/slave partitions among
 * espresso storage nodes.
 * Given a batch of storage nodes, the partition and replication factor, the algorithm first given a
 * initial state
 * When new batches of storage nodes are added, the algorithm will calculate the new ideal state
 * such that the total
 * partition movements are minimized.
 */
public class DefaultIdealStateCalculator {
  static final String _MasterAssignmentMap = "MasterAssignmentMap";
  static final String _SlaveAssignmentMap = "SlaveAssignmentMap";
  static final String _partitions = "partitions";
  static final String _replicas = "replicas";

  /**
   * Calculate the initial ideal state given a batch of storage instances, the replication factor
   * and
   * number of partitions
   * 1. Calculate the master assignment by random shuffling
   * 2. for each storage instance, calculate the 1st slave assignment map, by another random
   * shuffling
   * 3. for each storage instance, calculate the i-th slave assignment map
   * 4. Combine the i-th slave assignment maps together
   * @param instanceNames
   *          list of storage node instances
   * @param partitions
   *          number of partitions
   * @param replicas
   *          The number of replicas (slave partitions) per master partition
   * @param masterStateValue
   *          master state value: e.g. "MASTER" or "LEADER"
   * @param slaveStateValue
   *          slave state value: e.g. "SLAVE" or "STANDBY"
   * @param resourceName
   * @return a ZNRecord that contain the idealstate info
   */
  public static ZNRecord calculateIdealState(List<String> instanceNames, int partitions,
      int replicas, String resourceName, String masterStateValue, String slaveStateValue) {
    Collections.sort(instanceNames);
    if (instanceNames.size() < replicas + 1) {
      throw new HelixException("Number of instances must not be less than replicas + 1. "
          + "instanceNr:" + instanceNames.size() + ", replicas:" + replicas);
    } else if (partitions < instanceNames.size()) {
      ZNRecord idealState =
          IdealStateCalculatorByShuffling.calculateIdealState(instanceNames, partitions, replicas,
              resourceName, 12345, masterStateValue, slaveStateValue);
      int i = 0;
      for (String partitionId : idealState.getMapFields().keySet()) {
        Map<String, String> partitionAssignmentMap = idealState.getMapField(partitionId);
        List<String> partitionAssignmentPriorityList = new ArrayList<String>();
        String masterInstance = "";
        for (String instanceName : partitionAssignmentMap.keySet()) {
          if (partitionAssignmentMap.get(instanceName).equalsIgnoreCase(masterStateValue)
              && masterInstance.equals("")) {
            masterInstance = instanceName;
          } else {
            partitionAssignmentPriorityList.add(instanceName);
          }
        }
        Collections.shuffle(partitionAssignmentPriorityList, new Random(i++));
        partitionAssignmentPriorityList.add(0, masterInstance);
        idealState.setListField(partitionId, partitionAssignmentPriorityList);
      }
      return idealState;
    }

    Map<String, Object> result = calculateInitialIdealState(instanceNames, partitions, replicas);

    return convertToZNRecord(result, resourceName, masterStateValue, slaveStateValue);
  }

  public static ZNRecord calculateIdealStateBatch(List<List<String>> instanceBatches,
      int partitions, int replicas, String resourceName, String masterStateValue,
      String slaveStateValue) {
    Map<String, Object> result =
        calculateInitialIdealState(instanceBatches.get(0), partitions, replicas);

    for (int i = 1; i < instanceBatches.size(); i++) {
      result = calculateNextIdealState(instanceBatches.get(i), result);
    }

    return convertToZNRecord(result, resourceName, masterStateValue, slaveStateValue);
  }

  /**
   * Convert the internal result (stored as a Map<String, Object>) into ZNRecord.
   */
  public static ZNRecord convertToZNRecord(Map<String, Object> result, String resourceName,
      String masterStateValue, String slaveStateValue) {
    Map<String, List<Integer>> nodeMasterAssignmentMap =
        (Map<String, List<Integer>>) (result.get(_MasterAssignmentMap));
    Map<String, Map<String, List<Integer>>> nodeSlaveAssignmentMap =
        (Map<String, Map<String, List<Integer>>>) (result.get(_SlaveAssignmentMap));

    int partitions = (Integer) (result.get("partitions"));

    ZNRecord idealState = new ZNRecord(resourceName);
    idealState.setSimpleField(IdealStateProperty.NUM_PARTITIONS.toString(),
        String.valueOf(partitions));

    for (String instanceName : nodeMasterAssignmentMap.keySet()) {
      for (Integer partitionId : nodeMasterAssignmentMap.get(instanceName)) {
        String partitionName = resourceName + "_" + partitionId;
        if (!idealState.getMapFields().containsKey(partitionName)) {
          idealState.setMapField(partitionName, new TreeMap<String, String>());
        }
        idealState.getMapField(partitionName).put(instanceName, masterStateValue);
      }
    }

    for (String instanceName : nodeSlaveAssignmentMap.keySet()) {
      Map<String, List<Integer>> slaveAssignmentMap = nodeSlaveAssignmentMap.get(instanceName);

      for (String slaveNode : slaveAssignmentMap.keySet()) {
        List<Integer> slaveAssignment = slaveAssignmentMap.get(slaveNode);
        for (Integer partitionId : slaveAssignment) {
          String partitionName = resourceName + "_" + partitionId;
          idealState.getMapField(partitionName).put(slaveNode, slaveStateValue);
        }
      }
    }
    // generate the priority list of instances per partition. Master should be at front and slave
    // follows.

    for (String partitionId : idealState.getMapFields().keySet()) {
      Map<String, String> partitionAssignmentMap = idealState.getMapField(partitionId);
      List<String> partitionAssignmentPriorityList = new ArrayList<String>();
      String masterInstance = "";
      for (String instanceName : partitionAssignmentMap.keySet()) {
        if (partitionAssignmentMap.get(instanceName).equalsIgnoreCase(masterStateValue)
            && masterInstance.equals("")) {
          masterInstance = instanceName;
        } else {
          partitionAssignmentPriorityList.add(instanceName);
        }
      }
      Collections.shuffle(partitionAssignmentPriorityList);
      partitionAssignmentPriorityList.add(0, masterInstance);
      idealState.setListField(partitionId, partitionAssignmentPriorityList);
    }
    assert (result.containsKey("replicas"));
    idealState.setSimpleField(IdealStateProperty.REPLICAS.toString(), result.get("replicas")
        .toString());
    return idealState;
  }

  /**
   * Calculate the initial ideal state given a batch of storage instances, the replication factor
   * and
   * number of partitions
   * 1. Calculate the master assignment by random shuffling
   * 2. for each storage instance, calculate the 1st slave assignment map, by another random
   * shuffling
   * 3. for each storage instance, calculate the i-th slave assignment map
   * 4. Combine the i-th slave assignment maps together
   * @param instanceNames
   *          list of storage node instances
   * @param weight
   *          weight for the initial storage node (each node has the same weight)
   * @param partitions
   *          number of partitions
   * @param replicas
   *          The number of replicas (slave partitions) per master partition
   * @return a map that contain the idealstate info
   */
  public static Map<String, Object> calculateInitialIdealState(List<String> instanceNames,
      int partitions, int replicas) {
    Random r = new Random(54321);
    assert (replicas <= instanceNames.size() - 1);

    ArrayList<Integer> masterPartitionAssignment = new ArrayList<Integer>();
    for (int i = 0; i < partitions; i++) {
      masterPartitionAssignment.add(i);
    }
    // shuffle the partition id array
    Collections.shuffle(masterPartitionAssignment, new Random(r.nextInt()));

    // 1. Generate the random master partition assignment
    // instanceName -> List of master partitions on that instance
    Map<String, List<Integer>> nodeMasterAssignmentMap = new TreeMap<String, List<Integer>>();
    for (int i = 0; i < masterPartitionAssignment.size(); i++) {
      String instanceName = instanceNames.get(i % instanceNames.size());
      if (!nodeMasterAssignmentMap.containsKey(instanceName)) {
        nodeMasterAssignmentMap.put(instanceName, new ArrayList<Integer>());
      }
      nodeMasterAssignmentMap.get(instanceName).add(masterPartitionAssignment.get(i));
    }

    // instanceName -> slave assignment for its master partitions
    // slave assignment: instanceName -> list of slave partitions on it
    List<Map<String, Map<String, List<Integer>>>> nodeSlaveAssignmentMapsList =
        new ArrayList<Map<String, Map<String, List<Integer>>>>(replicas);

    Map<String, Map<String, List<Integer>>> firstNodeSlaveAssignmentMap =
        new TreeMap<String, Map<String, List<Integer>>>();
    Map<String, Map<String, List<Integer>>> combinedNodeSlaveAssignmentMap =
        new TreeMap<String, Map<String, List<Integer>>>();

    if (replicas > 0) {
      // 2. For each node, calculate the evenly distributed slave as the first slave assignment
      // We will figure out the 2nd ...replicas-th slave assignment based on the first level slave
      // assignment
      for (int i = 0; i < instanceNames.size(); i++) {
        List<String> slaveInstances = new ArrayList<String>();
        ArrayList<Integer> slaveAssignment = new ArrayList<Integer>();
        TreeMap<String, List<Integer>> slaveAssignmentMap = new TreeMap<String, List<Integer>>();

        for (int j = 0; j < instanceNames.size(); j++) {
          if (j != i) {
            slaveInstances.add(instanceNames.get(j));
            slaveAssignmentMap.put(instanceNames.get(j), new ArrayList<Integer>());
          }
        }
        // Get the number of master partitions on instanceName
        List<Integer> masterAssignment = nodeMasterAssignmentMap.get(instanceNames.get(i));
        // do a random shuffling as in step 1, so that the first-level slave are distributed among
        // rest instances

        for (int j = 0; j < masterAssignment.size(); j++) {
          slaveAssignment.add(j);
        }
        Collections.shuffle(slaveAssignment, new Random(r.nextInt()));

        Collections.shuffle(slaveInstances, new Random(instanceNames.get(i).hashCode()));

        // Get the slave assignment map of node instanceName
        for (int j = 0; j < masterAssignment.size(); j++) {
          String slaveInstanceName =
              slaveInstances.get(slaveAssignment.get(j) % slaveInstances.size());
          if (!slaveAssignmentMap.containsKey(slaveInstanceName)) {
            slaveAssignmentMap.put(slaveInstanceName, new ArrayList<Integer>());
          }
          slaveAssignmentMap.get(slaveInstanceName).add(masterAssignment.get(j));
        }
        firstNodeSlaveAssignmentMap.put(instanceNames.get(i), slaveAssignmentMap);
      }
      nodeSlaveAssignmentMapsList.add(firstNodeSlaveAssignmentMap);
      // From the first slave assignment map, calculate the rest slave assignment maps
      for (int replicaOrder = 1; replicaOrder < replicas; replicaOrder++) {
        // calculate the next slave partition assignment map
        Map<String, Map<String, List<Integer>>> nextNodeSlaveAssignmentMap =
            calculateNextSlaveAssignemntMap(firstNodeSlaveAssignmentMap, replicaOrder);
        nodeSlaveAssignmentMapsList.add(nextNodeSlaveAssignmentMap);
      }

      // Combine the calculated 1...replicas-th slave assignment map together

      for (String instanceName : nodeMasterAssignmentMap.keySet()) {
        Map<String, List<Integer>> combinedSlaveAssignmentMap =
            new TreeMap<String, List<Integer>>();

        for (Map<String, Map<String, List<Integer>>> slaveNodeAssignmentMap : nodeSlaveAssignmentMapsList) {
          Map<String, List<Integer>> slaveAssignmentMap = slaveNodeAssignmentMap.get(instanceName);

          for (String slaveInstance : slaveAssignmentMap.keySet()) {
            if (!combinedSlaveAssignmentMap.containsKey(slaveInstance)) {
              combinedSlaveAssignmentMap.put(slaveInstance, new ArrayList<Integer>());
            }
            combinedSlaveAssignmentMap.get(slaveInstance).addAll(
                slaveAssignmentMap.get(slaveInstance));
          }
        }
        migrateSlaveAssignMapToNewInstances(combinedSlaveAssignmentMap, new ArrayList<String>());
        combinedNodeSlaveAssignmentMap.put(instanceName, combinedSlaveAssignmentMap);
      }
    }
    /*
     * // Print the result master and slave assignment maps
     * System.out.println("Master assignment:");
     * for(String instanceName : nodeMasterAssignmentMap.keySet())
     * {
     * System.out.println(instanceName+":");
     * for(Integer x : nodeMasterAssignmentMap.get(instanceName))
     * {
     * System.out.print(x+" ");
     * }
     * System.out.println();
     * System.out.println("Slave assignment:");
     * int slaveOrder = 1;
     * for(Map<String, Map<String, List<Integer>>> slaveNodeAssignmentMap :
     * nodeSlaveAssignmentMapsList)
     * {
     * System.out.println("Slave assignment order :" + (slaveOrder++));
     * Map<String, List<Integer>> slaveAssignmentMap = slaveNodeAssignmentMap.get(instanceName);
     * for(String slaveName : slaveAssignmentMap.keySet())
     * {
     * System.out.print("\t" + slaveName +":\n\t" );
     * for(Integer x : slaveAssignmentMap.get(slaveName))
     * {
     * System.out.print(x + " ");
     * }
     * System.out.println("\n");
     * }
     * }
     * System.out.println("\nCombined slave assignment map");
     * Map<String, List<Integer>> slaveAssignmentMap =
     * combinedNodeSlaveAssignmentMap.get(instanceName);
     * for(String slaveName : slaveAssignmentMap.keySet())
     * {
     * System.out.print("\t" + slaveName +":\n\t" );
     * for(Integer x : slaveAssignmentMap.get(slaveName))
     * {
     * System.out.print(x + " ");
     * }
     * System.out.println("\n");
     * }
     * }
     */
    Map<String, Object> result = new TreeMap<String, Object>();
    result.put("MasterAssignmentMap", nodeMasterAssignmentMap);
    result.put("SlaveAssignmentMap", combinedNodeSlaveAssignmentMap);
    result.put("replicas", new Integer(replicas));
    result.put("partitions", new Integer(partitions));
    return result;
  }

  /**
   * In the case there are more than 1 slave, we use the following algorithm to calculate the n-th
   * slave
   * assignment map based on the first level slave assignment map.
   * @param firstInstanceSlaveAssignmentMap the first slave assignment map for all instances
   * @param order of the slave
   * @return the n-th slave assignment map for all the instances
   */
  static Map<String, Map<String, List<Integer>>> calculateNextSlaveAssignemntMap(
      Map<String, Map<String, List<Integer>>> firstInstanceSlaveAssignmentMap, int replicaOrder) {
    Map<String, Map<String, List<Integer>>> result =
        new TreeMap<String, Map<String, List<Integer>>>();

    for (String currentInstance : firstInstanceSlaveAssignmentMap.keySet()) {
      Map<String, List<Integer>> resultAssignmentMap = new TreeMap<String, List<Integer>>();
      result.put(currentInstance, resultAssignmentMap);
    }

    for (String currentInstance : firstInstanceSlaveAssignmentMap.keySet()) {
      Map<String, List<Integer>> previousSlaveAssignmentMap =
          firstInstanceSlaveAssignmentMap.get(currentInstance);
      Map<String, List<Integer>> resultAssignmentMap = result.get(currentInstance);
      int offset = replicaOrder - 1;
      for (String instance : previousSlaveAssignmentMap.keySet()) {
        List<String> otherInstances = new ArrayList<String>(previousSlaveAssignmentMap.size() - 1);
        // Obtain an array of other instances
        for (String otherInstance : previousSlaveAssignmentMap.keySet()) {
          otherInstances.add(otherInstance);
        }
        Collections.sort(otherInstances);
        int instanceIndex = -1;
        for (int index = 0; index < otherInstances.size(); index++) {
          if (otherInstances.get(index).equalsIgnoreCase(instance)) {
            instanceIndex = index;
          }
        }
        assert (instanceIndex >= 0);
        if (instanceIndex == otherInstances.size() - 1) {
          instanceIndex--;
        }
        // Since we need to evenly distribute the slaves on "instance" to other partitions, we
        // need to remove "instance" from the array.
        otherInstances.remove(instance);

        // distribute previous slave assignment to other instances.
        List<Integer> previousAssignmentList = previousSlaveAssignmentMap.get(instance);
        for (int i = 0; i < previousAssignmentList.size(); i++) {

          // Evenly distribute the previousAssignmentList to the remaining other instances
          int newInstanceIndex = (i + offset + instanceIndex) % otherInstances.size();
          String newInstance = otherInstances.get(newInstanceIndex);
          if (!resultAssignmentMap.containsKey(newInstance)) {
            resultAssignmentMap.put(newInstance, new ArrayList<Integer>());
          }
          resultAssignmentMap.get(newInstance).add(previousAssignmentList.get(i));
        }
      }
    }
    return result;
  }

  /**
   * Given the current idealState, and the list of new Instances needed to be added, calculate the
   * new Ideal state.
   * 1. Calculate how many master partitions should be moved to the new cluster of instances
   * 2. assign the number of master partitions px to be moved to each previous node
   * 3. for each previous node,
   * 3.1 randomly choose px nodes, move them to temp list
   * 3.2 for each px nodes, remove them from the slave assignment map; record the map position of
   * the partition;
   * 3.3 calculate # of new nodes that should be put in the slave assignment map
   * 3.4 even-fy the slave assignment map;
   * 3.5 randomly place # of new nodes that should be placed in
   * 4. from all the temp master node list get from 3.1,
   * 4.1 randomly assign them to nodes in the new cluster
   * 5. for each node in the new cluster,
   * 5.1 assemble the slave assignment map
   * 5.2 even-fy the slave assignment map
   * @param newInstances
   *          list of new added storage node instances
   * @param weight
   *          weight for the new storage nodes (each node has the same weight)
   * @param previousIdealState
   *          The previous ideal state
   * @return a map that contain the updated idealstate info
   */
  public static Map<String, Object> calculateNextIdealState(List<String> newInstances,
      Map<String, Object> previousIdealState) {
    // Obtain the master / slave assignment info maps
    Collections.sort(newInstances);
    Map<String, List<Integer>> previousMasterAssignmentMap =
        (Map<String, List<Integer>>) (previousIdealState.get("MasterAssignmentMap"));
    Map<String, Map<String, List<Integer>>> nodeSlaveAssignmentMap =
        (Map<String, Map<String, List<Integer>>>) (previousIdealState.get("SlaveAssignmentMap"));

    List<String> oldInstances = new ArrayList<String>();
    for (String oldInstance : previousMasterAssignmentMap.keySet()) {
      oldInstances.add(oldInstance);
    }

    int previousInstanceNum = previousMasterAssignmentMap.size();
    int partitions = (Integer) (previousIdealState.get("partitions"));

    // TODO: take weight into account when calculate this

    int totalMasterParitionsToMove =
        partitions * (newInstances.size()) / (previousInstanceNum + newInstances.size());
    int numMastersFromEachNode = totalMasterParitionsToMove / previousInstanceNum;
    int remain = totalMasterParitionsToMove % previousInstanceNum;

    // Note that when remain > 0, we should make [remain] moves with (numMastersFromEachNode + 1)
    // partitions.
    // And we should first choose those (numMastersFromEachNode + 1) moves from the instances that
    // has more
    // master partitions
    List<Integer> masterPartitionListToMove = new ArrayList<Integer>();

    // For corresponding moved slave partitions, keep track of their original location; the new node
    // does not
    // need to migrate all of them.
    Map<String, List<Integer>> slavePartitionsToMoveMap = new TreeMap<String, List<Integer>>();

    // Make sure that the instances that holds more master partitions are put in front
    List<String> bigList = new ArrayList<String>(), smallList = new ArrayList<String>();
    for (String oldInstance : previousMasterAssignmentMap.keySet()) {
      List<Integer> masterAssignmentList = previousMasterAssignmentMap.get(oldInstance);
      if (masterAssignmentList.size() > numMastersFromEachNode) {
        bigList.add(oldInstance);
      } else {
        smallList.add(oldInstance);
      }
    }
    // "sort" the list, such that the nodes that has more master partitions moves more partitions to
    // the
    // new added batch of instances.
    bigList.addAll(smallList);
    int totalSlaveMoves = 0;
    for (String oldInstance : bigList) {
      List<Integer> masterAssignmentList = previousMasterAssignmentMap.get(oldInstance);
      int numToChoose = numMastersFromEachNode;
      if (remain > 0) {
        numToChoose = numMastersFromEachNode + 1;
        remain--;
      }
      // randomly remove numToChoose of master partitions to the new added nodes
      ArrayList<Integer> masterPartionsMoved = new ArrayList<Integer>();
      randomSelect(masterAssignmentList, masterPartionsMoved, numToChoose);

      masterPartitionListToMove.addAll(masterPartionsMoved);
      Map<String, List<Integer>> slaveAssignmentMap = nodeSlaveAssignmentMap.get(oldInstance);
      removeFromSlaveAssignmentMap(slaveAssignmentMap, masterPartionsMoved,
          slavePartitionsToMoveMap);

      // Make sure that for old instances, the slave placement map is evenly distributed
      // Trace the "local slave moves", which should together contribute to most of the slave
      // migrations
      int movesWithinInstance =
          migrateSlaveAssignMapToNewInstances(slaveAssignmentMap, newInstances);
      // System.out.println("local moves: "+ movesWithinInstance);
      totalSlaveMoves += movesWithinInstance;
    }
    // System.out.println("local slave moves total: "+ totalSlaveMoves);
    // calculate the master /slave assignment for the new added nodes

    // We already have the list of master partitions that will migrate to new batch of instances,
    // shuffle the partitions and assign them to new instances
    Collections.shuffle(masterPartitionListToMove, new Random(12345));
    for (int i = 0; i < newInstances.size(); i++) {
      String newInstance = newInstances.get(i);
      List<Integer> masterPartitionList = new ArrayList<Integer>();
      for (int j = 0; j < masterPartitionListToMove.size(); j++) {
        if (j % newInstances.size() == i) {
          masterPartitionList.add(masterPartitionListToMove.get(j));
        }
      }

      Map<String, List<Integer>> slavePartitionMap = new TreeMap<String, List<Integer>>();
      for (String oldInstance : oldInstances) {
        slavePartitionMap.put(oldInstance, new ArrayList<Integer>());
      }
      // Build the slave assignment map for the new instance, based on the saved information
      // about those slave partition locations in slavePartitionsToMoveMap
      for (Integer x : masterPartitionList) {
        for (String oldInstance : slavePartitionsToMoveMap.keySet()) {
          List<Integer> slaves = slavePartitionsToMoveMap.get(oldInstance);
          if (slaves.contains(x)) {
            slavePartitionMap.get(oldInstance).add(x);
          }
        }
      }
      // add entry for other new instances into the slavePartitionMap
      List<String> otherNewInstances = new ArrayList<String>();
      for (String instance : newInstances) {
        if (!instance.equalsIgnoreCase(newInstance)) {
          otherNewInstances.add(instance);
        }
      }
      // Make sure that slave partitions are evenly distributed
      migrateSlaveAssignMapToNewInstances(slavePartitionMap, otherNewInstances);

      // Update the result in the result map. We can reuse the input previousIdealState map as
      // the result.
      previousMasterAssignmentMap.put(newInstance, masterPartitionList);
      nodeSlaveAssignmentMap.put(newInstance, slavePartitionMap);

    }
    /*
     * // Print content of the master/ slave assignment maps
     * for(String instanceName : previousMasterAssignmentMap.keySet())
     * {
     * System.out.println(instanceName+":");
     * for(Integer x : previousMasterAssignmentMap.get(instanceName))
     * {
     * System.out.print(x+" ");
     * }
     * System.out.println("\nmaster partition moved:");
     * System.out.println();
     * System.out.println("Slave assignment:");
     * Map<String, List<Integer>> slaveAssignmentMap = nodeSlaveAssignmentMap.get(instanceName);
     * for(String slaveName : slaveAssignmentMap.keySet())
     * {
     * System.out.print("\t" + slaveName +":\n\t" );
     * for(Integer x : slaveAssignmentMap.get(slaveName))
     * {
     * System.out.print(x + " ");
     * }
     * System.out.println("\n");
     * }
     * }
     * System.out.println("Master partitions migrated to new instances");
     * for(Integer x : masterPartitionListToMove)
     * {
     * System.out.print(x+" ");
     * }
     * System.out.println();
     * System.out.println("Slave partitions migrated to new instances");
     * for(String oldInstance : slavePartitionsToMoveMap.keySet())
     * {
     * System.out.print(oldInstance + ": ");
     * for(Integer x : slavePartitionsToMoveMap.get(oldInstance))
     * {
     * System.out.print(x+" ");
     * }
     * System.out.println();
     * }
     */
    return previousIdealState;
  }

  public ZNRecord calculateNextIdealState(List<String> newInstances,
      Map<String, Object> previousIdealState, String resourceName, String masterStateValue,
      String slaveStateValue) {
    return convertToZNRecord(calculateNextIdealState(newInstances, previousIdealState),
        resourceName, masterStateValue, slaveStateValue);
  }

  /**
   * Given the list of master partition that will be migrated away from the storage instance,
   * Remove their entries from the local instance slave assignment map.
   * @param slaveAssignmentMap the local instance slave assignment map
   * @param masterPartionsMoved the list of master partition ids that will be migrated away
   * @param removedAssignmentMap keep track of the removed slave assignment info. The info can be
   *          used by new added storage nodes.
   */
  static void removeFromSlaveAssignmentMap(Map<String, List<Integer>> slaveAssignmentMap,
      List<Integer> masterPartionsMoved, Map<String, List<Integer>> removedAssignmentMap) {
    for (String instanceName : slaveAssignmentMap.keySet()) {
      List<Integer> slaveAssignment = slaveAssignmentMap.get(instanceName);
      for (Integer partitionId : masterPartionsMoved) {
        if (slaveAssignment.contains(partitionId)) {
          slaveAssignment.remove(partitionId);
          if (!removedAssignmentMap.containsKey(instanceName)) {
            removedAssignmentMap.put(instanceName, new ArrayList<Integer>());
          }
          removedAssignmentMap.get(instanceName).add(partitionId);
        }
      }
    }
  }

  /**
   * Since some new storage instances are added, each existing storage instance should migrate some
   * slave partitions to the new added instances.
   * The algorithm keeps moving one partition to from the instance that hosts most slave partitions
   * to the instance that hosts least number of partitions, until max-min <= 1.
   * In this way we can guarantee that all instances hosts almost same number of slave partitions,
   * also
   * slave partitions are evenly distributed.
   * @param slaveAssignmentMap the local instance slave assignment map
   * @param masterPartionsMoved the list of master partition ids that will be migrated away
   * @param removedAssignmentMap keep track of the removed slave assignment info. The info can be
   *          used by new added storage nodes.
   */
  static int migrateSlaveAssignMapToNewInstances(Map<String, List<Integer>> nodeSlaveAssignmentMap,
      List<String> newInstances) {
    int moves = 0;
    boolean done = false;
    for (String newInstance : newInstances) {
      nodeSlaveAssignmentMap.put(newInstance, new ArrayList<Integer>());
    }
    while (!done) {
      List<Integer> maxAssignment = null, minAssignment = null;
      int minCount = Integer.MAX_VALUE, maxCount = Integer.MIN_VALUE;
      String minInstance = "";
      for (String instanceName : nodeSlaveAssignmentMap.keySet()) {
        List<Integer> slaveAssignment = nodeSlaveAssignmentMap.get(instanceName);
        if (minCount > slaveAssignment.size()) {
          minCount = slaveAssignment.size();
          minAssignment = slaveAssignment;
          minInstance = instanceName;
        }
        if (maxCount < slaveAssignment.size()) {
          maxCount = slaveAssignment.size();
          maxAssignment = slaveAssignment;
        }
      }
      if (maxCount - minCount <= 1) {
        done = true;
      } else {
        int indexToMove = -1;
        // find a partition that is not contained in the minAssignment list
        for (int i = 0; i < maxAssignment.size(); i++) {
          if (!minAssignment.contains(maxAssignment.get(i))) {
            indexToMove = i;
            break;
          }
        }

        minAssignment.add(maxAssignment.get(indexToMove));
        maxAssignment.remove(indexToMove);

        if (newInstances.contains(minInstance)) {
          moves++;
        }
      }
    }
    return moves;
  }

  /**
   * Randomly select a number of elements from original list and put them in the selectedList
   * The algorithm is used to select master partitions to be migrated when new instances are added.
   * @param originalList the original list
   * @param selectedList the list that contain selected elements
   * @param num number of elements to be selected
   */
  static void randomSelect(List<Integer> originalList, List<Integer> selectedList, int num) {
    assert (originalList.size() >= num);
    int[] indexArray = new int[originalList.size()];
    for (int i = 0; i < indexArray.length; i++) {
      indexArray[i] = i;
    }
    int numRemains = originalList.size();
    Random r = new Random(numRemains);
    for (int j = 0; j < num; j++) {
      int randIndex = r.nextInt(numRemains);
      selectedList.add(originalList.get(randIndex));
      originalList.remove(randIndex);
      numRemains--;
    }
  }

  public static void main(String args[]) {
    List<String> instanceNames = new ArrayList<String>();
    for (int i = 0; i < 10; i++) {
      instanceNames.add("localhost:123" + i);
    }
    int partitions = 48 * 3, replicas = 3;
    Map<String, Object> resultOriginal =
        DefaultIdealStateCalculator.calculateInitialIdealState(instanceNames, partitions, replicas);

  }
}
