/*

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

 import java.io.DataInput;

 import java.io.DataOutput;

 import java.io.IOException;

 import java.util.HashMap;

 import java.util.Map.Entry;

 import org.apache.giraph.aggregators.Aggregator;

 import org.apache.giraph.comm.aggregators.AggregatorUtils;

 import org.apache.giraph.conf.ImmutableClassesGiraphConfiguration;

 import org.apache.giraph.utils.MasterLoggingAggregator;

 import org.apache.giraph.utils.WritableUtils;

 import org.apache.hadoop.io.Writable;

 import org.apache.log4j.Logger;

 import com.google.common.base.Preconditions;

 /**

  * Class that translates aggregator handling on the master to

  * reduce and broadcast operations supported by the MasterAggregatorHandler.

  */

 public class AggregatorToGlobalCommTranslation

     implements MasterAggregatorUsage, Writable {

   /** Class logger */

   private static final Logger LOG =

       Logger.getLogger(AggregatorToGlobalCommTranslation.class);

   /** Class providing reduce and broadcast interface to use */

   private final MasterGlobalCommUsage globalComm;

   /** List of registered aggregators */

   private final HashMap<String, AggregatorWrapper<Writable>>

   registeredAggregators = new HashMap<>();

   /**

    * List of init aggregator values, in case someone tries to

    * access aggregator immediatelly after registering it.

    *

    * Instead of simply returning value, we need to store it during

    * that superstep, so consecutive calls will return identical object,

    * which they can modify.

    */

   private final HashMap<String, Writable>

   initAggregatorValues = new HashMap<>();

   /** Conf */

   private final ImmutableClassesGiraphConfiguration<?, ?, ?> conf;

   /**

    * Constructor

    * @param conf Configuration

    * @param globalComm Global communication interface

    */

   public AggregatorToGlobalCommTranslation(

       ImmutableClassesGiraphConfiguration<?, ?, ?> conf,

       MasterGlobalCommUsage globalComm) {

     this.conf = conf;

     this.globalComm = globalComm;

     MasterLoggingAggregator.registerAggregator(this, conf);

   }

   @Override

   public <A extends Writable> A getAggregatedValue(String name) {

     AggregatorWrapper<Writable> agg = registeredAggregators.get(name);

     if (agg == null) {

       LOG.warn("getAggregatedValue: " +

         AggregatorUtils.getUnregisteredAggregatorMessage(name,

             registeredAggregators.size() != 0, conf));

       // to make sure we are not accessing reducer of the same name.

       return null;

     }

     A value = globalComm.getReduced(name);

     if (value == null) {

       value = (A) initAggregatorValues.get(name);

     }

     if (value == null) {

       value = (A) agg.getReduceOp().createInitialValue();

       initAggregatorValues.put(name, value);

     }

     Preconditions.checkState(value != null);

     return value;

   }

   @Override

   public <A extends Writable> void setAggregatedValue(String name, A value) {

     AggregatorWrapper<Writable> aggregator = registeredAggregators.get(name);

     if (aggregator == null) {

       throw new IllegalArgumentException("setAggregatedValue: "  +

           AggregatorUtils.getUnregisteredAggregatorMessage(name,

               registeredAggregators.size() != 0, conf));

     }

     aggregator.setCurrentValue(value);

   }

   /**

    * Called after master compute, to do aggregator->reduce/broadcast

    * translation

    */

   public void postMasterCompute() {

     // broadcast what master set, or if it didn't broadcast reduced value

     // register reduce with the same value

     for (Entry<String, AggregatorWrapper<Writable>> entry :

         registeredAggregators.entrySet()) {

       Writable value = entry.getValue().getCurrentValue();

       if (value == null) {

         value = globalComm.getReduced(entry.getKey());

       }

       Preconditions.checkState(value != null);

       globalComm.broadcast(entry.getKey(), new AggregatorBroadcast<>(

           entry.getValue().getReduceOp().getAggregatorClass(), value));

       // Always register clean instance of reduceOp, not to conflict with

       // reduceOp from previous superstep.

       AggregatorReduceOperation<Writable> cleanReduceOp =

           entry.getValue().createReduceOp();

       if (entry.getValue().isPersistent()) {

         globalComm.registerReducer(

             entry.getKey(), cleanReduceOp, value);

       } else {

         globalComm.registerReducer(

             entry.getKey(), cleanReduceOp);

       }

       entry.getValue().setCurrentValue(null);

     }

     initAggregatorValues.clear();

   }

   /** Prepare before calling master compute */

   public void prepareSuperstep() {

     MasterLoggingAggregator.logAggregatedValue(this, conf);

   }

   @Override

   public <A extends Writable> boolean registerAggregator(String name,

       Class<? extends Aggregator<A>> aggregatorClass) throws

       InstantiationException, IllegalAccessException {

     registerAggregator(name, aggregatorClass, false);

     return true;

   }

   @Override

   public <A extends Writable> boolean registerPersistentAggregator(String name,

       Class<? extends Aggregator<A>> aggregatorClass) throws

       InstantiationException, IllegalAccessException {

     registerAggregator(name, aggregatorClass, true);

     return true;

   }

   @Override

   public void write(DataOutput out) throws IOException {

     out.writeInt(registeredAggregators.size());

     for (Entry<String, AggregatorWrapper<Writable>> entry :

         registeredAggregators.entrySet()) {

       out.writeUTF(entry.getKey());

       entry.getValue().write(out);

     }

   }

   @Override

   public void readFields(DataInput in) throws IOException {

     registeredAggregators.clear();

     int numAggregators = in.readInt();

     for (int i = 0; i < numAggregators; i++) {

       String name = in.readUTF();

       AggregatorWrapper<Writable> agg = new AggregatorWrapper<>();

       agg.readFields(in);

       registeredAggregators.put(name, agg);

     }

     initAggregatorValues.clear();

   }

   /**

    * Helper function for registering aggregators.

    *

    * @param name            Name of the aggregator

    * @param aggregatorClass Aggregator class

    * @param persistent      Whether aggregator is persistent or not

    * @param <A>             Aggregated value type

    * @return Newly registered aggregator or aggregator which was previously

    *         created with selected name, if any

    */

   private <A extends Writable> AggregatorWrapper<A> registerAggregator

   (String name, Class<? extends Aggregator<A>> aggregatorClass,

       boolean persistent) throws InstantiationException,

       IllegalAccessException {

     AggregatorWrapper<A> aggregatorWrapper =

         (AggregatorWrapper<A>) registeredAggregators.get(name);

     if (aggregatorWrapper == null) {

       aggregatorWrapper =

           new AggregatorWrapper<A>(aggregatorClass, persistent);

       // postMasterCompute uses previously reduced value to broadcast,

       // unless current value is set. After aggregator is registered,

       // there was no previously reduced value, so set current value

       // to default to avoid calling getReduced() on unregistered reducer.

       // (which logs unnecessary warnings)

       aggregatorWrapper.setCurrentValue(

           aggregatorWrapper.getReduceOp().createInitialValue());

       registeredAggregators.put(

           name, (AggregatorWrapper<Writable>) aggregatorWrapper);

     }

     return aggregatorWrapper;

   }

   /**

    * Object holding all needed data related to single Aggregator

    * @param <A> Aggregated value type

    */

   private class AggregatorWrapper<A extends Writable>

       implements Writable {

     /** False iff aggregator should be reset at the end of each super step */

     private boolean persistent;

     /** Translation of aggregator to reduce operations */

     private AggregatorReduceOperation<A> reduceOp;

     /** Current value, set by master manually */

     private A currentValue;

     /** Constructor */

     public AggregatorWrapper() {

     }

     /**

      * Constructor

      * @param aggregatorClass Aggregator class

      * @param persistent Is persistent

      */

     public AggregatorWrapper(

         Class<? extends Aggregator<A>> aggregatorClass,

         boolean persistent) {

       this.persistent = persistent;

       this.reduceOp = new AggregatorReduceOperation<>(aggregatorClass, conf);

     }

     public AggregatorReduceOperation<A> getReduceOp() {

       return reduceOp;

     }

     /**

      * Create a fresh instance of AggregatorReduceOperation

      * @return fresh instance of AggregatorReduceOperation

      */

     public AggregatorReduceOperation<A> createReduceOp() {

       return reduceOp.createCopy();

     }

     public A getCurrentValue() {

       return currentValue;

     }

     public void setCurrentValue(A currentValue) {

       this.currentValue = currentValue;

     }

     public boolean isPersistent() {

       return persistent;

     }

     @Override

     public void write(DataOutput out) throws IOException {

       out.writeBoolean(persistent);

       reduceOp.write(out);

       Preconditions.checkState(currentValue == null, "AggregatorWrapper " +

           "shouldn't have value at the end of the superstep");

     }

     @Override

     public void readFields(DataInput in) throws IOException {

       persistent = in.readBoolean();

       reduceOp = WritableUtils.createWritable(

           AggregatorReduceOperation.class, conf);

       reduceOp.readFields(in);

       currentValue = null;

     }

   }

 }