#ifndef __AVERAGE_H__
#define __AVERAGE_H__

#include <map>
#include <vector>

#include <A++.h>

namespace AverageVariables{
  /**
   * Contains enum of variable ID's used with the Average class.
   * 
   * \Warning Do *not* use the integers directly in this enum!
   * Instead, use the variable.  That is, don't use "5" for BZ.
   * Instead, use "AverageVariables::BZ".
   *
   * \Warning: Average::getVariableIds refers directly to these enum
   * values.  Change these & make sure to update the getVariableIds
   * method!!!
   */
  enum Variables{
    BX = 0,  BY = 1,  BZ = 2,
    BI = 3,  BJ = 4,  BK = 5,
    RHO = 6,
    CSOUND = 7,
    PRESSURE = 8,
    EI = 9,  EJ = 10,  EK = 11,
    VX = 12, VY = 13, VZ = 14,
    V = 15
  };
  static const int nVariables = 16;
}

/*******************************************************************************
 * Implements time averaging of MHD variables.  Algorithm developed by
 * John Lyon.  Implementation by Pete Schmitt.
 *
 * Compatible with fundamental A++/P++ typeArray's, typically
 * floatArray or doubleArray.
 *
 * Description of the algorithm:
 *
 * Let's assume we want a time average of a quantity \f[$q$\f] (which
 * can be density, pressure, velocity, etc).  Let the LFM timestep at
 * step \f[$N$\f] be defined as \f[$dt_N$\f]. Let's say we have done
 * \f[$N$\f] steps (since the last call to the RCM) in accumulating,
 * the timestep weighed average is then
 * 
 * \f[\begin{equation}
 * <q>_N=\frac{\sum_{i=1}^N q_i dt_i}{\sum_{i=1}^N dt_i}=\frac{\sum_{i=1}^N q_i dt_i}{T_N}
 * \end{equation}\f]
 * 
 * where \f[$T_N=\sum_{i=1}^N dt_i$\f] and is the time elapsed. Now we
 * are at timestep \f[$N+1$\f] and want to compute a new average that
 * takes into account (1) and the new value, so we want to compute
 * 
 * \f[\begin{equation}\setcounter{equation}{2}
 * <q>_{N+1}=\frac{\sum_{i=1}^{N+1} q_i dt_i}{T_{N+1}}
 * \end{equation}\f]
 * 
 * Since we don't have the information for all the steps but we do
 * have (1) we can rewrite (2) as
 * 
 * \f[\begin{equation}\setcounter{equation}{3}
 * <q>_{N+1}=\frac{\sum_{i=1}^N q_i dt_i + q_{N+1}dt_{N+1}}{T_N+dt_{N+1}}
 * =\frac{<q>_N (T_N-dt_{N+1})}{T_{N+1}}+\frac{q_{N+1}dt_{N+1}}{T_{N+1}}
 * \end{equation}\f]
 * 
 * This also saves use from redoing the averaging at each timestep. 
 * 
 * In the code 
 * 
 * \f[\begin{eqnarray*}
 * \verb q_ac  &\equiv & <q>_{N+1} \\
 * \verb tau_rcm &\equiv & T_{N+1} \\
 * \verb saved &\equiv & T_{N+1}-dt \\
 * \verb delta &\equiv & \frac{dt}{T_{N+1}} \\
 * \end{eqnarray*}\f]
 * 
 * where \f[$dt$\f] is the current timestep (\f[$dt_{N+1}$\f]).
 ******************************************************************************/
template <class typeArray>
class Average{
public:
  /// Default constructor.
  Average();

  /// Overloaded Constructor (must be called before you can use this class).
  /**
   * Allocates variables for averaging.
   */
  Average(const int &nip1, const int &njp1, const int &nkp1,
	  const Partitioning_Type &mhdGhostPartitioning,
	  const std::vector<int> variableIds);
  
