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.HashMap;
  import java.util.List;
  import java.util.Map;
  import java.util.Random;
  import java.util.TreeMap;
  
  import org.apache.helix.ZNRecord;
  import org.apache.helix.model.IdealState.IdealStateProperty;
  
  
  public class IdealStateCalculatorByRush
  {
    /**
     * Build the config map for RUSH algorithm. The input of RUSH algorithm groups
     * nodes into "cluster"s, and different clusters can be assigned with
     * different weights.
     *
     * @param numClusters
     *          number of node clusters
     * @param instancesPerCluster
     *          List of clusters, each contain a list of node name strings.
     * @param replicationDegree
     *          the replication degree
     * @param clusterWeights
     *          the weight for each node cluster
     * @return this config map structure for RUSH algorithm.
     */
    static HashMap<String, Object> buildRushConfig(int numClusters,
        List<List<String>> instancesPerCluster, int replicationDegree,
        List<Integer> clusterWeights)
    {
      HashMap<String, Object> config = new HashMap<String, Object>();
      config.put("replicationDegree", replicationDegree);
  
      HashMap[] clusterList = new HashMap[numClusters];
      config.put("subClusters", clusterList);
  
      HashMap[] nodes;
      HashMap<String, String> node;
      HashMap<String, Object> clusterData;
      for (int n = 0; n < numClusters; n++)
      {
        int numNodes = instancesPerCluster.get(n).size();
        nodes = new HashMap[numNodes];
        for (int i = 0; i < numNodes; i++)
        {
          node = new HashMap<String, String>();
          node.put("partition", instancesPerCluster.get(n).get(i));
          nodes[i] = node;
        }
        clusterData = new HashMap<String, Object>();
        clusterData.put("weight", clusterWeights.get(n));
        clusterData.put("nodes", nodes);
        clusterList[n] = clusterData;
      }
      return config;
    }
  
    /**
     * Calculate the ideal state for list of instances clusters.
     *
     * @param numClusters
     *          number of node clusters
     * @param instanceClusters
     *          List of clusters, each contain a list of node name strings.
     * @param instanceClusterWeights
     *          the weight for each instance cluster
     * @param partitions
     *          the partition number of the database
     * @param replicas
     *          the replication degree
     * @param resourceName
     *          the name of the database
     * @return The ZNRecord that contains the ideal state
     */
    public static ZNRecord calculateIdealState(
        List<List<String>> instanceClusters,
       List<Integer> instanceClusterWeights, int partitions, int replicas,
       String resourceName) throws Exception
   {
     ZNRecord result = new ZNRecord(resourceName);
 
     int numberOfClusters = instanceClusters.size();
     List<List<String>> nodesInClusters = instanceClusters;
     List<Integer> clusterWeights = instanceClusterWeights;
 
     HashMap<String, Object> rushConfig = buildRushConfig(numberOfClusters,
         nodesInClusters, replicas + 1, clusterWeights);
     RUSHrHash rushHash = new RUSHrHash(rushConfig);
 
     Random r = new Random(0);
     for (int i = 0; i < partitions; i++)
     {
       int partitionId = i;
       String partitionName = resourceName + ".partition-" + partitionId;
 
       ArrayList<HashMap> partitionAssignmentResult = rushHash
           .findNode(i);
       List<String> nodeNames = new ArrayList<String>();
       for (HashMap<?, ?> p : partitionAssignmentResult)
       {
         for (Object key : p.keySet())
         {
           if (p.get(key) instanceof String)
           {
             nodeNames.add(p.get(key).toString());
           }
         }
       }
       Map<String, String> partitionAssignment = new TreeMap<String, String>();
 
       for (int j = 0; j < nodeNames.size(); j++)
       {
         partitionAssignment.put(nodeNames.get(j), "SLAVE");
       }
       int master = r.nextInt(nodeNames.size());
       //master = nodeNames.size()/2;
       partitionAssignment.put(nodeNames.get(master), "MASTER");
 
 
       result.setMapField(partitionName, partitionAssignment);
     }
     result.setSimpleField(IdealStateProperty.NUM_PARTITIONS.toString(), String.valueOf(partitions));
     return result;
   }
 
   public static ZNRecord calculateIdealState(
       List<String> instanceClusters,
       int instanceClusterWeight, int partitions, int replicas,
       String resourceName) throws Exception
   {
     List<List<String>> instanceClustersList = new ArrayList<List<String>>();
     instanceClustersList.add(instanceClusters);
 
     List<Integer> instanceClusterWeightList = new ArrayList<Integer>();
     instanceClusterWeightList.add(instanceClusterWeight);
 
     return calculateIdealState(
         instanceClustersList,
         instanceClusterWeightList, partitions, replicas,
         resourceName);
   }
   /**
    * Helper function to see how many partitions are mapped to different
    * instances in two ideal states
    * */
   public static void printDiff(ZNRecord record1, ZNRecord record2)
   {
     int diffCount = 0;
     int diffCountMaster = 0;
     for (String key : record1.getMapFields().keySet())
     {
       Map<String, String> map1 = record1.getMapField(key);
       Map<String, String> map2 = record2.getMapField(key);
 
       for (String k : map1.keySet())
       {
         if (!map2.containsKey(k))
         {
           diffCount++;
         }
         else if (!map1.get(k).equalsIgnoreCase(map2.get(k)))
         {
           diffCountMaster++;
         }
       }
     }
     System.out.println("\ndiff count = " + diffCount);
     System.out.println("\nmaster diff count:"+ diffCountMaster);
   }
 
   /**
    * Helper function to calculate and print the standard deviation of the
    * partition assignment ideal state.
    * */
   public static void printIdealStateStats(ZNRecord record)
   {
     Map<String, Integer> countsMap = new TreeMap<String, Integer>();
     Map<String, Integer> masterCountsMap = new TreeMap<String, Integer>();
     for (String key : record.getMapFields().keySet())
     {
       Map<String, String> map1 = record.getMapField(key);
       for (String k : map1.keySet())
       {
         if (!countsMap.containsKey(k))
         {
           countsMap.put(k, new Integer(0));
         }
         else
         {
           countsMap.put(k, countsMap.get(k).intValue() + 1);
         }
         if (!masterCountsMap.containsKey(k))
         {
           masterCountsMap.put(k, new Integer(0));
 
         }
         else if (map1.get(k).equalsIgnoreCase("MASTER"))
         {
           masterCountsMap.put(k, masterCountsMap.get(k).intValue() + 1);
         }
       }
     }
     double sum = 0;
     int maxCount = 0;
     int minCount = Integer.MAX_VALUE;
     for (String k : countsMap.keySet())
     {
       int count = countsMap.get(k);
       sum += count;
       if (maxCount < count)
       {
         maxCount = count;
       }
       if (minCount > count)
       {
         minCount = count;
       }
       System.out.print(count + " ");
     }
     System.out.println("\nMax count: " + maxCount + " min count:" + minCount);
     System.out.println("\n master:");
     double sumMaster = 0;
     int maxCountMaster = 0;
     int minCountMaster = Integer.MAX_VALUE;
     for (String k : masterCountsMap.keySet())
     {
       int count = masterCountsMap.get(k);
       sumMaster += count;
       if (maxCountMaster < count)
       {
         maxCountMaster = count;
       }
       if (minCountMaster > count)
       {
         minCountMaster = count;
       }
       System.out.print(count + " ");
     }
     System.out.println("\nMean master: "+ 1.0*sumMaster/countsMap.size());
     System.out.println("Max master count: " + maxCountMaster + " min count:" + minCountMaster);
     double mean = sum / (countsMap.size());
     // calculate the deviation of the node distribution
     double deviation = 0;
     for (String k : countsMap.keySet())
     {
       double count = countsMap.get(k);
       deviation += (count - mean) * (count - mean);
     }
     System.out.println("Mean: " + mean + " normal deviation:"
         + Math.sqrt(deviation / countsMap.size()) / mean);
 
     //System.out.println("Max count: " + maxCount + " min count:" + minCount);
     int steps = 10;
     int stepLen = (maxCount - minCount) / steps;
     if(stepLen == 0) return;
     List<Integer> histogram = new ArrayList<Integer>((maxCount - minCount)
         / stepLen + 1);
 
     for (int i = 0; i < (maxCount - minCount) / stepLen + 1; i++)
     {
       histogram.add(0);
     }
     for (String k : countsMap.keySet())
     {
       int count = countsMap.get(k);
       int stepNo = (count - minCount) / stepLen;
       histogram.set(stepNo, histogram.get(stepNo) + 1);
     }
     System.out.println("histogram:");
     for (Integer x : histogram)
     {
       System.out.print(x + " ");
     }
   }
 
   public static void main(String args[]) throws Exception
   {
     int partitions = 4096, replicas = 2;
     String resourceName = "espressoDB1";
     List<String> instanceNames = new ArrayList<String>();
     List<List<String>> instanceCluster1 = new ArrayList<List<String>>();
     for (int i = 0; i < 20; i++)
     {
       instanceNames.add("local"+i+"host_123" + i);
     }
     instanceCluster1.add(instanceNames);
     List<Integer> weights1 = new ArrayList<Integer>();
     weights1.add(1);
     ZNRecord result = IdealStateCalculatorByRush.calculateIdealState(
         instanceCluster1, weights1, partitions, replicas, resourceName);
 
     printIdealStateStats(result);
 
     List<String> instanceNames2 = new ArrayList<String>();
     for (int i = 400; i < 405; i++)
     {
       instanceNames2.add("localhost_123" + i);
     }
     instanceCluster1.add(instanceNames2);
     weights1.add(1);
     ZNRecord result2 = IdealStateCalculatorByRush.calculateIdealState(
         instanceCluster1, weights1, partitions, replicas, resourceName);
 
     printDiff(result, result2);
     printIdealStateStats(result2);
   }
 }