/*

  * 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.block_app.framework.api.local;

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

 import java.util.ArrayList;

 import java.util.Collections;

 import java.util.Iterator;

 import java.util.LinkedList;

 import java.util.List;

 import java.util.Map;

 import java.util.Queue;

 import java.util.concurrent.ConcurrentHashMap;

 import java.util.concurrent.ThreadLocalRandom;

 import org.apache.giraph.aggregators.Aggregator;

 import org.apache.giraph.block_app.framework.api.BlockMasterApi;

 import org.apache.giraph.block_app.framework.api.BlockOutputHandleAccessor;

 import org.apache.giraph.block_app.framework.api.BlockWorkerContextReceiveApi;

 import org.apache.giraph.block_app.framework.api.BlockWorkerContextSendApi;

 import org.apache.giraph.block_app.framework.api.BlockWorkerReceiveApi;

 import org.apache.giraph.block_app.framework.api.BlockWorkerSendApi;

 import org.apache.giraph.block_app.framework.api.BlockWorkerValueAccessor;

 import org.apache.giraph.block_app.framework.api.Counter;

 import org.apache.giraph.block_app.framework.api.local.InternalMessageStore.InternalChecksMessageStore;

 import org.apache.giraph.block_app.framework.api.local.InternalMessageStore.InternalWrappedMessageStore;

 import org.apache.giraph.block_app.framework.internal.BlockCounters;

 import org.apache.giraph.block_app.framework.internal.BlockWorkerContextLogic;

 import org.apache.giraph.block_app.framework.internal.BlockWorkerPieces;

 import org.apache.giraph.block_app.framework.output.BlockOutputDesc;

 import org.apache.giraph.block_app.framework.output.BlockOutputHandle;

 import org.apache.giraph.block_app.framework.output.BlockOutputWriter;

 import org.apache.giraph.block_app.framework.piece.global_comm.BroadcastHandle;

 import org.apache.giraph.block_app.framework.piece.global_comm.internal.ReducersForPieceHandler.BroadcastHandleImpl;

 import org.apache.giraph.comm.SendMessageCache.TargetVertexIdIterator;

 import org.apache.giraph.comm.messages.PartitionSplitInfo;

 import org.apache.giraph.conf.GiraphConstants;

 import org.apache.giraph.conf.ImmutableClassesGiraphConfiguration;

 import org.apache.giraph.conf.MessageClasses;

 import org.apache.giraph.edge.Edge;

 import org.apache.giraph.edge.OutEdges;

 import org.apache.giraph.graph.Vertex;

 import org.apache.giraph.graph.VertexMutations;

 import org.apache.giraph.graph.VertexResolver;

 import org.apache.giraph.master.AggregatorToGlobalCommTranslation;

 import org.apache.giraph.partition.GraphPartitionerFactory;

 import org.apache.giraph.partition.Partition;

 import org.apache.giraph.reducers.ReduceOperation;

 import org.apache.giraph.utils.TestGraph;

 import org.apache.giraph.utils.WritableUtils;

 import org.apache.giraph.worker.WorkerAggregatorDelegator;

 import org.apache.giraph.worker.WorkerGlobalCommUsage;

 import org.apache.hadoop.io.Writable;

 import org.apache.hadoop.io.WritableComparable;

 import com.google.common.base.Preconditions;

 import it.unimi.dsi.fastutil.ints.IntArrayList;

 import it.unimi.dsi.fastutil.ints.IntList;

 /**

  * Internal implementation of Block API interfaces - representing an in-memory

  * giraph instance.

  *

  * @param <I> Vertex id type

  * @param <V> Vertex value type

  * @param <E> Edge value type

  */

 @SuppressWarnings({ "rawtypes", "unchecked" })

 class InternalApi<I extends WritableComparable, V extends Writable,

     E extends Writable> implements BlockMasterApi, BlockOutputHandleAccessor {

   private final TestGraph<I, V, E> inputGraph;

   private final List<Partition<I, V, E>> partitions;

   private final GraphPartitionerFactory<I, V, E> partitionerFactory;

   private final ImmutableClassesGiraphConfiguration conf;

   private final boolean runAllChecks;

   private final InternalAggregators globalComm;

   private final AggregatorToGlobalCommTranslation aggregators;

   private final boolean createVertexOnMsgs;

   private final ConcurrentHashMap<I, VertexMutations<I, V, E>> mutations;

   private InternalMessageStore previousMessages;

   private InternalMessageStore nextMessages;

   private MessageClasses previousMessageClasses;

   private MessageClasses nextMessageClasses;

   private final InternalWorkerApi workerApi;

   private final BlockWorkerContextLogic workerContextLogic;

   private List<Writable> previousWorkerMessages;

   private List<Writable> nextWorkerMessages;

   public InternalApi(

       TestGraph<I, V, E> graph,

       ImmutableClassesGiraphConfiguration conf,

       int numPartitions,

       boolean runAllChecks) {

     this.inputGraph = graph;

     this.partitions = new ArrayList<>(numPartitions);

     for (int i = 0; i < numPartitions; i++) {

       this.partitions.add(conf.createPartition(i, null));

     }

     this.partitionerFactory = conf.createGraphPartitioner();

     Preconditions.checkNotNull(this.partitionerFactory);

     Preconditions.checkState(this.partitions.size() == numPartitions);

     for (Vertex<I, V, E> vertex : graph) {

       getPartition(vertex.getId()).putVertex(vertex);

     }

     graph.clear();

     this.conf = conf;

     this.runAllChecks = runAllChecks;

     this.globalComm = new InternalAggregators(runAllChecks);

     this.aggregators = new AggregatorToGlobalCommTranslation(conf, globalComm);

     this.mutations = new ConcurrentHashMap<>();

     this.workerApi = new InternalWorkerApi();

     this.workerApi.setConf(conf);

     this.workerApi.setWorkerGlobalCommUsage(this.globalComm);

     this.createVertexOnMsgs =

         GiraphConstants.RESOLVER_CREATE_VERTEX_ON_MSGS.get(conf);

     workerContextLogic = new BlockWorkerContextLogic();

   }

   /**

    * Wrapper for calling Worker API interface.

    * Needs to be separate from Master API, since getAggregatedValue

    * has different implementation on worker and on master.

    */

   class InternalWorkerApi extends WorkerAggregatorDelegator<I, V, E>

       implements BlockWorkerSendApi<I, V, E, Writable>,

       BlockWorkerReceiveApi<I>, BlockWorkerContextSendApi<I, Writable>,

       BlockWorkerContextReceiveApi<I>, BlockWorkerValueAccessor,

       WorkerGlobalCommUsage {

     @Override

     public void addVertexRequest(I id, V value) {

       addVertexRequest(id, value, conf.createAndInitializeOutEdges());

     }

     @Override

     public void addVertexRequest(I id, V value, OutEdges<I, E> edges) {

       Vertex<I, V, E> vertex = conf.createVertex();

       vertex.initialize(id, value, edges);

       getMutationFor(id).addVertex(vertex);

     }

     @Override

     public void removeVertexRequest(I vertexId) {

       getMutationFor(vertexId).removeVertex();

     }

     @Override

     public void addEdgeRequest(I sourceVertexId, Edge<I, E> edge) {

       getMutationFor(sourceVertexId).addEdge(edge);

     }

     @Override

     public void removeEdgesRequest(I sourceVertexId, I targetVertexId) {

       getMutationFor(sourceVertexId).removeEdge(targetVertexId);

     }

     @Override

     public void sendMessage(I id, Writable message) {

       nextMessages.sendMessage(id, message);

     }

     @Override

     public void sendMessageToAllEdges(

         Vertex<I, V, E> vertex, Writable message) {

       sendMessageToMultipleEdges(

           new TargetVertexIdIterator<>(vertex),

           message);

     }

     @Override

     public void sendMessageToMultipleEdges(

         Iterator<I> vertexIdIterator, Writable message) {

       nextMessages.sendMessageToMultipleEdges(vertexIdIterator, message);

     }

     @Override

     public int getMyWorkerIndex() {

       return 0;

     }

     @Override

     public int getWorkerCount() {

       return 1;

     }

     @Override

     public int getWorkerForVertex(I vertexId) {

       return 0;

     }

     @Override

     public void sendMessageToWorker(Writable message, int workerIndex) {

       Preconditions.checkArgument(workerIndex == getMyWorkerIndex(),

           "With just one worker you can only send worker message to itself, " +

               "but tried to send to " + workerIndex);

       nextWorkerMessages.add(message);

     }

     @Override

     public Object getWorkerValue() {

       return workerContextLogic.getWorkerValue();

     }

     @Override

     public long getTotalNumVertices() {

       return InternalApi.this.getTotalNumVertices();

     }

     @Override

     public long getTotalNumEdges() {

       return InternalApi.this.getTotalNumEdges();

     }

     @Override

     public <OW extends BlockOutputWriter, OD extends BlockOutputDesc<OW>>

     OD getOutputDesc(String confOption) {

       return workerContextLogic.getOutputHandle().<OW, OD>getOutputDesc(

           confOption);

     }

     @Override

     public <OW extends BlockOutputWriter> OW getWriter(String confOption) {

       return workerContextLogic.getOutputHandle().getWriter(confOption);

     }

     @Override

     public void setStatus(String status) {

     }

     @Override

     public void progress() {

     }

     @Override

     public Counter getCounter(final String group, final String name) {

       return BlockCounters.getNoOpCounter();

     }

   }

   @Override

   public void broadcast(String name, Writable value) {

     globalComm.broadcast(name, value);

   }

   @Override

   public <T extends Writable> BroadcastHandle<T> broadcast(T object) {

     BroadcastHandleImpl<T> handle = new BroadcastHandleImpl<>();

     broadcast(handle.getName(), object);

     return handle;

   }

   @Override

   public <S, R extends Writable> void registerReducer(

       String name, ReduceOperation<S, R> reduceOp) {

     globalComm.registerReducer(name, reduceOp);

   }

   @Override

   public <S, R extends Writable> void registerReducer(

       String name, ReduceOperation<S, R> reduceOp,

       R globalInitialValue) {

     globalComm.registerReducer(name, reduceOp, globalInitialValue);

   }

   @Override

   public <R extends Writable> R getReduced(String name) {

     return globalComm.getReduced(name);

   }

   @Override

   public <A extends Writable> A getAggregatedValue(String name) {

     return aggregators.getAggregatedValue(name);

   }

   @Override

   public <A extends Writable> void setAggregatedValue(String name, A value) {

     aggregators.setAggregatedValue(name, value);

   }

   @Override

   public <A extends Writable>

   boolean registerAggregator(

       String name, Class<? extends Aggregator<A>> aggregatorClass)

       throws InstantiationException, IllegalAccessException {

     return aggregators.registerAggregator(name, aggregatorClass);

   }

   @Override

   public <A extends Writable>

   boolean registerPersistentAggregator(

       String name, Class<? extends Aggregator<A>> aggregatorClass)

       throws InstantiationException, IllegalAccessException {

     return aggregators.registerPersistentAggregator(name, aggregatorClass);

   }

   @Override

   public ImmutableClassesGiraphConfiguration<I, V, E> getConf() {

     return conf;

   }

   @Override

   public void setStatus(String status) {

   }

   @Override

   public void progress() {

   }

   @Override

   public Counter getCounter(final String group, final String name) {

     return BlockCounters.getNoOpCounter();

   }

   private VertexMutations<I, V, E> getMutationFor(I vertexId) {

     VertexMutations<I, V, E> curMutations = new VertexMutations<>();

     VertexMutations<I, V, E> prevMutations =

         mutations.putIfAbsent(vertexId, curMutations);

     if (prevMutations != null) {

       curMutations = prevMutations;

     }

     return curMutations;

   }

   public Iterable takeMessages(I id) {

     if (previousMessages != null) {

       Iterable result = previousMessages.takeMessages(id);

       if (result != null) {

         return result;

       }

     }

     return Collections.emptyList();

   }

   public Iterable<I> getPartitionDestinationVertices(int partitionId) {

     if (previousMessages != null) {

       Iterable result =

           previousMessages.getPartitionDestinationVertices(partitionId);

       if (result != null) {

         return result;

       }

     }

     return Collections.emptyList();

   }

   public List<Writable> takeWorkerMessages() {

     if (previousWorkerMessages != null) {

       List<Writable> ret = new ArrayList<>(previousWorkerMessages.size());

       for (Writable message : previousWorkerMessages) {

         // Use message copies probabilistically, to catch both not serializing

         // some fields, and storing references from message object itself

         // (which can be reusable).

         ret.add(runAllChecks && ThreadLocalRandom.current().nextBoolean() ?

             WritableUtils.createCopy(message) : message);

       }

       previousWorkerMessages = null;

       if (runAllChecks) {

         Collections.shuffle(ret);

       }

       return ret;

     }

     return Collections.emptyList();

   }

   public void afterWorkerBeforeMaster() {

     globalComm.afterWorkerBeforeMaster();

     aggregators.prepareSuperstep();

   }

   public void afterMasterBeforeWorker() {

     aggregators.postMasterCompute();

   }

   public void afterMasterBeforeWorker(BlockWorkerPieces computation) {

     afterMasterBeforeWorker();

     previousMessages = nextMessages;

     previousMessageClasses = nextMessageClasses;

     previousWorkerMessages = nextWorkerMessages;

     nextMessageClasses = computation.getOutgoingMessageClasses(conf);

     nextMessages = createMessageStore(

       conf,

       nextMessageClasses,

       createPartitionInfo(),

       runAllChecks

     );

     nextWorkerMessages = new ArrayList<>();

     // finalize previous messages

     if (previousMessages != null) {

       previousMessages.finalizeStore();

     }

     boolean ignoreExistingVertices = ignoreExistingVertices();

     // process mutations:

     VertexResolver<I, V, E> vertexResolver = conf.createVertexResolver();

     for (Map.Entry<I, VertexMutations<I, V, E>> entry : mutations.entrySet()) {

       I vertexIndex = entry.getKey();

       Vertex<I, V, E> originalVertex =

           getPartition(vertexIndex).getVertex(vertexIndex);

       VertexMutations<I, V, E> curMutations = entry.getValue();

       Vertex<I, V, E> vertex = vertexResolver.resolve(

         vertexIndex,

         originalVertex,

         curMutations,

         !ignoreExistingVertices && previousMessages != null &&

         previousMessages.hasMessage(vertexIndex)

       );

       if (vertex != null) {

         getPartition(vertex.getId()).putVertex(vertex);

       } else if (originalVertex != null) {

         getPartition(originalVertex.getId()).removeVertex(

             originalVertex.getId());

         if (!ignoreExistingVertices && previousMessages != null) {

           previousMessages.takeMessages(originalVertex.getId());

         }

       }

     }

     mutations.clear();

     if (!ignoreExistingVertices && createVertexOnMsgs &&

         previousMessages != null) {

       Iterator<I> iter = previousMessages.targetVertexIds();

       while (iter.hasNext()) {

         I target = iter.next();

         if (getPartition(target).getVertex(target) == null) {

           // need a copy as the key might be reusable

           I copyId = WritableUtils.createCopy(target);

           Vertex<I, V, E> vertex =

               vertexResolver.resolve(copyId, null, null, true);

           if (vertex != null) {

             getPartition(vertex.getId()).putVertex(vertex);

           }

         }

       }

     }

   }

   public boolean ignoreExistingVertices() {

     return previousMessageClasses != null &&

         previousMessageClasses.ignoreExistingVertices();

   }

   private <M extends Writable>

   InternalMessageStore<I, M> createMessageStore(

     ImmutableClassesGiraphConfiguration<I, ?, ?> conf,

     MessageClasses<I, M> messageClasses,

     PartitionSplitInfo<I> partitionInfo,

     boolean runAllChecks

   ) {

     InternalMessageStore<I, M> messageStore =

       InternalWrappedMessageStore.create(conf, messageClasses, partitionInfo);

     if (runAllChecks) {

       return new InternalChecksMessageStore<I, M>(

           messageStore, conf, messageClasses.createMessageValueFactory(conf));

     } else {

       return messageStore;

     }

   }

   private PartitionSplitInfo<I> createPartitionInfo() {

     return new PartitionSplitInfo<I>() {

       /** Ids of partitions */

       private IntList partitionIds;

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

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

       @Override

       public int getPartitionId(I vertexId) {

         return partitionerFactory.getPartition(vertexId, partitions.size(), 1);

       }

       @Override

       public Iterable<Integer> getPartitionIds() {

         if (partitionIds == null) {

           partitionIds = new IntArrayList(partitions.size());

           for (int i = 0; i < partitions.size(); i++) {

             partitionIds.add(i);

           }

         }

         Preconditions.checkState(partitionIds.size() == partitions.size());

         return partitionIds;

       }

       @Override

       public long getPartitionVertexCount(Integer partitionId) {

         return partitions.get(partitionId).getVertexCount();

       }

       @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 LinkedList<Partition<I, V, E>>();

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

           partitionQueue.add(partition);

         }

       }

       @Override

       public Partition getNextPartition() {

         return partitionQueue.poll();

       }

       @Override

       public void putPartition(Partition partition) {

       }

     };

   }

   public List<Partition<I, V, E>> getPartitions() {

     return partitions;

   }

   public InternalWorkerApi getWorkerApi() {

     return workerApi;

   }

   @Override

   public long getTotalNumEdges() {

     int numEdges = 0;

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

       numEdges += partition.getEdgeCount();

     }

     return numEdges;

   }

   @Override

   public long getTotalNumVertices() {

     int numVertices = 0;

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

       numVertices += partition.getVertexCount();

     }

     return numVertices;

   }

   @Override

   public void logToCommandLine(String line) {

     System.err.println("Command line: " + line);

   }

   @Override

   public BlockOutputHandle getBlockOutputHandle() {

     return workerContextLogic.getOutputHandle();

   }

   @Override

   public <OW extends BlockOutputWriter,

       OD extends BlockOutputDesc<OW>> OD getOutputDesc(String confOption) {

     return workerContextLogic.getOutputHandle().<OW, OD>getOutputDesc(

         confOption);

   }

   @Override

   public <OW extends BlockOutputWriter> OW getWriter(String confOption) {

     return workerContextLogic.getOutputHandle().getWriter(confOption);

   }

   public BlockWorkerContextLogic getWorkerContextLogic() {

     return workerContextLogic;

   }

   @Override

   public int getWorkerCount() {

     return 1;

   }

   private int getPartitionId(I id) {

     Preconditions.checkNotNull(id);

     return partitionerFactory.getPartition(id, partitions.size(), 1);

   }

   private Partition<I, V, E> getPartition(I id) {

     return partitions.get(getPartitionId(id));

   }

   public void postApplication() {

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

       for (Vertex<I, V, E> vertex : partition) {

         inputGraph.setVertex(vertex);

       }

     }

   }

 }