// ============================================================================
// ==                                                                        ==
// == 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;

  struct MultiIndex
  {
    int local, global;
  };

  /** \brief
   * Defines for each system component a mapping for sets of global DOF indices
   * to global matrix indices. The mapping is defined for all DOFs in rank's 
   * subdomain. When periodic boundary conditions are used, then the mapping
   * stores also information for the periodic associations of rank's DOF on
   * periodic boundaries.
   */
  class DofToMatIndex
  {
  public:
    DofToMatIndex() {}

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

    /// Add a new mapping.
    inline void add(int component, DegreeOfFreedom dof, int globalMatIndex)
    {
      data[component][dof] = globalMatIndex;
    }

    /// Map 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;
  };


  class GlobalDofMap
  {
  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.
    GlobalDofMap() 
    {
      ERROR_EXIT("Should not be called!\n");
    }
      
    GlobalDofMap(MPI::Intracomm* m)
      : mpiComm(m),
	sendDofs(NULL),
	recvDofs(NULL),
	feSpace(NULL),
	needGlobalMapping(false),
	nRankDofs(0),
	nLocalDofs(0),
	nOverallDofs(0),
	rStartDofs(0),
	overlap(false)
    {}
    
    void clear();
    
    MultiIndex& operator[](DegreeOfFreedom d)
    {
      TEST_EXIT_DBG(dofMap.count(d))("Should not happen!\n");

      return dofMap[d];
    }
    
    void insertRankDof(DegreeOfFreedom dof0, DegreeOfFreedom dof1 = -1)
    {
      FUNCNAME("GlobalDofMap::insertRankDof()");
      
      TEST_EXIT_DBG(dofMap.count(dof0) == 0)("Should not happen!\n");
      
      dofMap[dof0].local = dof1;
      nLocalDofs++;
      if (dof1 != -1)
	nRankDofs++;
    }
    
    void insert(DegreeOfFreedom dof0, DegreeOfFreedom dof1 = -1)
    {
      FUNCNAME("GlobalDofMap::insert()");
      
      TEST_EXIT_DBG(dofMap.count(dof0) == 0)("Should not happen!\n");
      
      dofMap[dof0].local = dof1;
      nLocalDofs++;
      nonRankDofs.insert(dof0);
    }
    
    bool isSet(DegreeOfFreedom dof)
    {
      return static_cast<bool>(dofMap.count(dof));
    }

    bool isRankDof(DegreeOfFreedom dof)
    {
      return !(static_cast<bool>(nonRankDofs.count(dof)));
    }
    
    unsigned int size()
    {
      return dofMap.size();
    }
    
    map<DegreeOfFreedom, MultiIndex>& getMap()
    {
      return dofMap;
    }
    
    void update();

    void addOffset(int offset);

    void computeNonLocalIndices();

    void print();

    void setFeSpace(const FiniteElemSpace *fe)
    {
      feSpace = fe;
    }

    void setOverlap(bool b)
    {
      overlap = b;
    }

    void setNeedGlobalMapping(bool b)
    {
      needGlobalMapping = b;
    }

    void setDofComm(DofComm &pSend, DofComm &pRecv)
    {
      sendDofs = &pSend;
      recvDofs = &pRecv;
    }

  private:
    MPI::Intracomm* mpiComm;

    DofComm *sendDofs;
    
    DofComm *recvDofs;

    const FiniteElemSpace *feSpace;

    /// 
    map<DegreeOfFreedom, MultiIndex> dofMap;

    std::set<DegreeOfFreedom> nonRankDofs;

    bool needGlobalMapping;
  public:
    /// 
    int nRankDofs, nLocalDofs, nOverallDofs, rStartDofs;

    bool overlap;
  };
  

  template<typename T>
  class FeSpaceData
  {
  public:
    FeSpaceData() 
      : mpiComm(NULL),
	sendDofs(NULL),
	recvDofs(NULL),
	overlap(false),
	feSpaces(0),
	nRankDofs(-1),
	nOverallDofs(-1),
	rStartDofs(-1)
    {} 

    T& operator[](const FiniteElemSpace* feSpace)
    {
      FUNCNAME("FeSpaceData::operator[]()");

      TEST_EXIT_DBG(data.count(feSpace))("Should not happen!\n");

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

    void addFeSpace(const FiniteElemSpace* feSpace)
    {
      FUNCNAME("FeSpaceData::addFeSpace()");
      
      if (data.count(feSpace))
	data.find(feSpace)->second.clear();
      else
	data.insert(make_pair(feSpace, T(mpiComm)));

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

    int getRankDofs(vector<const FiniteElemSpace*> &fe)
    {
      FUNCNAME("FeSpaceData::getRankDofs()");

      int result = 0;
      for (unsigned int i = 0; i < fe.size(); i++) {
	TEST_EXIT_DBG(data.count(fe[i]))("Cannot find FE space: %p\n", fe[i]);
	result += data[fe[i]].nRankDofs;
      }

      return result;
    }

    inline int getRankDofs()
    {
      TEST_EXIT_DBG(nRankDofs >= 0)("Should not happen!\n");

      return nRankDofs;
    }

    int getLocalDofs(vector<const FiniteElemSpace*> &fe)
    {
      FUNCNAME("FeSpaceData::getLocalDofs()");

      int result = 0;
      for (unsigned int i = 0; i < fe.size(); i++) {
	TEST_EXIT_DBG(data.count(fe[i]))("Cannot find FE space: %p\n", fe[i]);
	result += data[fe[i]].nLocalDofs;
      }

      return result;
    }

    inline int getLocalDofs()
    {
      TEST_EXIT_DBG(nLocalDofs >= 0)("Should not happen!\n");

      return nLocalDofs;
    }

    int getOverallDofs(vector<const FiniteElemSpace*> &feSpaces)
    {
      int result = 0;
      for (unsigned int i = 0; i < feSpaces.size(); i++) {
	TEST_EXIT_DBG(data.count(feSpaces[i]))("Should not happen!\n");
	result += data.find(feSpaces[i])->second.nOverallDofs;
      }

      return result;
    }

    inline int getOverallDofs()
    {
      TEST_EXIT_DBG(nOverallDofs >= 0)("Should not happen!\n");

      return nOverallDofs;
    }

    int getStartDofs(vector<const FiniteElemSpace*> &feSpaces)
    {
      int result = 0;
      for (unsigned int i = 0; i < feSpaces.size(); i++) {
	TEST_EXIT_DBG(data.count(feSpaces[i]))("Should not happen!\n");
	result += data.find(feSpaces[i])->second.rStartDofs;
      }

      return result;
    }

    int getStartDofs()
    {
      TEST_EXIT_DBG(rStartDofs >= 0)("Should not happen!\n");

      return rStartDofs;
    }

    void init(MPI::Intracomm *m,
	      vector<const FiniteElemSpace*> &fe,
	      bool needGlobalMapping,
	      bool ol)
    {
      mpiComm = m;
      feSpaces = fe;
      overlap = ol;
      for (unsigned int i = 0; i < feSpaces.size(); i++) {
	feSpacesUnique.insert(feSpaces[i]);
	
	addFeSpace(feSpaces[i]);
	data[feSpaces[i]].setNeedGlobalMapping(needGlobalMapping);
	data[feSpaces[i]].setOverlap(overlap);
      }
    }

    void update()
    {
      for (std::set<const FiniteElemSpace*>::iterator it = feSpacesUnique.begin();
	   it != feSpacesUnique.end(); ++it)
	data[*it].update();

      nRankDofs = getRankDofs(feSpaces);
      nLocalDofs = getLocalDofs(feSpaces);
      nOverallDofs = getOverallDofs(feSpaces);
      rStartDofs = getStartDofs(feSpaces);

      computeMatIndex();
    }


    void computeMatIndex()
    {
      dofToMatIndex.clear();

      int offset = rStartDofs;

      for (unsigned int i = 0; i < feSpaces.size(); i++) {

	map<DegreeOfFreedom, MultiIndex>& dofMap = data[feSpaces[i]].getMap();
	typedef map<DegreeOfFreedom, MultiIndex>::iterator ItType;
	for (ItType it = dofMap.begin(); it != dofMap.end(); ++it) {
	  if (data[feSpaces[i]].isRankDof(it->first)) {
	    int globalMatIndex = 
	      it->second.local - data[feSpaces[i]].rStartDofs + offset;
	    dofToMatIndex.add(i, it->first, globalMatIndex);
	  }
	}
	
	if (!overlap)
	  continue;
	
	TEST_EXIT_DBG(sendDofs != NULL && recvDofs != NULL)
	  ("No communicator given!\n");
	
	StdMpi<vector<DegreeOfFreedom> > stdMpi(*mpiComm);
	for (DofComm::Iterator it(*sendDofs, feSpaces[i]); 
	     !it.end(); it.nextRank()) {
	  vector<DegreeOfFreedom> sendGlobalDofs;
	  
	  for (; !it.endDofIter(); it.nextDof())
	    if (dofMap.count(it.getDofIndex()))
	      sendGlobalDofs.push_back(dofToMatIndex.get(i, it.getDofIndex()));
	  
	  stdMpi.send(it.getRank(), sendGlobalDofs);
	}
	
	for (DofComm::Iterator it(*recvDofs, feSpaces[i]); 
	     !it.end(); it.nextRank())
	  stdMpi.recv(it.getRank());
	
	stdMpi.startCommunication();
	
	{
	  int counter = 0;
	  for (DofComm::Iterator it(*recvDofs, feSpaces[i]); 
	       !it.end(); it.nextRank()) {
	    for (; !it.endDofIter(); it.nextDof()) {
	      if (dofMap.count(it.getDofIndex())) {
		DegreeOfFreedom d = stdMpi.getRecvData(it.getRank())[counter++];
		dofToMatIndex.add(i, it.getDofIndex(), d);
	      }
	    }
	  }
	}

	offset += data[feSpaces[i]].nRankDofs;
      }
    }


    inline int mapLocal(int index, int ithFeSpace)
    {
      int result = 0;
      for (int i = 0; i < ithFeSpace; i++)
	result += data[feSpaces[i]].nRankDofs;
      result += data[feSpaces[ithFeSpace]][index].local;

      return result;
    }

    inline int mapLocal(int index, const FiniteElemSpace *feSpace)
    {
      for (unsigned int i = 0; i < feSpaces.size(); i++)
	if (feSpaces[i] == feSpace)
	  return mapLocal(index, feSpace, i);

      return -1;
    }

    inline int mapGlobal(int index, int ithFeSpace)
    {
      int result = rStartDofs;
      for (int i = 0; i < ithFeSpace; i++)
	result += data[feSpaces[i]].nRankDofs;
      result += data[feSpaces[ithFeSpace]][index].local;

      return result;
    }

    inline int mapGlobal(int index, const FiniteElemSpace *feSpace)
    {
      for (unsigned int i = 0; i < feSpaces.size(); i++)
	if (feSpaces[i] == feSpace)
	  return mapGlobal(index, feSpace, i);

      return -1;
    }

    void setDofComm(DofComm &pSend, DofComm &pRecv)
    {
      sendDofs = &pSend;
      recvDofs = &pRecv;

      for (std::set<const FiniteElemSpace*>::iterator it = feSpacesUnique.begin();
	   it != feSpacesUnique.end(); ++it)
	data[*it].setDofComm(pSend, pRecv);
    }

  private:
    MPI::Intracomm* mpiComm;

    DofComm *sendDofs;
    
    DofComm *recvDofs;

    bool overlap;
    
    map<const FiniteElemSpace*, T> data;

    vector<const FiniteElemSpace*> feSpaces;

    std::set<const FiniteElemSpace*> feSpacesUnique;

    int nRankDofs, nLocalDofs, nOverallDofs, rStartDofs;

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

}

#endif