/*

  * 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.library;

 import java.util.ArrayList;

 import java.util.Iterator;

 import java.util.List;

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

 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.CreateReducersApi;

 import org.apache.giraph.block_app.framework.piece.Piece;

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

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

 import org.apache.giraph.block_app.framework.piece.global_comm.array.BroadcastArrayHandle;

 import org.apache.giraph.block_app.framework.piece.interfaces.VertexReceiver;

 import org.apache.giraph.block_app.framework.piece.interfaces.VertexSender;

 import org.apache.giraph.block_app.library.internal.SendMessagePiece;

 import org.apache.giraph.block_app.library.internal.SendMessageWithCombinerPiece;

 import org.apache.giraph.block_app.reducers.array.ArrayOfHandles;

 import org.apache.giraph.combiner.MessageCombiner;

 import org.apache.giraph.function.Consumer;

 import org.apache.giraph.function.PairConsumer;

 import org.apache.giraph.function.Supplier;

 import org.apache.giraph.function.vertex.ConsumerWithVertex;

 import org.apache.giraph.function.vertex.SupplierFromVertex;

 import org.apache.giraph.graph.Vertex;

 import org.apache.giraph.reducers.ReduceOperation;

 import org.apache.giraph.reducers.impl.SumReduce;

 import org.apache.giraph.types.NoMessage;

 import org.apache.hadoop.io.LongWritable;

 import org.apache.hadoop.io.Writable;

 import org.apache.hadoop.io.WritableComparable;

 import org.apache.log4j.Logger;

 /**

  * Utility class for creating common Pieces and computations for processing

  * graphs.

  */

 public class Pieces {

   private static final Logger LOG = Logger.getLogger(Pieces.class);

   private Pieces() { }

   /**

    * For each vertex execute given process function.

    * Computation is happening in send phase of the returned Piece.

    */

   public static

   <I extends WritableComparable, V extends Writable, E extends Writable>

   Piece<I, V, E, NoMessage, Object> forAllVertices(

       final String pieceName, final Consumer<Vertex<I, V, E>> process) {

     return new Piece<I, V, E, NoMessage, Object>() {

       @Override

       public VertexSender<I, V, E> getVertexSender(

           BlockWorkerSendApi<I, V, E, NoMessage> workerApi,

           Object executionStage) {

         return new InnerVertexSender() {

           @Override

           public void vertexSend(Vertex<I, V, E> vertex) {

             process.apply(vertex);

           }

         };

       }

       @Override

       public String toString() {

         return pieceName;

       }

     };

   }

   /**

    * Execute given function on master.

    */

   public static

   Piece<WritableComparable, Writable,  Writable, NoMessage,

     Object> masterCompute(

       final String pieceName, final Consumer<BlockMasterApi> process) {

     return new Piece<WritableComparable, Writable,  Writable, NoMessage,

         Object>() {

       @Override

       public void masterCompute(

           BlockMasterApi masterApi, Object executionStage) {

         process.apply(masterApi);

       }

     };

   }

   /**

    * For each vertex execute given process function.

    * Computation is happening in the receive phase of the returned Piece.

    * This function should be used if you need returned Piece to interact with

    * subsequent Piece, as that requires passed function to be executed

    * during receive phase,

    */

   public static

   <I extends WritableComparable, V extends Writable, E extends Writable>

   Piece<I, V, E, NoMessage, Object> forAllVerticesOnReceive(

       final String pieceName, final Consumer<Vertex<I, V, E>> process) {

     return new Piece<I, V, E, NoMessage, Object>() {

       @Override

       public VertexReceiver<I, V, E, NoMessage> getVertexReceiver(

           BlockWorkerReceiveApi<I> workerApi, Object executionStage) {

         return new InnerVertexReceiver() {

           @Override

           public void vertexReceive(

               Vertex<I, V, E> vertex, Iterable<NoMessage> messages) {

             process.apply(vertex);

           }

         };

       }

       @Override

       public String toString() {

         return pieceName;

       }

     };

   }

   /**

    * Creates Piece which removes vertices for which supplier returns true.

    */

   public static

   <I extends WritableComparable, V extends Writable, E extends Writable>

   Piece<I, V, E, NoMessage, Object> removeVertices(

       final String pieceName,

       final SupplierFromVertex<I, V, E, Boolean> shouldRemoveVertex) {

     return new Piece<I, V, E, NoMessage, Object>() {

       private ReducerHandle<LongWritable, LongWritable> countRemovedAgg;

       @Override

       public void registerReducers(

           CreateReducersApi reduceApi, Object executionStage) {

         countRemovedAgg = reduceApi.createLocalReducer(SumReduce.LONG);

       }

       @Override

       public VertexSender<I, V, E> getVertexSender(

           final BlockWorkerSendApi<I, V, E, NoMessage> workerApi,

           Object executionStage) {

         return new InnerVertexSender() {

           @Override

           public void vertexSend(Vertex<I, V, E> vertex) {

             if (shouldRemoveVertex.get(vertex)) {

               workerApi.removeVertexRequest(vertex.getId());

               reduceLong(countRemovedAgg, 1);

             }

           }

         };

       }

       @Override

       public void masterCompute(BlockMasterApi master, Object executionStage) {

         LOG.info("Removed " + countRemovedAgg.getReducedValue(master) +

             " vertices from the graph, during stage " + executionStage);

       }

       @Override

       public String toString() {

         return pieceName;

       }

     };

   }

   /**

    * Creates single reducer piece - given reduce class, supplier of values on

    * worker, reduces and passes the result to given consumer on master.

    *

    * @param <S> Single value type, objects passed on workers

    * @param <R> Reduced value type

    * @param <I> Vertex id type

    * @param <V> Vertex value type

    * @param <E> Edge value type

    */

   public static

   <S, R extends Writable, I extends WritableComparable, V extends Writable,

   E extends Writable>

   Piece<I, V, E, NoMessage, Object> reduce(

       String name,

       ReduceOperation<S, R> reduceOp,

       SupplierFromVertex<I, V, E, S> valueSupplier,

       final Consumer<R> reducedValueConsumer) {

     return reduceWithMaster(

         name, reduceOp, valueSupplier,

         new PairConsumer<R, BlockMasterApi>() {

           @Override

           public void apply(R input, BlockMasterApi master) {

             reducedValueConsumer.apply(input);

           }

         });

   }

   /**

    * Creates single reducer piece - given reduce class, supplier of values on

    * worker, reduces and passes the result to given consumer on master.

    *

    * @param <S> Single value type, objects passed on workers

    * @param <R> Reduced value type

    * @param <I> Vertex id type

    * @param <V> Vertex value type

    * @param <E> Edge value type

    */

   public static

   <S, R extends Writable, I extends WritableComparable, V extends Writable,

   E extends Writable>

   Piece<I, V, E, NoMessage, Object> reduceWithMaster(

       final String name,

       final ReduceOperation<S, R> reduceOp,

       final SupplierFromVertex<I, V, E, S> valueSupplier,

       final PairConsumer<R, BlockMasterApi> reducedValueConsumer) {

     return new Piece<I, V, E, NoMessage, Object>() {

       private ReducerHandle<S, R> handle;

       @Override

       public void registerReducers(

           CreateReducersApi reduceApi, Object executionStage) {

         handle = reduceApi.createLocalReducer(reduceOp);

       }

       @Override

       public VertexSender<I, V, E> getVertexSender(

           BlockWorkerSendApi<I, V, E, NoMessage> workerApi,

           Object executionStage) {

         return new InnerVertexSender() {

           @Override

           public void vertexSend(Vertex<I, V, E> vertex) {

             handle.reduce(valueSupplier.get(vertex));

           }

         };

       }

       @Override

       public void masterCompute(BlockMasterApi master, Object executionStage) {

         reducedValueConsumer.apply(handle.getReducedValue(master), master);

       }

       @Override

       public String toString() {

         return name;

       }

     };

   }

   /**

    * Creates single reducer and broadcast piece - given reduce class, supplier

    * of values on worker, reduces and broadcasts the value, passing it to the

    * consumer on worker for each vertex.

    *

    * @param <S> Single value type, objects passed on workers

    * @param <R> Reduced value type

    * @param <I> Vertex id type

    * @param <V> Vertex value type

    * @param <E> Edge value type

    */

   public static

   <S, R extends Writable, I extends WritableComparable, V extends Writable,

   E extends Writable>

   Piece<I, V, E, NoMessage, Object> reduceAndBroadcast(

       final String name,

       final ReduceOperation<S, R> reduceOp,

       final SupplierFromVertex<I, V, E, S> valueSupplier,

       final ConsumerWithVertex<I, V, E, R> reducedValueConsumer) {

     return new Piece<I, V, E, NoMessage, Object>() {

       private final ReducerAndBroadcastWrapperHandle<S, R> handle =

           new ReducerAndBroadcastWrapperHandle<>();

       @Override

       public void registerReducers(

           CreateReducersApi reduceApi, Object executionStage) {

         handle.registeredReducer(reduceApi.createLocalReducer(reduceOp));

       }

       @Override

       public VertexSender<I, V, E> getVertexSender(

           BlockWorkerSendApi<I, V, E, NoMessage> workerApi,

           Object executionStage) {

         return new InnerVertexSender() {

           @Override

           public void vertexSend(Vertex<I, V, E> vertex) {

             handle.reduce(valueSupplier.get(vertex));

           }

         };

       }

       @Override

       public void masterCompute(BlockMasterApi master, Object executionStage) {

         handle.broadcastValue(master);

       }

       @Override

       public VertexReceiver<I, V, E, NoMessage> getVertexReceiver(

           BlockWorkerReceiveApi<I> workerApi, Object executionStage) {

         final R value = handle.getBroadcast(workerApi);

         return new InnerVertexReceiver() {

           @Override

           public void vertexReceive(

               Vertex<I, V, E> vertex, Iterable<NoMessage> messages) {

             reducedValueConsumer.apply(vertex, value);

           }

         };

       }

       @Override

       public String toString() {

         return name;

       }

     };

   }

   /**

    * Like reduceAndBroadcast, but uses array of handles for reducers and

    * broadcasts, to make it feasible and performant when values are large.

    * Each supplied value to reduce will be reduced in the handle defined by

    * handleHashSupplier%numHandles

    *

    * @param <S> Single value type, objects passed on workers

    * @param <R> Reduced value type

    * @param <I> Vertex id type

    * @param <V> Vertex value type

    * @param <E> Edge value type

    */

   public static

   <S, R extends Writable, I extends WritableComparable, V extends Writable,

       E extends Writable>

   Piece<I, V, E, NoMessage, Object> reduceAndBroadcastWithArrayOfHandles(

       final String name,

       final int numHandles,

       final Supplier<ReduceOperation<S, R>> reduceOp,

       final SupplierFromVertex<I, V, E, Long> handleHashSupplier,

       final SupplierFromVertex<I, V, E, S> valueSupplier,

       final ConsumerWithVertex<I, V, E, R> reducedValueConsumer) {

     return new Piece<I, V, E, NoMessage, Object>() {

       protected ArrayOfHandles.ArrayOfReducers<S, R> reducers;

       protected BroadcastArrayHandle<R> broadcasts;

       private int getHandleIndex(Vertex<I, V, E> vertex) {

         return (int) Math.abs(handleHashSupplier.get(vertex) % numHandles);

       }

       @Override

       public void registerReducers(

           final CreateReducersApi reduceApi, Object executionStage) {

         reducers = new ArrayOfHandles.ArrayOfReducers<>(

             numHandles,

             new Supplier<ReducerHandle<S, R>>() {

               @Override

               public ReducerHandle<S, R> get() {

                 return reduceApi.createLocalReducer(reduceOp.get());

               }

             });

       }

       @Override

       public VertexSender<I, V, E> getVertexSender(

           BlockWorkerSendApi<I, V, E, NoMessage> workerApi,

           Object executionStage) {

         return new InnerVertexSender() {

           @Override

           public void vertexSend(Vertex<I, V, E> vertex) {

             reducers.get(getHandleIndex(vertex)).reduce(

                 valueSupplier.get(vertex));

           }

         };

       }

       @Override

       public void masterCompute(BlockMasterApi master, Object executionStage) {

         broadcasts = reducers.broadcastValue(master);

       }

       @Override

       public VertexReceiver<I, V, E, NoMessage> getVertexReceiver(

           BlockWorkerReceiveApi<I> workerApi, Object executionStage) {

         final List<R> values = new ArrayList<>();

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

           values.add(broadcasts.get(i).getBroadcast(workerApi));

         }

         return new InnerVertexReceiver() {

           @Override

           public void vertexReceive(

               Vertex<I, V, E> vertex, Iterable<NoMessage> messages) {

             reducedValueConsumer.apply(

                 vertex, values.get(getHandleIndex(vertex)));

           }

         };

       }

       @Override

       public String toString() {

         return name;

       }

     };

   }

   /**

    * Creates Piece that for each vertex, sends message provided by

    * messageSupplier to all targets provided by targetsSupplier.

    * Received messages are then passed to and processed by provided

    * messagesConsumer.

    *

    * If messageSupplier or targetsSupplier returns null, current vertex

    * is not going to send any messages.

    */

   public static

   <I extends WritableComparable, V extends Writable, E extends Writable,

   M extends Writable>

   SendMessagePiece<I, V, E, M> sendMessage(

       String name,

       Class<M> messageClass,

       SupplierFromVertex<I, V, E, M> messageSupplier,

       SupplierFromVertex<I, V, E, Iterator<I>> targetsSupplier,

       ConsumerWithVertex<I, V, E, Iterable<M>> messagesConsumer) {

     return new SendMessagePiece<>(

         name, messageClass, messageSupplier, targetsSupplier, messagesConsumer);

   }

   /**

    * Creates Piece that for each vertex, sends message provided by

    * messageSupplier to all neighbors of current vertex.

    * Received messages are then passed to and processed by provided

    * messagesConsumer.

    *

    * If messageSupplier returns null, current vertex

    * is not going to send any messages.

    */

   public static

   <I extends WritableComparable, V extends Writable, E extends Writable,

   M extends Writable>

   SendMessagePiece<I, V, E, M> sendMessageToNeighbors(

       String name,

       Class<M> messageClass,

       SupplierFromVertex<I, V, E, M> messageSupplier,

       ConsumerWithVertex<I, V, E, Iterable<M>> messagesConsumer) {

     return sendMessage(

         name, messageClass, messageSupplier,

         VertexSuppliers.<I, V, E>vertexNeighborsSupplier(),

         messagesConsumer);

   }

   /**

    * Creates Piece that for each vertex, sends message provided by

    * messageSupplier to all targets provided by targetsSupplier,

    * and uses given messageCombiner to combine messages together.

    * Received combined message is then passed to and processed by provided

    * messageConsumer. (null is passed to it, if vertex received no messages)

    *

    * If messageSupplier or targetsSupplier returns null, current vertex

    * is not going to send any messages.

    */

   public static

   <I extends WritableComparable, V extends Writable, E extends Writable,

   M extends Writable>

   SendMessageWithCombinerPiece<I, V, E, M> sendMessage(

       String name,

       MessageCombiner<? super I, M> messageCombiner,

       SupplierFromVertex<I, V, E, M> messageSupplier,

       SupplierFromVertex<I, V, E, Iterator<I>> targetsSupplier,

       ConsumerWithVertex<I, V, E, M> messagesConsumer) {

     return new SendMessageWithCombinerPiece<>(

         name, messageCombiner,

         messageSupplier, targetsSupplier, messagesConsumer);

   }

   /**

    * Creates Piece that for each vertex, sends message provided by

    * messageSupplier to all neighbors of current vertex,

    * and uses given messageCombiner to combine messages together.

    * Received combined message is then passed to and processed by provided

    * messageConsumer. (null is passed to it, if vertex received no messages)

    *

    * If messageSupplier returns null, current vertex

    * is not going to send any messages.

    */

   public static

   <I extends WritableComparable, V extends Writable, E extends Writable,

   M extends Writable>

   SendMessageWithCombinerPiece<I, V, E, M> sendMessageToNeighbors(

       String name,

       MessageCombiner<? super I, M> messageCombiner,

       SupplierFromVertex<I, V, E, M> messageSupplier,

       ConsumerWithVertex<I, V, E, M> messagesConsumer) {

     return sendMessage(

         name, messageCombiner, messageSupplier,

         VertexSuppliers.<I, V, E>vertexNeighborsSupplier(),

         messagesConsumer);

   }

 }