//
// 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.


#include <fstream>
#include <stdint.h>
#include <boost/filesystem.hpp>

#include "ArhReader.h"
#include "Mesh.h"
#include "MeshStructure.h"
#include "Traverse.h"
#include "DOFVector.h"
#include "Debug.h"

namespace AMDiS {

  using namespace std;


  void ArhReader::read(string filename, Mesh *mesh,
		       DOFVector<double>* vec0,
		       DOFVector<double>* vec1,
		       DOFVector<double>* vec2,
		       bool writeParallel,
		       int nProcs)
  {
    vector<DOFVector<double>*> vecs(0);
    if (vec0)
      vecs.push_back(vec0);
    if (vec1)
      vecs.push_back(vec1);
    if (vec2)
      vecs.push_back(vec2);

    ArhReader::read(filename, mesh, vecs, writeParallel, nProcs);
  }


  void ArhReader::read(string filename, Mesh *mesh,
		       vector<DOFVector<double>*> vecs,
		       bool writeParallel,
		       int nProcs)
  {
    FUNCNAME("ArhReader::read()");

    if (writeParallel) {
      using boost::lexical_cast;
      int sPos = filename.find(".arh");
      TEST_EXIT(sPos >= 0)("Failed to find file postfix!\n");
      string name = filename.substr(0, sPos);

      if (nProcs == -1) {
#ifdef HAVE_PARALLEL_DOMAIN_AMDIS
	string procFilename = name + "-p" + lexical_cast<string>(MPI::COMM_WORLD.Get_rank()) + "-.arh";
	readFile(procFilename, mesh, vecs);
#else
	ERROR_EXIT("Reading parallel ARH files in sequential computations requires to specify the number of nodes on which the ARH file was created!\n");
#endif
      } else {
	for (int i = 0; i < nProcs; i++) {
	  string procFilename = name + "-p" + lexical_cast<string>(i) + "-.arh";
	  readFile(procFilename, mesh, vecs);
	}
      }
    } else {
      readFile(filename, mesh, vecs);
    }
    MSG("ARH file read from: %s\n", filename.c_str());
  }


  void ArhReader::readMeta(string filename,
			   Mesh *mesh,
			   DOFVector<double>* vec0,
			   DOFVector<double>* vec1,
			   DOFVector<double>* vec2)
  {
    vector<DOFVector<double>*> vecs;
    if (vec0)
      vecs.push_back(vec0);
    if (vec1)
      vecs.push_back(vec1);
    if (vec2)
      vecs.push_back(vec2);

    readMeta(filename, mesh, vecs);
  }


  void ArhReader::readMeta(string filename,
			   Mesh *mesh,
			   vector<DOFVector<double>*> vecs)
  {
    FUNCNAME("ArhReader::readMeta()");

    
    // === Read the meta arh file. ===

    string arhPrefix = "";
    map<int, int> elInRank;
    map<int, int> elCodeSize;
    int nProc = readMetaData(filename, elInRank, elCodeSize, arhPrefix);
    

    // === Check which arh files must be read by current rank. ===

    // Set of all file indices which should be read to restore all macro elements.
    std::set<int> readArhFiles;

    TraverseStack stack;
    ElInfo *elInfo = stack.traverseFirst(mesh, 0, Mesh::CALL_EL_LEVEL);
    while (elInfo) {
      int macroElIndex = elInfo->getElement()->getIndex();
      TEST_EXIT(elInRank.count(macroElIndex))("Should not happen!\n");
      readArhFiles.insert(elInRank[macroElIndex]);
      elInfo = stack.traverseNext(elInfo);
    } 


    // === Read the individual arh files. ===

    boost::filesystem::path p(filename.c_str());
    boost::filesystem::path directory = p.parent_path();

    for (std::set<int>::iterator it = readArhFiles.begin();
	 it != readArhFiles.end(); ++it) {
      string arhFilename = directory.native() + "/" + arhPrefix;
      if (nProc == 1)
	arhFilename += ".arh";
      else
	arhFilename += "-p" + boost::lexical_cast<string>(*it) + "-.arh";

      MSG("ARH file read from: %s\n", arhFilename.c_str());
      readFile(arhFilename, mesh, vecs);
    }
  }


  int ArhReader::readMetaData(string filename,
			      map<int, int> &elInRank,
			      map<int, int> &elCodeSize,
			      string &arhPrefix)
  {
    ifstream file;
    file.open(filename.c_str());
    TEST_EXIT(file.is_open())
      ("Cannot open arh meta file \"%s\"\n", filename.c_str());

    string readStr = "";
    file >> readStr;
    arhPrefix = "";
    file >> arhPrefix;
    int nProc;
    file >> nProc;

    // Maps to each macro element index the arh file index it is stored in.
    for (int i = 0; i < nProc; i++) {
      int tmp;
      file >> tmp;
      TEST_EXIT(tmp == i)("Should not happen!\n");
      int nMacroEl;
      file >> nMacroEl;
      for (int j = 0; j < nMacroEl; j++) {
	int elIndex, codeSize;
	file >> elIndex;
	file >> codeSize;
	elInRank[elIndex] = i;
	elCodeSize[elIndex] = codeSize;
      }
    }

    file.close();

    return nProc;
  }