  /// Overload assignment when default constructor is called before the overloaded constructor.
  /**
   *  P++ Partitioning_Type and typeArrays use dynamically allocated
   *  memory.  Without this function, allocating an Average object on
   *  the stack by first calling the default constructor, then
   *  assignment:
   *     Average abc;
   *     Average abc = Average(nip1,njp1,nkp1, mhdGhost, variables);
   *  generates run-time errors when calling variables[i].partition(mhdGhost):
   *
   *   NODE 0: ERROR in floatArray::Test_Conformability ( const floatArray & X )
   *   NODE 0: Operands have different width ghost boundaries -- sorry -- not yet supported!
   *   NODE 0: Internal Ghost Boundary widths: Lhs (0, 0, 0, 0) Rhs (4, 4, 4, 0)
   *   NODE 0: Exiting ...
   *   Exiting P++ program from inside of APP_ABORT()! (Calling system abort function ...)
   *   Exiting P++ Virtual Machine!
   *
   * I believe the source of this problem is that the standard
   * non-overloaded C++ Average::operator= does a shallow copy of P++
   * objects.  This may not increment the P++ reference counting, and
   * the Partitioning_Type may be inappropriately deallocated.  In any
   * case, let's explicitly copy the P++ typeArray and
   * Partitioning_Type.  Find more details at these tutorials:
   *
   * http://www.learncpp.com/cpp-tutorial/911-the-copy-constructor-and-overloading-the-assignment-operator/
   * http://www.learncpp.com/cpp-tutorial/912-shallow-vs-deep-copying/
   */
  Average<typeArray> &operator=(const Average<typeArray> &copy);

  /// Set the timestep that will be used in the running average
  void setDeltaTime(const double &dt);

  /// Increment the average for this variable
  /**
   * \note Must set the timestep using Average::setDeltaTime(...)
   * prior to incrementing a variable.
   *
   * \param variable: Instantaneous value of variable to be summed in
   *                  the average.
   * \param variableId: Must be one of the variables defined in the
   *                    AverageVariables namespace.
   */
  void increment(const typeArray &variable, const int &variableId);

  /// Obtain a map of allocated variable indices & names
  /*
   * Output:  map containing
   *    Integer key: variable Id.
   *    String value: name of the variable
   */
  map<int, string> getVariableIds() const;

  /// Obtain the averaged value.  
  /** 
   * variableId should be chosen via the enum in the AverageVariables
   * namespace.
   */
  const typeArray & getVariable(const int &variableId) const;

  /// Set an averaged value.  Usefule for restarting the code
  /** 
   * variableId should be chosen via the enum in the AverageVariables
   * namespace.
   */
  void setVariable(const typeArray &variable, const int &variableId);

  ///  Reset averaging arrays to zero
  /**
   * Loops over each variable & calls Average::reset(varId).  Set to
   * zero at initialization and/or after the time interval you wish to
   * average.
   *
   * FIXME: Each A++/P++ zeroing is O(n^3)!  We may be able to
   * optimize by a triple for-loop setting each variable to zero.
   */
  void reset();
  
private:

  /// Reset a particular averaging array to zero
  /**
   * Obtain list of variable ID's in enum in AverageVariables
   * namespace.
   */
  void reset(const int &variableId);

  /// List of variables to be averaged
  std::vector<int> variableIds;

  /// Size of timestep being averaged
  double dt;
  double tau;

  /// Variables used in averaging
  //@{
  int nip1, njp1, nkp1;
  Partitioning_Type mhdPartitioning;

  /// Array of typeArray's used to store data.
  /**
   * We cannot simply create a std::vector<typeArray> thanks to
   * A++/P++.  The compiler throws an error in the STL constructor
   * that looks like:
   * 
   *     no instance of overloaded "floatArray::operator new" matches 
   *     the argument list argument types are: 
   *         (unsigned long, floatArray *)
   *
   * Likewise, dynamically creating an array of typeArrays does not
   * work (I get segfaults while manipulating the arrays).  I think
   * these segfaults are related to A++/P++ dynamically allocating
   * memory when typeArray::redim is called.  typeArrays become
   * non-continugous.
   * 
   * Thanks again, A++/P++ developers!
   */
  typeArray variables[AverageVariables::nVariables];
  //typeArray Bx, By, Bz;
  //typeArray Bi, Bj, Bk;
  //typeArray Rho;
  //typeArray Csound;
  //typeArray Ei, Ej, Ek;
  //typeArray Vx, Vy, Vz;  
  //@}
};

template <class typeArray>
Average<typeArray>::Average()
{
#ifdef DEBUG_MODE_ON
  std::cerr << __FILE__ <<  " (L " << __LINE__ << ") Inside Function  " << __FUNCTION__ << "(...)" << endl;
#endif
}

/******************************************************************************/

