// ============================================================================
// ==                                                                        ==
// == AMDiS - Adaptive multidimensional simulations                          ==
// ==                                                                        ==
// ==  http://www.amdis-fem.org                                              ==
// ==                                                                        ==
// ============================================================================
//
// Software License for AMDiS
//
// Copyright (c) 2010 Dresden University of Technology 
// All rights reserved.
// Authors: Simon Vey, Thomas Witkowski et al.
//
// This file is part of AMDiS
//
// See also license.opensource.txt in the distribution.



/** \file FeSpaceMapping.h */

#include <vector>
#include <map>
#include <set>
#include "parallel/DofComm.h"
#include "parallel/MpiHelper.h"
#include "parallel/ParallelTypes.h"
#include "parallel/StdMpi.h"

#ifndef AMDIS_FE_SPACE_MAPPING_H
#define AMDIS_FE_SPACE_MAPPING_H

namespace AMDiS {

  using namespace std;

  /// Is used if a DOF index is mapped to multiple indices, i.e., to both, a local
  /// and a global one.
  struct MultiIndex
  {
    int local, global;
  };


  /** \brief
   * Is used to store matrix indices to all DOFs in rank's subdomain. Thus, the
   * class defines a mapping from component number and DOF index to a global
   * matrix index. This class does not calculate the indices by itself!
   */
  class DofToMatIndex
  {
  public:
    DofToMatIndex() {}

    /// Reset the data structure.
    inline void clear()
    {
      data.clear();
    }

    /** Add a new mapping for a given DOF.
     * 
     * \param[in]   component       Component number for which the mapping 
     *                              is defined.
     * \param[in]   dof             DOF index
     * \param[in]   globalMatIndex  Global matrix index.
     */
    inline void add(int component, DegreeOfFreedom dof, int globalMatIndex)
    {
      data[component][dof] = globalMatIndex;
    }

    /// Maps a global DOF index to the global matrix index for a specific 
    /// system component number.
    inline int get(int component, DegreeOfFreedom dof)
    {
      FUNCNAME("DofToMatIndex::get()");

      TEST_EXIT_DBG(data.count(component))
	("No mapping data for component %d available!\n", component);

      TEST_EXIT_DBG(data[component].count(dof))
	("Mapping for DOF %d in component %d does not exists!\n",
	 dof, component);

      return data[component][dof];
    }

  private:
    /// The mapping data. For each system component there is a specific map that
    /// maps global DOF indices to global matrix indices.
    map<int, map<DegreeOfFreedom, int> > data;
  };


  /**
   * This class defines the parallel mapping of DOFs for one FE space. It is used
   * by the class \ref ParallelDofMapping to specifiy the mapping for a set of 
   * FE spaces.
   */
  class FeSpaceDofMap
  {
  public:
    /// This constructor exists only to create std::map of this class and make
    /// use of the operator [] for read access. Should never be called.
    FeSpaceDofMap() 
    {
      ERROR_EXIT("Should not be called!\n");
    }
     
    /// This is the only valid constructur to be used. 
    FeSpaceDofMap(MPI::Intracomm m)
      : mpiComm(m),
	sendDofs(NULL),
	recvDofs(NULL),
	feSpace(NULL),
	needGlobalMapping(false),
	hasNonLocalDofs(false),
	nRankDofs(0),
	nLocalDofs(0),
	nOverallDofs(0),
	rStartDofs(0)
    {}

    /// Clears all data of the mapping.
    void clear();
    
    /// Maps a DOF index to both, the local and global index of the mapping. The
    /// global index must not be set.
    MultiIndex& operator[](DegreeOfFreedom d)
    {
      TEST_EXIT_DBG(dofMap.count(d))("Should not happen!\n");

      return dofMap[d];
    }
    
    /// Inserts a new DOF to rank's mapping. The DOF is assumed to be owend by
    /// the rank.
    void insertRankDof(DegreeOfFreedom dof0, DegreeOfFreedom dof1 = -1)
    {
      FUNCNAME("FeSpaceDofMap::insertRankDof()");
      
      TEST_EXIT_DBG(dofMap.count(dof0) == 0)("Should not happen!\n");
      
      dofMap[dof0].local = dof1;
      nLocalDofs++;
      if (dof1 != -1)
	nRankDofs++;
    }
    
