/*

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

 import java.util.Iterator;

 import java.util.Spliterator;

 import java.util.Spliterators;

 import java.util.stream.StreamSupport;

 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.block.Block;

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

 import org.apache.giraph.block_app.framework.piece.delegate.FilteringPiece;

 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.striping.StripingUtils;

 import org.apache.giraph.combiner.MessageCombiner;

 import org.apache.giraph.conf.ImmutableClassesGiraphConfiguration;

 import org.apache.giraph.function.Function;

 import org.apache.giraph.function.Predicate;

 import org.apache.giraph.function.primitive.Int2ObjFunction;

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

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

 import org.apache.giraph.graph.Vertex;

 import org.apache.hadoop.io.Writable;

 import org.apache.hadoop.io.WritableComparable;

 import com.google.common.base.Preconditions;

 /**

  * Piece that sends a message provided through messageProducer to given set of

  * neighbors, uses a message combiner and passes them to messagesConsumer.

  *

  * @param <I> Vertex id type

  * @param <V> Vertex value type

  * @param <E> Edge value type

  * @param <M> Message type

  */

 public class SendMessageWithCombinerPiece<I extends WritableComparable,

     V extends Writable, E extends Writable, M extends Writable>

     extends Piece<I, V, E, M, Object> {

   private final String name;

   private final MessageCombiner<? super I, M> messageCombiner;

   private final SupplierFromVertex<I, V, E, M> messageSupplier;

   private final SupplierFromVertex<I, V, E, Iterator<I>> targetsSupplier;

   private final ConsumerWithVertex<I, V, E, M> messagesConsumer;

   public SendMessageWithCombinerPiece(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) {

     Preconditions.checkNotNull(messageCombiner);

     this.name = name;

     this.messageCombiner = messageCombiner;

     this.messageSupplier = messageSupplier;

     this.targetsSupplier = targetsSupplier;

     this.messagesConsumer = messagesConsumer;

   }

   /**

    * Stripe message sending computation across multiple stripes, in

    * each stripe only part of the vertices will receive messages.

    *

    * @param stripes Number of stripes

    * @param stripeSupplier Stripe supplier function, if IDs are Longs, you can

    *                       use StripingUtils::fastHashStripingPredicate

    * @return Resulting block

    */

   public Block stripeByReceiver(

       int stripes,

       Int2ObjFunction<Int2ObjFunction<Predicate<I>>> stripeSupplier) {

     return StripingUtils.generateStripedBlock(

         stripes,

         new Function<Predicate<I>, Block>() {

           @Override

           public Block apply(final Predicate<I> stripePredicate) {

             return FilteringPiece.createReceiveFiltering(

                 new SupplierFromVertex<I, V, E, Boolean>() {

                   @Override

                   public Boolean get(Vertex<I, V, E> vertex) {

                     return stripePredicate.apply(vertex.getId());

                   }

                 },

                 new SendMessageWithCombinerPiece<>(

                   name,

                   messageCombiner,

                   messageSupplier,

                   new SupplierFromVertex<I, V, E, Iterator<I>>() {

                     @Override

                     public Iterator<I> get(Vertex<I, V, E> vertex) {

                       return StreamSupport.stream(

                         Spliterators.spliteratorUnknownSize(

                           targetsSupplier.get(vertex), Spliterator.ORDERED),

                         false).filter(stripePredicate::apply).iterator();

                     }

                   },

                   messagesConsumer));

           }

         },

         stripeSupplier);

   }

   @Override

   public VertexSender<I, V, E> getVertexSender(

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

       Object executionStage) {

     return new InnerVertexSender() {

       @Override

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

         Iterator<I> targets = targetsSupplier.get(vertex);

         M message = messageSupplier.get(vertex);

         if (message != null && targets != null && targets.hasNext()) {

           workerApi.sendMessageToMultipleEdges(targets, message);

         }

       }

     };

   }

   @Override

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

       BlockWorkerReceiveApi<I> workerApi,

       Object executionStage) {

     return new InnerVertexReceiver() {

       @Override

       public void vertexReceive(Vertex<I, V, E> vertex, Iterable<M> messages) {

         Iterator<M> iter = messages.iterator();

         M combinedMessage = null;

         if (iter.hasNext()) {

           combinedMessage = iter.next();

           // When message combiner is used, there is never more then one message

           Preconditions.checkArgument(!iter.hasNext());

         }

         messagesConsumer.apply(vertex, combinedMessage);

       }

     };

   }

   @Override

   public MessageCombiner<? super I, M> getMessageCombiner(

       ImmutableClassesGiraphConfiguration conf) {

     return messageCombiner;

   }

   @Override

   public String toString() {

     return name;

   }

 }