template <class typeArray>
Average<typeArray>::Average(const int &nip1, const int &njp1, const int &nkp1, 
			    const Partitioning_Type &mhdGhostPartitioning,
			    const std::vector<int> variableIds)
{
#ifdef DEBUG_MODE_ON
  std::cerr << __FILE__ <<  " (L " << __LINE__ << ") Inside Function  " << __FUNCTION__ << "(...)" << endl;
#endif

  this->variableIds = variableIds;
  this->dt = 0.0;

  this->nip1 = nip1;
  this->njp1 = njp1;
  this->nkp1 = nkp1;

  mhdPartitioning = mhdGhostPartitioning;
  
  // Configure arrays used to store results of averaging.  See
  // comments in class definition for an explanation of why we
  // populate worker arrays in this manner.
  vector<int>::iterator it;
  for ( it=this->variableIds.begin() ; it < this->variableIds.end(); it++ ){
    if ( ((*it) < 0) || (*it >= AverageVariables::nVariables) ){
      std::cerr << "ERROR in " << __FILE__ <<  "(L" << __LINE__ << ")" << __FUNCTION__ << "(...): " 
		<< "Did not understand variable with index \"" << *it << "\".  "
		<< "Did you try using the variables defined in the \"AverageVariables\" namespace?"
		<< endl;
      Communication_Manager::Sync();
      Optimization_Manager::Exit_Virtual_Machine();
    }
    variables[*it].redim(nip1, njp1, nkp1);
    variables[*it].partition(mhdPartitioning);
    variables[*it].setBase(1);
  }

  // Initialize average arrays to zero.
  this->reset();
}

/******************************************************************************/

template <class typeArray>
Average<typeArray> & Average<typeArray>::operator=(const Average<typeArray> &copy)
{
#ifdef DEBUG_MODE_ON
  std::cerr << __FILE__ <<  " (L " << __LINE__ << ") Inside Function  " << __FUNCTION__ << "(...)" << endl;
#endif
  // Check for self-assignment
  if (this == &copy)
    return *this;

  dt = copy.dt;
  tau = copy.tau;
  nip1 = copy.nip1;
  njp1 = copy.njp1;
  nkp1 = copy.nkp1;
  // forces deep-copy of Partitioning_Type
  mhdPartitioning = copy.mhdPartitioning;
  variableIds = copy.variableIds;

  vector<int>::iterator it;
  for ( it=this->variableIds.begin() ; it < this->variableIds.end(); it++ ){
    variables[*it].redim(nip1, njp1, nkp1);
    variables[*it].partition(mhdPartitioning);
    variables[*it].setBase(1);
 
    // forces deep-copy of variable[*it]
   variables[*it] = copy.variables[*it];
  }

  return *this;
}

/******************************************************************************/

template <class typeArray>
void Average<typeArray>::setDeltaTime(const double &dt)
{ 
#ifdef DEBUG_MODE_ON
  std::cerr << __FILE__ <<  " (L " << __LINE__ << ") Inside Function  " << __FUNCTION__ << "(...)" << endl;
#endif

  this->dt = dt; 
  tau += this->dt;
}

/******************************************************************************/

template <class typeArray>
void Average<typeArray>::increment(const typeArray &variable, const int &variableId)
{
#ifdef DEBUG_MODE_ON
  std::cerr << __FILE__ <<  " (L " << __LINE__ << ") Inside Function  " << __FUNCTION__ << "(...)" << endl;
#endif

  if ( (variableId < 0) || (variableId >= AverageVariables::nVariables) ){
    std::cerr << "ERROR in " << __FILE__ <<  "(L" << __LINE__ << ")" << __FUNCTION__ << "(...): " 
	      << "Did not understand variable with index \"" << variableId << "\".  "
	      << "Did you try using the variables defined in the \"AverageVariables\" namespace?"
	      << endl;
    Communication_Manager::Sync();
    Optimization_Manager::Exit_Virtual_Machine();
  }

  const float saved = (tau-this->dt) / tau;
  const float delta = this->dt / tau;

  variables[variableId] = variables[variableId]*saved + variable*delta;
}

/******************************************************************************/