    /// Inserts a new DOF to rank's mapping. The DOF exists in rank's subdomain
    /// but is owned by a different rank, thus it is part of an interior boundary.
    void insert(DegreeOfFreedom dof0, DegreeOfFreedom dof1 = -1)
    {
      FUNCNAME("FeSpaceDofMap::insert()");
      
      TEST_EXIT_DBG(dofMap.count(dof0) == 0)("Should not happen!\n");
      
      dofMap[dof0].local = dof1;
      nLocalDofs++;
      nonRankDofs.insert(dof0);
    }
    
    /// Checks if a given DOF is in the DOF mapping.
    bool isSet(DegreeOfFreedom dof)
    {
      return static_cast<bool>(dofMap.count(dof));
    }

    /// Checks if a given DOF is a rank owned DOF of the DOF mapping. The DOF must
    /// a DOF of the mapping (this is not checked here), otherwise the result is
    /// meaningsless.
    bool isRankDof(DegreeOfFreedom dof)
    {
      return !(static_cast<bool>(nonRankDofs.count(dof)));
    }
    
    /// Returns number of DOFs in the mapping.
    unsigned int size()
    {
      return dofMap.size();
    }
    
    /// Returns the raw data of the mapping.
    map<DegreeOfFreedom, MultiIndex>& getMap()
    {
      return dofMap;
    }

    /// Recomputes the mapping.
    void update();

    /// Sets the FE space this mapping corresponds to.
    void setFeSpace(const FiniteElemSpace *fe)
    {
      feSpace = fe;
    }

    /// Informs the mapping whether the mapping will include DOFs that are not
    /// owned by the rank.
    void setNonLocalDofs(bool b)
    {
      hasNonLocalDofs = b;
    }

    /// Informs the mapping whether a global index must be computed.
    void setNeedGlobalMapping(bool b)
    {
      needGlobalMapping = b;
    }

    /// Sets the DOF communicators.
    void setDofComm(DofComm &pSend, DofComm &pRecv)
    {
      sendDofs = &pSend;
      recvDofs = &pRecv;
    }

  private:
    /// Computes a global mapping from the local one.
    void computeGlobalMapping();

    /// Computes the global indices of all DOFs in the mapping that are not owned
    /// by the rank.
    void computeNonLocalIndices();

  private:
    /// MPI communicator object;
    MPI::Intracomm mpiComm;

    /// DOF communicators for all DOFs on interior boundaries.
    DofComm *sendDofs, *recvDofs;

    /// The FE space this mapping belongs to. This is used only the get the
    /// correct DOF communicator in \ref sendDofs and \ref recvDofs.
    const FiniteElemSpace *feSpace;

    /// Mapping data from DOF indices to local and global indices.
    map<DegreeOfFreedom, MultiIndex> dofMap;

    /// Set of all DOFs that are in mapping but are not owned by the rank.
    std::set<DegreeOfFreedom> nonRankDofs;

    /// If true, a global index mapping will be computed for all DOFs.
    bool needGlobalMapping;

    /// Is true if there are DOFs in at least one subdomain that are not owned
    /// by the rank. If the value is false, each rank contains only DOFs that
    /// are also owned by this rank.
    bool hasNonLocalDofs;
  public:
    /// 
    int nRankDofs, nLocalDofs, nOverallDofs, rStartDofs;
  };
  
  
  /**
   * Implements the mapping from sets of distributed DOF indices to local and
   * global indices. The mapping works for a given set of FE spaces. Furthermore,
   * this class may compute the matrix indices of the set of DOF indices.
   */
  class ParallelDofMapping
  {
  public:
    ParallelDofMapping() 
      : mpiComm(NULL),
	sendDofs(NULL),
	recvDofs(NULL),
	hasNonLocalDofs(false),
	feSpaces(0),
	nRankDofs(-1),
	nLocalDofs(-1),
	nOverallDofs(-1),
	rStartDofs(-1)
    {} 

    /** \brief Initialize the parallel DOF mapping.
     *
     * \param[in]  m                  MPI communicator.
     * \param[in]  fe                 The FE spaces of all components of the 
     *                                PDE to be solved.
     * \param[in]  uniqueFe           Unique list of FE spaces. Thus, two
     *                                arbitrary elements of this list are always
     *                                different.
     * \param[in]  needGlobalMapping  If true, the mapping computes also a global
     *                                index for the DOFs.
     * \param[in]  bNonLocalDofs      If true, at least one rank's mapping con-
     *                                taines DOFs that are not owend by the rank.
     */
    void init(MPI::Intracomm m,
	      vector<const FiniteElemSpace*> &fe,
	      vector<const FiniteElemSpace*> &uniqueFe,
	      bool needGlobalMapping,
	      bool bNonLocalDofs);
    
