/*
    * 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 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.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;
  import com.google.common.collect.Iterators;
  
  /**
   * Piece that sends a message provided through messageProducer to given set of
   * neighbors, and passes them to messagesConsumer.
   *
   * @param <I> Vertex id type
   * @param <V> Vertex value type
   * @param <E> Edge value type
   * @param <M> Message type
   */
  @SuppressWarnings("rawtypes")
  public class SendMessagePiece<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 Class<M> messageClass;
    private final SupplierFromVertex<I, V, E, M> messageSupplier;
    private final SupplierFromVertex<I, V, E, Iterator<I>> targetsSupplier;
    private final ConsumerWithVertex<I, V, E, Iterable<M>> messagesConsumer;
  
    public SendMessagePiece(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) {
      Preconditions.checkNotNull(messageClass);
      this.name = name;
      this.messageClass = messageClass;
      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 SendMessagePiece<>(
                    name,
                    messageClass,
                   messageSupplier,
                   new SupplierFromVertex<I, V, E, Iterator<I>>() {
                     @Override
                     public Iterator<I> get(Vertex<I, V, E> vertex) {
                       return Iterators.filter(
                           targetsSupplier.get(vertex),
                           new com.google.common.base.Predicate<I>() {
                             @Override
                             public boolean apply(I targetId) {
                               return stripePredicate.apply(targetId);
                             }
                           });
                     }
                   },
                   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) {
         messagesConsumer.apply(vertex, messages);
       }
     };
   }
 
   @Override
   public Class<M> getMessageClass() {
     return messageClass;
   }
 
   @Override
   public String toString() {
     return name;
   }
 }