template <class typeArray> 
map<int, string> Average<typeArray>::getVariableIds() const
{
#ifdef DEBUG_MODE_ON
  std::cerr << __FILE__ <<  " (L " << __LINE__ << ") Inside Function  " << __FUNCTION__ << "(...)" << endl;
#endif

  map<int, string> vars;
  
  vector<int>::const_iterator it;
  for ( it=this->variableIds.begin(); it < this->variableIds.end(); it++ ){
    string varName;
    if (*it == AverageVariables::BX)
      varName = "avgBx";
    else if (*it == AverageVariables::BY)
      varName = "avgBy";
    else if (*it == AverageVariables::BZ)
      varName = "avgBz";
    else if (*it == AverageVariables::BI)
      varName = "avgBi";
    else if (*it == AverageVariables::BJ)
      varName = "avgBj";
    else if (*it == AverageVariables::BK)
      varName = "avgBk";
    else if (*it == AverageVariables::RHO)
      varName = "avgRho";
    else if (*it == AverageVariables::CSOUND)
      varName = "avgCSound";
    else if (*it == AverageVariables::PRESSURE)
      varName = "avgPressure";
    else if (*it == AverageVariables::EI)
      varName = "avgEi";
    else if (*it == AverageVariables::EJ)
      varName = "avgEj";
    else if (*it == AverageVariables::EK)
      varName = "avgEk";
    else if (*it == AverageVariables::VX)
      varName = "avgVx";
    else if (*it == AverageVariables::VY)
      varName = "avgVy";
    else if (*it == AverageVariables::VZ)
      varName = "avgVz";
    else if (*it == AverageVariables::V)
      varName = "avgV";
    else{
      std::cerr << "ERROR in " << __FILE__ <<  "(L" << __LINE__ << ")" << __FUNCTION__ << "(...): " 
		<< "Did not understand variable with index \"" << *it << "\".  "
		<< "Did you try using the variables defined in the \"AverageVariables\" namespace?"
		<< endl;
      Communication_Manager::Sync();
      Optimization_Manager::Exit_Virtual_Machine();
    }
    
    vars.insert( pair<int, string>(*it, varName) );
  }
  
  return vars;
}

/******************************************************************************/

template <class typeArray>
const typeArray & Average<typeArray>::getVariable(const int &variableId) const
{
#ifdef DEBUG_MODE_ON
  std::cerr << __FILE__ <<  " (L " << __LINE__ << ") Inside Function  " << __FUNCTION__ << "(...)" << endl;
#endif

  if ( (variableId < 0) || (variableId >= AverageVariables::nVariables) ){
    std::cerr << "ERROR in " << __FILE__ <<  "(L" << __LINE__ << ")" << __FUNCTION__ << "(...): " 
	      << "Did not understand variable with index \"" << variableId << "\".  "
	      << "Did you try using the variables defined in the \"AverageVariables\" namespace?"
	      << endl;
    Communication_Manager::Sync();
    Optimization_Manager::Exit_Virtual_Machine();
  }

  return variables[variableId];
}

/******************************************************************************/

template <class typeArray>
void Average<typeArray>::setVariable(const typeArray &variable, const int &variableId)
{
#ifdef DEBUG_MODE_ON
  std::cerr << __FILE__ <<  " (L " << __LINE__ << ") Inside Function  " << __FUNCTION__ << "(...)" << endl;
#endif

  if ( (variableId < 0) || (variableId >= AverageVariables::nVariables) ){
    std::cerr << "ERROR in " << __FILE__ <<  "(L" << __LINE__ << ")" << __FUNCTION__ << "(...): " 
	      << "Did not understand variable with index \"" << variableId << "\".  "
	      << "Did you try using the variables defined in the \"AverageVariables\" namespace?"
	      << endl;
    Communication_Manager::Sync();
    Optimization_Manager::Exit_Virtual_Machine();
  }
  variables[variableId] = variable;
}


/******************************************************************************/

template <class typeArray>
void Average<typeArray>::reset()
{
#ifdef DEBUG_MODE_ON
   std::cerr << __FILE__ <<  " (L " << __LINE__ << ") Inside Function  " << __FUNCTION__ << "(...)" << endl;
#endif

 tau = 0.0;
  vector<int>::iterator it;
  for ( it=this->variableIds.begin() ; it < this->variableIds.end(); it++ ){
    reset(*it);
  }
}

/******************************************************************************/

template <class typeArray>
void Average<typeArray>::reset(const int &variableId)
{
#ifdef DEBUG_MODE_ON
  std::cerr << __FILE__ <<  " (L " << __LINE__ << ") Inside Function  " << __FUNCTION__ << "(...)" << endl;
#endif

  if ( (variableId < 0) || (variableId >= AverageVariables::nVariables) ){
    std::cerr << "ERROR in " << __FILE__ <<  "(L" << __LINE__ << ")" << __FUNCTION__ << "(...): " 
	      << "Did not understand variable with index \"" << variableId << "\".  "
	      << "Did you try using the variables defined in the \"AverageVariables\" namespace?"
	      << endl;
    Communication_Manager::Sync();
    Optimization_Manager::Exit_Virtual_Machine();
  }

  variables[variableId] = 0.0;
}

/******************************************************************************/

#endif