/*

  * 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.giraph.partition;

 import com.google.common.collect.Maps;

 import org.apache.giraph.conf.ImmutableClassesGiraphConfiguration;

 import org.apache.giraph.utils.ExtendedDataOutput;

 import org.apache.giraph.utils.VertexIterator;

 import org.apache.hadoop.io.Writable;

 import org.apache.hadoop.io.WritableComparable;

 import org.apache.hadoop.mapreduce.Mapper;

 import java.util.concurrent.ArrayBlockingQueue;

 import java.util.concurrent.BlockingQueue;

 import java.util.concurrent.ConcurrentMap;

 import static com.google.common.base.Preconditions.checkState;

 /**

  * A simple in-memory partition store.

  *

  * @param <I> Vertex id

  * @param <V> Vertex data

  * @param <E> Edge data

  */

 public class SimplePartitionStore<I extends WritableComparable,

     V extends Writable, E extends Writable>

     implements PartitionStore<I, V, E> {

   /** Configuration */

   private final ImmutableClassesGiraphConfiguration<I, V, E> conf;

   /** Job context (for progress) */

   private final Mapper<?, ?, ?, ?>.Context context;

   /** Map of stored partitions. */

   private final ConcurrentMap<Integer, Partition<I, V, E>> partitions =

       Maps.newConcurrentMap();

   /** Queue of partitions to be precessed in a superstep */

   private BlockingQueue<Partition<I, V, E>> partitionQueue;

   /**

    * Constructor.

    * @param conf Configuration

    * @param context Mapper context

    */

   public SimplePartitionStore(ImmutableClassesGiraphConfiguration<I, V, E> conf,

       Mapper<?, ?, ?, ?>.Context context) {

     this.conf = conf;

     this.context = context;

   }

   @Override

   public boolean addPartition(Partition<I, V, E> partition) {

     return partitions.putIfAbsent(partition.getId(), partition) == null;

   }

   @Override

   public Partition<I, V, E> removePartition(Integer partitionId) {

     return partitions.remove(partitionId);

   }

   @Override

   public boolean hasPartition(Integer partitionId) {

     return partitions.containsKey(partitionId);

   }

   @Override

   public Iterable<Integer> getPartitionIds() {

     return partitions.keySet();

   }

   @Override

   public int getNumPartitions() {

     return partitions.size();

   }

   @Override

   public long getPartitionVertexCount(Integer partitionId) {

     Partition partition = partitions.get(partitionId);

     if (partition == null) {

       return 0;

     } else {

       return partition.getVertexCount();

     }

   }

   @Override

   public long getPartitionEdgeCount(Integer partitionId) {

     Partition partition = partitions.get(partitionId);

     if (partition == null) {

       return 0;

     } else {

       return partition.getEdgeCount();

     }

   }

   @Override

   public boolean isEmpty() {

     return partitions.size() == 0;

   }

   @Override

   public void startIteration() {

     checkState(partitionQueue == null || partitionQueue.isEmpty(),

         "startIteration: It seems that some of " +

           "of the partitions from previous iteration over partition store are" +

           " not yet processed.");

     partitionQueue =

         new ArrayBlockingQueue<Partition<I, V, E>>(getNumPartitions());

     for (Partition<I, V, E> partition : partitions.values()) {

       partitionQueue.add(partition);

     }

   }

   @Override

   public Partition<I, V, E> getNextPartition() {

     return partitionQueue.poll();

   }

   @Override

   public void putPartition(Partition<I, V, E> partition) { }

   /**

    * Get or create a partition.

    * @param partitionId Partition Id

    * @return The requested partition (never null)

    */

   private Partition<I, V, E> getOrCreatePartition(Integer partitionId) {

     Partition<I, V, E> oldPartition = partitions.get(partitionId);

     if (oldPartition == null) {

       Partition<I, V, E> newPartition =

           conf.createPartition(partitionId, context);

       oldPartition = partitions.putIfAbsent(partitionId, newPartition);

       if (oldPartition == null) {

         return newPartition;

       }

     }

     return oldPartition;

   }

   @Override

   public void addPartitionVertices(Integer partitionId,

       ExtendedDataOutput extendedDataOutput) {

     VertexIterator<I, V, E> vertexIterator =

         new VertexIterator<I, V, E>(extendedDataOutput, conf);

     Partition<I, V, E> partition = getOrCreatePartition(partitionId);

     partition.addPartitionVertices(vertexIterator);

     putPartition(partition);

   }

   @Override

   public void shutdown() { }

   @Override

   public void initialize() { }

 }