  int ArhReader::readMetaData(string filename)
  {
    FUNCNAME("ArhReader::readMetaData()");

    ifstream file;
    file.open(filename.c_str());
    TEST_EXIT(file.is_open())
      ("Cannot open arh meta file \"%s\"\n", filename.c_str());

    string readStr = "";
    file >> readStr;
    file >> readStr;
    int nProc;
    file >> nProc;
    file.close();
    return nProc;
  }


  void ArhReader::readFile(string filename,
			   Mesh *mesh,
			   vector<DOFVector<double>*> vecs)
  {
    FUNCNAME("ArhReader::readFile()");

    // === Get set of all macro elements in mesh. ===
    std::set<int> macroInMesh;
    for (std::deque<MacroElement*>::iterator it = mesh->getMacroElements().begin();
	 it != mesh->getMacroElements().end(); ++it)
      macroInMesh.insert((*it)->getIndex());


    RefinementManager *refManager = NULL;
    switch (mesh->getDim()) {
    case 2:
      refManager = new RefinementManager2d();
      break;
    case 3:
      refManager = new RefinementManager3d();
      break;
    default:
      ERROR_EXIT("Should not happen!\n");
    }


    ifstream file;
    file.open(filename.c_str(), ios::in | ios::binary);
    TEST_EXIT(file.is_open())
      ("Cannot open file %s\n", filename.c_str());

    string typeId(4, ' ');
    uint32_t nMacroElements = 0;
    uint32_t nValueVectors = 0;
    uint32_t nAllValues = 0;

    file.read(const_cast<char*>(typeId.data()), 4);
    file.read(reinterpret_cast<char*>(&nMacroElements), 4);
    file.read(reinterpret_cast<char*>(&nValueVectors), 4);
    file.read(reinterpret_cast<char*>(&nAllValues), 4);

    TEST_EXIT(nValueVectors == vecs.size())
      ("File %s has %d vectors, but %d DOFVectors are provided!\n",
       filename.c_str(), nValueVectors, vecs.size());

    for (unsigned int i = 0; i < nMacroElements; i++) {
      uint32_t elIndex = 0;
      uint32_t nStructureCodes = 0;
      uint32_t codeSize = 0;

      file.read(reinterpret_cast<char*>(&elIndex), 4);
      file.read(reinterpret_cast<char*>(&nStructureCodes), 4);
      file.read(reinterpret_cast<char*>(&codeSize), 4);

      vector<uint64_t> structureCode(nStructureCodes);
      file.read(reinterpret_cast<char*>(&(structureCode[0])), 8 * nStructureCodes);

      MeshStructure elementStructure;
      elementStructure.init(structureCode, codeSize);
      if (macroInMesh.count(elIndex))
	elementStructure.fitMeshToStructure(mesh, refManager, false, elIndex);

      if (nValueVectors > 0) {
	uint32_t nValuesPerVector = 0;
	file.read(reinterpret_cast<char*>(&nValuesPerVector), 4);

	for (unsigned int j = 0; j < nValueVectors; j++) {
	  vector<double> values(nValuesPerVector);
	  file.read(reinterpret_cast<char*>(&(values[0])), 8 * nValuesPerVector);
	  if (vecs[j] != NULL) {
	    if (macroInMesh.count(elIndex))
	      setDofValues(elIndex, mesh, values, vecs[j]);
	  }
	}
      }
    }

    file.close();

    delete refManager;
  }


  void ArhReader::setDofValues(int macroElIndex, Mesh *mesh,
			       vector<double>& values, DOFVector<double>* vec)
  {
    FUNCNAME("ArhReader::setDofValues()");

    bool macroElement = false;
    int valuePos = 0;
    TraverseStack stack;
    ElInfo *elInfo = stack.traverseFirstOneMacro(mesh, macroElIndex, -1,
						 Mesh::CALL_EVERY_EL_PREORDER);
    while (elInfo) {
      if (!macroElement) {
	Element *mEl = elInfo->getMacroElement()->getElement();
	for (int i = 0; i < mesh->getGeo(VERTEX); i++)
	  (*vec)[mEl->getDof(i, 0)] = values[valuePos++];
	macroElement = true;
      }

      Element *el = elInfo->getElement();
      if (!el->isLeaf())
	(*vec)[el->getChild(0)->getDof(mesh->getDim(), 0)] = values[valuePos++];

      elInfo = stack.traverseNext(elInfo);
    }
  }


  int ArhReader::getNumValueVectors(string filename)
  {
    ifstream file;
    file.open(filename.c_str(), ios::in | ios::binary);
    
    string typeId = "";
    uint32_t nMacroElements = 0;
    uint32_t nValueVectors = 0;
    
    file.read(const_cast<char*>(typeId.data()), 4);
    file.read(reinterpret_cast<char*>(&nMacroElements), 4);
    file.read(reinterpret_cast<char*>(&nValueVectors), 4);
    
    file.close();
    
    return nValueVectors;
  }