    /// Clear all data.
    void clear();

    /// Set the DOF communicator objects that are required to exchange information
    /// about DOFs that are on interior boundaries.
    void setDofComm(DofComm &pSend, DofComm &pRecv);

    /// Access the DOF mapping for a given FE space.
    inline FeSpaceDofMap& operator[](const FiniteElemSpace* feSpace)
    {
      FUNCNAME("ParallelDofMapping::operator[]()");

      TEST_EXIT_DBG(data.count(feSpace))
	("No data for FE space at address %p!\n", feSpace);

      return data.find(feSpace)->second;
    }

    /// Returns \ref nRankDofs, thus the number of DOFs owned by the rank.
    inline int getRankDofs()
    {
      TEST_EXIT_DBG(nRankDofs >= 0)("Should not happen!\n");

      return nRankDofs;
    }

    /// Returns \ref nLocalDofs, thus the number of DOFs in ranks subdomain.
    inline int getLocalDofs()
    {
      TEST_EXIT_DBG(nLocalDofs >= 0)("Should not happen!\n");

      return nLocalDofs;
    }

    /// Returns \ref nOverallDofs, thus the number of all DOFs in this mapping.
    inline int getOverallDofs()
    {
      TEST_EXIT_DBG(nOverallDofs >= 0)("Should not happen!\n");

      return nOverallDofs;
    }

    /// Returns \ref rStartDofs, thus the smallest global index of a DOF that is
    /// owned by the rank.
    inline int getStartDofs()
    {
      TEST_EXIT_DBG(rStartDofs >= 0)("Should not happen!\n");

      return rStartDofs;
    }

    /// Update the mapping.
    void update();

    /// Returns the global matrix index of a given DOF for a given 
    /// component number.
    inline int getMatIndex(int ithComponent, DegreeOfFreedom d)
    {
      return dofToMatIndex.get(ithComponent, d);
    }

    /// Returns the local matrix index of a given DOF for a given 
    /// component number.
    inline int getLocalMatIndex(int ithComponent, DegreeOfFreedom d)
    {
      FUNCNAME("ParallelDofMapping::getLocalMatIndex()");

      TEST_EXIT_DBG(data[feSpaces[ithComponent]].isRankDof(d))
	("Should not happen!\n");

      return dofToMatIndex.get(ithComponent, d) - rStartDofs;
    }

    // Writes all data of this object to an output stream.
    void serialize(ostream &out)
    {
      ERROR_EXIT("MUST BE IMPLEMENTED!\n");
    }

    // Reads the object data from an input stream.
    void deserialize(istream &in)
    {
      ERROR_EXIT("MUST BE IMPLEMENTED!\n");
    }

  protected:
    /// Insert a new FE space DOF mapping for a given FE space.
    void addFeSpace(const FiniteElemSpace* feSpace);

    /// Compute \ref nRankDofs.
    int computeRankDofs();

    /// Compute \ref nLocalDofs.
    int computeLocalDofs();

    /// Compute \ref nOverallDofs.
    int computeOverallDofs();

    /// Compute \ref rStartDofs.
    int computeStartDofs();

    /// Compute local and global matrix indices.
    void computeMatIndex();

  private:
    /// MPI communicator object;
    MPI::Intracomm mpiComm;

    /// DOF communicators for all DOFs on interior boundaries.
    DofComm *sendDofs, *recvDofs;

    /// Is true if there are DOFs in at least one subdomain that are not owned
    /// by the rank. If the value is false, each rank contains only DOFs that
    /// are also owned by this rank.
    bool hasNonLocalDofs;

    /// Maps from FE space pointers to DOF mappings.
    map<const FiniteElemSpace*, FeSpaceDofMap> data;

    /// The FE spaces for all components.
    vector<const FiniteElemSpace*> feSpaces;

    /// The set of all FE spaces. It uniquly contains all different FE spaces
    /// from \ref feSpaces.
    vector<const FiniteElemSpace*> feSpacesUnique;

    /// Number of DOFs owned by rank.
    int nRankDofs;

    /// Number of DOFs in rank's subdomain.
    int nLocalDofs;

    /// Number of global DOFs (this value is thus the same on all ranks).
    int nOverallDofs;

    /// Smallest global index of a DOF owned by the rank.
    int rStartDofs;

    /// Mapping from global DOF indices to global matrix indices under 
    /// consideration of possibly multiple components.
    DofToMatIndex dofToMatIndex;
  };

}

#endif