/*
    * 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;
   }
 }