/*
    * 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.internal;
  
  import java.util.HashSet;
  import java.util.Iterator;
  import java.util.Map;
  import java.util.Map.Entry;
  import java.util.TreeMap;
  
  import org.apache.commons.lang3.time.DurationFormatUtils;
  import org.apache.giraph.block_app.framework.BlockFactory;
  import org.apache.giraph.block_app.framework.BlockUtils;
  import org.apache.giraph.block_app.framework.api.BlockApiHandle;
  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.block.Block;
  import org.apache.giraph.block_app.framework.block.BlockWithApiHandle;
  import org.apache.giraph.block_app.framework.piece.AbstractPiece;
  import org.apache.giraph.conf.GiraphConfiguration;
  import org.apache.giraph.function.Consumer;
  import org.apache.giraph.writable.tuple.IntLongWritable;
  import org.apache.log4j.Logger;
  import com.google.common.base.Preconditions;
  
  /**
   * Block execution logic on master, iterating over Pieces of the
   * application Block, executing master logic, and providing what needs to be
   * executed on the workers.
   *
   * @param <S> Execution stage type
   */
  @SuppressWarnings("rawtypes")
  public class BlockMasterLogic<S> {
    private static final Logger LOG = Logger.getLogger(BlockMasterLogic.class);
  
    private Iterator<AbstractPiece> pieceIterator;
    private PairedPieceAndStage<S> previousPiece;
    private transient BlockMasterApi masterApi;
    private long lastTimestamp = -1;
    private BlockWorkerPieces previousWorkerPieces;
    private boolean computationDone;
    private BlockApiHandle blockApiHandle;
  
    /** Tracks elapsed time on master for each distinct Piece */
    private final TimeStatsPerEvent masterPerPieceTimeStats =
        new TimeStatsPerEvent("master");
    /** Tracks elapsed time on workers for each pair of recieve/send pieces. */
    private final TimeStatsPerEvent workerPerPieceTimeStats =
        new TimeStatsPerEvent("worker");
  
    /**
     * Initialize master logic to execute BlockFactory defined in
     * the configuration.
     */
    public void initialize(
        GiraphConfiguration conf, final BlockMasterApi masterApi) {
      BlockFactory<S> factory = BlockUtils.createBlockFactory(conf);
      initialize(factory.createBlock(conf), factory.createExecutionStage(conf),
          masterApi);
    }
  
    /**
     * Initialize Master Logic to execute given block, starting
     * with given executionStage.
     */
    public void initialize(
        Block executionBlock, S executionStage, final BlockMasterApi masterApi) {
      this.masterApi = masterApi;
      this.computationDone = false;
  
      LOG.info("Executing application - " + executionBlock);
      if (executionBlock instanceof BlockWithApiHandle) {
        blockApiHandle =
          ((BlockWithApiHandle) executionBlock).getBlockApiHandle();
      }
      if (blockApiHandle == null) {
        blockApiHandle = new BlockApiHandle();
      }
      blockApiHandle.setMasterApi(masterApi);
  
      // We register all possible aggregators at the beginning
      executionBlock.forAllPossiblePieces(new Consumer<AbstractPiece>() {
        private final HashSet<AbstractPiece> registeredPieces = new HashSet<>();
       @SuppressWarnings("deprecation")
       @Override
       public void apply(AbstractPiece piece) {
         // no need to register the same piece twice.
         if (registeredPieces.add(piece)) {
           try {
             piece.registerAggregators(masterApi);
           } catch (InstantiationException | IllegalAccessException e) {
             throw new RuntimeException(e);
           }
         }
       }
     });
 
     pieceIterator = executionBlock.iterator();
     // Invariant is that ReceiveWorkerPiece of previousPiece has already been
     // executed and that previousPiece.nextExecutionStage() should be used for
     // iterating. So passing piece as null, and initial state as current state,
     // so that nothing get's executed in first half, and calculateNextState
     // returns initial state.
     previousPiece = new PairedPieceAndStage<>(null, executionStage);
   }
 
   /**
    * Initialize object after deserializing it.
    * BlockMasterApi is not serializable, so it is transient, and set via this
    * method afterwards.
    */
   public void initializeAfterRead(BlockMasterApi masterApi) {
     this.masterApi = masterApi;
   }
 
   /**
    * Executes operations on master (master compute and registering reducers),
    * and calculates next pieces to be exectued on workers.
    *
    * @param superstep Current superstep
    * @return Next BlockWorkerPieces to be executed on workers, or null
    *         if computation should be halted.
    */
   public BlockWorkerPieces<S> computeNext(long superstep) {
     long beforeMaster = System.currentTimeMillis();
     if (lastTimestamp != -1) {
       BlockCounters.setWorkerTimeCounter(
           previousWorkerPieces, superstep - 1,
           beforeMaster - lastTimestamp, masterApi, workerPerPieceTimeStats);
     }
 
     if (previousPiece == null) {
       postApplication();
       return null;
     } else {
       boolean logExecutionStatus =
           BlockUtils.LOG_EXECUTION_STATUS.get(masterApi.getConf());
       if (logExecutionStatus) {
         LOG.info("Master executing " + previousPiece +
             ", in superstep " + superstep);
       }
       previousPiece.masterCompute(masterApi);
       ((BlockOutputHandleAccessor) masterApi).getBlockOutputHandle().
           returnAllWriters();
       long afterMaster = System.currentTimeMillis();
 
       if (previousPiece.getPiece() != null) {
         BlockCounters.setMasterTimeCounter(
             previousPiece, superstep, afterMaster - beforeMaster, masterApi,
             masterPerPieceTimeStats);
       }
 
       PairedPieceAndStage<S> nextPiece;
       if (pieceIterator.hasNext()) {
         nextPiece = new PairedPieceAndStage<S>(
             pieceIterator.next(), previousPiece.nextExecutionStage());
         nextPiece.registerReducers(masterApi);
       } else {
         nextPiece = null;
       }
       BlockCounters.setStageCounters(
           "Master finished stage: ", previousPiece.getExecutionStage(),
           masterApi);
       if (logExecutionStatus) {
         LOG.info(
             "Master passing next " + nextPiece + ", in superstep " + superstep);
       }
 
       // if there is nothing more to compute, no need for additional superstep
       // this can only happen if application uses no pieces.
       BlockWorkerPieces<S> result;
       if (previousPiece.getPiece() == null && nextPiece == null) {
         postApplication();
         result = null;
       } else {
         result = new BlockWorkerPieces<>(
           previousPiece, nextPiece, blockApiHandle);
         if (logExecutionStatus) {
           LOG.info("Master in " + superstep + " superstep passing " +
               result + " to be executed");
         }
       }
 
       previousPiece = nextPiece;
       lastTimestamp = afterMaster;
       previousWorkerPieces = result;
       return result;
     }
   }
 
   /**
    * Clean up any master state, after application has finished.
    */
   private void postApplication() {
     ((BlockOutputHandleAccessor) masterApi).getBlockOutputHandle().
         closeAllWriters();
     Preconditions.checkState(!computationDone);
     computationDone = true;
     IntLongWritable masterTimes = masterPerPieceTimeStats.logTimeSums();
     IntLongWritable workerTimes = workerPerPieceTimeStats.logTimeSums();
     LOG.info("Time split:\n" +
         TimeStatsPerEvent.header() +
         TimeStatsPerEvent.line(
             masterTimes.getLeft().get(),
             100.0 * masterTimes.getRight().get() /
               (masterTimes.getRight().get() + workerTimes.getRight().get()),
             masterTimes.getRight().get(),
             "master") +
         TimeStatsPerEvent.line(
             workerTimes.getLeft().get(),
             100.0 * workerTimes.getRight().get() /
               (masterTimes.getRight().get() + workerTimes.getRight().get()),
             workerTimes.getRight().get(),
             "worker"));
   }
 
   /**
    * Class tracking invocation count and elapsed time for a set of events,
    * each event being having a String name.
    */
   public static class TimeStatsPerEvent {
     private final String groupName;
     private final Map<String, IntLongWritable> keyToCountAndTime =
         new TreeMap<>();
 
     public TimeStatsPerEvent(String groupName) {
       this.groupName = groupName;
     }
 
     public void inc(String name, long millis) {
       IntLongWritable val = keyToCountAndTime.get(name);
       if (val == null) {
         val = new IntLongWritable();
         keyToCountAndTime.put(name, val);
       }
       val.getLeft().set(val.getLeft().get() + 1);
       val.getRight().set(val.getRight().get() + millis);
     }
 
     public IntLongWritable logTimeSums() {
       StringBuilder sb = new StringBuilder("Time sums " + groupName + ":\n");
       sb.append(header());
       long total = 0;
       int count = 0;
       for (Entry<String, IntLongWritable> entry :
             keyToCountAndTime.entrySet()) {
         total += entry.getValue().getRight().get();
         count += entry.getValue().getLeft().get();
       }
 
       for (Entry<String, IntLongWritable> entry :
             keyToCountAndTime.entrySet()) {
         sb.append(line(
             entry.getValue().getLeft().get(),
             (100.0 * entry.getValue().getRight().get()) / total,
             entry.getValue().getRight().get(),
             entry.getKey()));
       }
       LOG.info(sb);
       return new IntLongWritable(count, total);
     }
 
     public static String header() {
       return String.format(
           "%10s%10s%11s   %s%n", "count", "time %", "time", "name");
     }
 
     public static String line(
         int count, double percTime, long time, String name) {
       return String.format("%10d%9.2f%%%11s   %s%n",
           count,
           percTime,
           DurationFormatUtils.formatDuration(time, "HH:mm:ss"),
           name);
     }
   }
 }