/*
    * 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.worker;
  
  import java.io.IOException;
  import java.util.Map;
  import java.util.Map.Entry;
  import java.util.Set;
  
  import org.apache.giraph.bsp.CentralizedServiceWorker;
  import org.apache.giraph.comm.GlobalCommType;
  import org.apache.giraph.comm.aggregators.AggregatorUtils;
  import org.apache.giraph.comm.aggregators.AllAggregatorServerData;
  import org.apache.giraph.comm.aggregators.GlobalCommValueOutputStream;
  import org.apache.giraph.comm.aggregators.OwnerAggregatorServerData;
  import org.apache.giraph.comm.aggregators.WorkerAggregatorRequestProcessor;
  import org.apache.giraph.conf.ImmutableClassesGiraphConfiguration;
  import org.apache.giraph.reducers.ReduceOperation;
  import org.apache.giraph.reducers.Reducer;
  import org.apache.giraph.utils.UnsafeByteArrayOutputStream;
  import org.apache.giraph.utils.UnsafeReusableByteArrayInput;
  import org.apache.giraph.utils.WritableUtils;
  import org.apache.hadoop.io.Writable;
  import org.apache.hadoop.util.Progressable;
  import org.apache.log4j.Logger;
  
  import com.google.common.collect.Maps;
  import com.google.common.collect.Sets;
  
  /** Handler for reduce/broadcast on the workers */
  public class WorkerAggregatorHandler implements WorkerThreadGlobalCommUsage {
    /** Class logger */
    private static final Logger LOG =
        Logger.getLogger(WorkerAggregatorHandler.class);
    /** Map of broadcasted values */
    private final Map<String, Writable> broadcastedMap =
        Maps.newHashMap();
    /** Map of reducers currently being reduced */
    private final Map<String, Reducer<Object, Writable>> reducerMap =
        Maps.newHashMap();
  
    /** Service worker */
    private final CentralizedServiceWorker<?, ?, ?> serviceWorker;
    /** Progressable for reporting progress */
    private final Progressable progressable;
    /** How big a single aggregator request can be */
    private final int maxBytesPerAggregatorRequest;
    /** Giraph configuration */
    private final ImmutableClassesGiraphConfiguration conf;
  
    /**
     * Constructor
     *
     * @param serviceWorker Service worker
     * @param conf          Giraph configuration
     * @param progressable  Progressable for reporting progress
     */
    public WorkerAggregatorHandler(
        CentralizedServiceWorker<?, ?, ?> serviceWorker,
        ImmutableClassesGiraphConfiguration conf,
        Progressable progressable) {
      this.serviceWorker = serviceWorker;
      this.progressable = progressable;
      this.conf = conf;
      maxBytesPerAggregatorRequest = conf.getInt(
          AggregatorUtils.MAX_BYTES_PER_AGGREGATOR_REQUEST,
          AggregatorUtils.MAX_BYTES_PER_AGGREGATOR_REQUEST_DEFAULT);
    }
  
    @Override
    public <B extends Writable> B getBroadcast(String name) {
      B value = (B) broadcastedMap.get(name);
      if (value == null) {
        LOG.warn("getBroadcast: " +
            AggregatorUtils.getUnregisteredBroadcastMessage(name,
                broadcastedMap.size() != 0, conf));
      }
      return value;
    }
  
    @Override
    public void reduce(String name, Object value) {
      Reducer<Object, Writable> reducer = reducerMap.get(name);
      if (reducer != null) {
       progressable.progress();
       synchronized (reducer) {
         reducer.reduce(value);
       }
     } else {
       throw new IllegalStateException("reduce: " +
           AggregatorUtils.getUnregisteredReducerMessage(name,
               reducerMap.size() != 0, conf));
     }
   }
 
   /**
    * Combine partially reduced value into currently reduced value.
    * @param name Name of the reducer
    * @param valueToReduce Partial value to reduce
    */
   @Override
   public void reduceMerge(String name, Writable valueToReduce) {
     Reducer<Object, Writable> reducer = reducerMap.get(name);
     if (reducer != null) {
       progressable.progress();
       synchronized (reducer) {
         reducer.reduceMerge(valueToReduce);
       }
     } else {
       throw new IllegalStateException("reduce: " +
           AggregatorUtils.getUnregisteredReducerMessage(name,
               reducerMap.size() != 0, conf));
     }
   }
 
   /**
    * Prepare aggregators for current superstep
    *
    * @param requestProcessor Request processor for aggregators
    */
   public void prepareSuperstep(
       WorkerAggregatorRequestProcessor requestProcessor) {
     broadcastedMap.clear();
     reducerMap.clear();
 
     if (LOG.isDebugEnabled()) {
       LOG.debug("prepareSuperstep: Start preparing aggregators");
     }
     AllAggregatorServerData allGlobalCommData =
         serviceWorker.getServerData().getAllAggregatorData();
     // Wait for my aggregators
     Iterable<byte[]> dataToDistribute =
         allGlobalCommData.getDataFromMasterWhenReady(
             serviceWorker.getMasterInfo());
     try {
       // Distribute my aggregators
       requestProcessor.distributeReducedValues(dataToDistribute);
     } catch (IOException e) {
       throw new IllegalStateException("prepareSuperstep: " +
           "IOException occurred while trying to distribute aggregators", e);
     }
     // Wait for all other aggregators and store them
     allGlobalCommData.fillNextSuperstepMapsWhenReady(
         getOtherWorkerIdsSet(), broadcastedMap,
         reducerMap);
     if (LOG.isDebugEnabled()) {
       LOG.debug("prepareSuperstep: Aggregators prepared");
     }
   }
 
   /**
    * Send aggregators to their owners and in the end to the master
    *
    * @param requestProcessor Request processor for aggregators
    */
   public void finishSuperstep(
       WorkerAggregatorRequestProcessor requestProcessor) {
     if (LOG.isInfoEnabled()) {
       LOG.info("finishSuperstep: Start gathering aggregators, " +
           "workers will send their aggregated values " +
           "once they are done with superstep computation");
     }
     OwnerAggregatorServerData ownerGlobalCommData =
         serviceWorker.getServerData().getOwnerAggregatorData();
     // First send partial aggregated values to their owners and determine
     // which aggregators belong to this worker
     for (Map.Entry<String, Reducer<Object, Writable>> entry :
         reducerMap.entrySet()) {
       try {
         boolean sent = requestProcessor.sendReducedValue(entry.getKey(),
             entry.getValue().getCurrentValue());
         if (!sent) {
           // If it's my aggregator, add it directly
           ownerGlobalCommData.reduce(entry.getKey(),
               entry.getValue().getCurrentValue());
         }
       } catch (IOException e) {
         throw new IllegalStateException("finishSuperstep: " +
             "IOException occurred while sending aggregator " +
             entry.getKey() + " to its owner", e);
       }
       progressable.progress();
     }
     try {
       // Flush
       requestProcessor.flush();
     } catch (IOException e) {
       throw new IllegalStateException("finishSuperstep: " +
           "IOException occurred while sending aggregators to owners", e);
     }
 
     // Wait to receive partial aggregated values from all other workers
     Iterable<Map.Entry<String, Writable>> myReducedValues =
         ownerGlobalCommData.getMyReducedValuesWhenReady(
             getOtherWorkerIdsSet());
 
     // Send final aggregated values to master
     GlobalCommValueOutputStream globalOutput =
         new GlobalCommValueOutputStream(false);
     for (Map.Entry<String, Writable> entry : myReducedValues) {
       try {
         int currentSize = globalOutput.addValue(entry.getKey(),
             GlobalCommType.REDUCED_VALUE,
             entry.getValue());
         if (currentSize > maxBytesPerAggregatorRequest) {
           requestProcessor.sendReducedValuesToMaster(
               globalOutput.flush());
         }
         progressable.progress();
       } catch (IOException e) {
         throw new IllegalStateException("finishSuperstep: " +
             "IOException occurred while writing aggregator " +
             entry.getKey(), e);
       }
     }
     try {
       requestProcessor.sendReducedValuesToMaster(globalOutput.flush());
     } catch (IOException e) {
       throw new IllegalStateException("finishSuperstep: " +
           "IOException occured while sending aggregators to master", e);
     }
     // Wait for master to receive aggregated values before proceeding
     serviceWorker.getWorkerClient().waitAllRequests();
 
     ownerGlobalCommData.reset();
     if (LOG.isDebugEnabled()) {
       LOG.debug("finishSuperstep: Aggregators finished");
     }
   }
 
   /**
    * Create new aggregator usage which will be used by one of the compute
    * threads.
    *
    * @return New aggregator usage
    */
   public WorkerThreadGlobalCommUsage newThreadAggregatorUsage() {
     if (AggregatorUtils.useThreadLocalAggregators(conf)) {
       return new ThreadLocalWorkerGlobalCommUsage();
     } else {
       return this;
     }
   }
 
   @Override
   public void finishThreadComputation() {
     // If we don't use thread-local aggregators, all the aggregated values
     // are already in this object
   }
 
   /**
    * Get set of all worker task ids except the current one
    *
    * @return Set of all other worker task ids
    */
   public Set<Integer> getOtherWorkerIdsSet() {
     Set<Integer> otherWorkers = Sets.newHashSetWithExpectedSize(
         serviceWorker.getWorkerInfoList().size());
     for (WorkerInfo workerInfo : serviceWorker.getWorkerInfoList()) {
       if (workerInfo.getTaskId() != serviceWorker.getWorkerInfo().getTaskId()) {
         otherWorkers.add(workerInfo.getTaskId());
       }
     }
     return otherWorkers;
   }
 
   /**
   * Not thread-safe implementation of {@link WorkerThreadGlobalCommUsage}.
   * We can use one instance of this object per thread to prevent
   * synchronizing on each aggregate() call. In the end of superstep,
   * values from each of these will be aggregated back to {@link
   * WorkerThreadGlobalCommUsage}
   */
   public class ThreadLocalWorkerGlobalCommUsage
     implements WorkerThreadGlobalCommUsage {
     /** Thread-local reducer map */
     private final Map<String, Reducer<Object, Writable>> threadReducerMap;
 
     /**
     * Constructor
     *
     * Creates new instances of all reducers from
     * {@link WorkerAggregatorHandler}
     */
     public ThreadLocalWorkerGlobalCommUsage() {
       threadReducerMap = Maps.newHashMapWithExpectedSize(
           WorkerAggregatorHandler.this.reducerMap.size());
 
       UnsafeByteArrayOutputStream out = new UnsafeByteArrayOutputStream();
       UnsafeReusableByteArrayInput in = new UnsafeReusableByteArrayInput();
 
       for (Entry<String, Reducer<Object, Writable>> entry :
           reducerMap.entrySet()) {
         ReduceOperation<Object, Writable> globalReduceOp =
             entry.getValue().getReduceOp();
 
         ReduceOperation<Object, Writable> threadLocalCopy =
             WritableUtils.createCopy(out, in, globalReduceOp, conf);
 
         threadReducerMap.put(entry.getKey(), new Reducer<>(threadLocalCopy));
       }
     }
 
     @Override
     public void reduce(String name, Object value) {
       Reducer<Object, Writable> reducer = threadReducerMap.get(name);
       if (reducer != null) {
         progressable.progress();
         reducer.reduce(value);
       } else {
         throw new IllegalStateException("reduce: " +
             AggregatorUtils.getUnregisteredAggregatorMessage(name,
                 threadReducerMap.size() != 0, conf));
       }
     }
 
     @Override
     public void reduceMerge(String name, Writable value) {
       Reducer<Object, Writable> reducer = threadReducerMap.get(name);
       if (reducer != null) {
         progressable.progress();
         reducer.reduceMerge(value);
       } else {
         throw new IllegalStateException("reduceMerge: " +
             AggregatorUtils.getUnregisteredAggregatorMessage(name,
                 threadReducerMap.size() != 0, conf));
       }
     }
 
     @Override
     public <B extends Writable> B getBroadcast(String name) {
       return WorkerAggregatorHandler.this.getBroadcast(name);
     }
 
     @Override
     public void finishThreadComputation() {
       // Aggregate the values this thread's vertices provided back to
       // WorkerAggregatorHandler
       for (Entry<String, Reducer<Object, Writable>> entry :
           threadReducerMap.entrySet()) {
         WorkerAggregatorHandler.this.reduceMerge(entry.getKey(),
             entry.getValue().getCurrentValue());
       }
     }
   }
 
 }