  //the following three functions are identical to read/readfile except that they read from
  //a block of memory instead of from a file
  void ArhReader::readFromMemoryBlock(vector<char> &data, Mesh *mesh,
				      DOFVector<double>* vec0,
				      DOFVector<double>* vec1,
				      DOFVector<double>* vec2,
				      bool writeParallel,
				      int nProcs)
  {
    uint32_t nValueVectors;
    memcpy(reinterpret_cast<char*>(&nValueVectors), &data[8], 4);
    vector<DOFVector<double>*> vecs(0);
    if (nValueVectors > 0)
      vecs.push_back(vec0);
    if (nValueVectors > 1)
      vecs.push_back(vec1);
    if (nValueVectors > 2)
      vecs.push_back(vec2);
    for (uint32_t i = 3; i < nValueVectors; i++)
      vecs.push_back(NULL);

    readFromMemoryBlock(data, mesh, vecs, writeParallel, nProcs);
  }


  void ArhReader::readFromMemoryBlock(vector<char> &data, Mesh *mesh,
				      vector<DOFVector<double>*> vecs,
				      bool writeParallel,
				      int nProcs)
  {
    FUNCNAME("ArhReader::readFromMemoryBlock()");

    if (writeParallel) {
      ERROR_EXIT("Reading a hierarchy from several memory blocks is not implemented yet!\n");
    } else {
      readBlock(data, mesh, vecs);
    }
  }


  void ArhReader::readBlock(vector<char> &data,
			    Mesh *mesh,
			    vector<DOFVector<double>*> vecs)
  {
    FUNCNAME("ArhReader::readBlock()");

    // === Get set of all macro elements in mesh. ===
    std::set<int> macroInMesh;
    for (std::deque<MacroElement*>::iterator it = mesh->getMacroElements().begin();
                                         it != mesh->getMacroElements().end(); ++it)
      macroInMesh.insert((*it)->getIndex());


    RefinementManager *refManager = NULL;
    switch (mesh->getDim()) {
    case 2:
      refManager = new RefinementManager2d();
      break;
    case 3:
      refManager = new RefinementManager3d();
      break;
    default:
      ERROR_EXIT("Should not happen!\n");
    }

    string typeId = "";
    uint32_t nMacroElements = 0;
    uint32_t nValueVectors = 0;
    uint32_t nAllValues = 0;
    uint32_t mem_index = 0;

    memcpy(const_cast<char*>(typeId.data()), &data[mem_index], 4);
    mem_index += 4;
    memcpy(reinterpret_cast<char*>(&nMacroElements), &data[mem_index], 4);
    mem_index += 4;
    memcpy(reinterpret_cast<char*>(&nValueVectors), &data[mem_index], 4);
    mem_index += 4;
    memcpy(reinterpret_cast<char*>(&nAllValues), &data[mem_index], 4);
    mem_index += 4;

    TEST_EXIT(nValueVectors == vecs.size())
      ("data has %d vectors, but %d DOFVectors are provided!\n", nValueVectors, vecs.size());

    for (unsigned int i = 0; i < nMacroElements; i++) {
      uint32_t elIndex = 0;
      uint32_t nStructureCodes = 0;
      uint32_t codeSize = 0;

      memcpy(reinterpret_cast<char*>(&elIndex), &data[mem_index], 4);
      mem_index += 4;
      memcpy(reinterpret_cast<char*>(&nStructureCodes), &data[mem_index], 4);
      mem_index += 4;
      memcpy(reinterpret_cast<char*>(&codeSize), &data[mem_index], 4);
      mem_index += 4;

      vector<uint64_t> structureCode(nStructureCodes);
      memcpy(reinterpret_cast<char*>(&structureCode[0]), &data[mem_index], 8 * nStructureCodes);
      mem_index += 8 * nStructureCodes;

      MeshStructure elementStructure;
      elementStructure.init(structureCode, codeSize);
      if (macroInMesh.count(elIndex) == 1)
        elementStructure.fitMeshToStructure(mesh, refManager, false, elIndex);

      if(nValueVectors > 0){
        uint32_t nValuesPerVector = 0;
        memcpy(reinterpret_cast<char*>(&nValuesPerVector), &data[mem_index], 4);
        mem_index += 4;

        for (unsigned int j = 0; j < nValueVectors; j++) {
          vector<double> values(nValuesPerVector);
          memcpy(reinterpret_cast<char*>(&values[0]), &data[mem_index], 8 * nValuesPerVector);
          mem_index += 8 * nValuesPerVector;
          if (vecs[j] != NULL) {
            if (macroInMesh.count(elIndex) == 1)
              setDofValues(elIndex, mesh, values, vecs[j]);
          }
        }
      }
    }

    delete refManager;
  }

}