/******************************************************************************
 *
 * AMDiS - Adaptive multidimensional simulations
 *
 * Copyright (C) 2013 Dresden University of Technology. All Rights Reserved.
 * Web: https://fusionforge.zih.tu-dresden.de/projects/amdis
 *
 * Authors: 
 * Simon Vey, Thomas Witkowski, Andreas Naumann, Simon Praetorius, et al.
 *
 * This file is provided AS IS with NO WARRANTY OF ANY KIND, INCLUDING THE
 * WARRANTY OF DESIGN, MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE.
 *
 *
 * This file is part of AMDiS
 *
 * See also license.opensource.txt in the distribution.
 * 
 ******************************************************************************/

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

#include "Arh2Reader.h"
#include "detail/Arh2Reader.h"
#include <boost/iostreams/filtering_streambuf.hpp>
#include <boost/iostreams/copy.hpp>
#include <boost/iostreams/filter/zlib.hpp>

namespace AMDiS { namespace io {

  namespace Arh2Reader
  {
    using namespace std;
 
    void readByName(string filename, 
		      DOFVector<double>& vec,
		      bool writeParallel,
		      int nProcs)
    {
      readByName(filename, &vec, writeParallel, nProcs);
    }

    void readByName(string filename, 
		      DOFVector<double>* vec,
		      bool writeParallel,
		      int nProcs)
    {
      vector<DOFVector<double>*> vecs(0);
      vecs.push_back(vec);
      detail::readFile(filename, nullptr, vecs, writeParallel, nProcs, true);
    }
    
    void readByName(string filename,
		    vector<DOFVector<double>*> vecs,
		    bool writeParallel,
		    int nProcs)
    {
      detail::readFile(filename, nullptr, vecs, writeParallel, nProcs, true);
    }
    
    void readByName(string filename,
		    SystemVector *sysVec,
		    bool writeParallel,
		    int nProcs)
    {
      vector<DOFVector<double>*> vecs(0);
      for(int i = 0 ; i < sysVec->getSize(); i++)
      {
	vecs.push_back(sysVec->getDOFVector(i));
      }
      detail::readFile(filename, nullptr, vecs, writeParallel, nProcs, true);
    }

    void readFile(string filename, 
			  SystemVector *sysVec, 
			  bool writeParallel,
			  int nProcs)
    {
      vector<DOFVector<double>*> vecs(0);
      for(int i = 0 ; i < sysVec->getSize(); i++)
      {
	vecs.push_back(sysVec->getDOFVector(i));
      }
      detail::readFile(filename, nullptr, vecs, writeParallel, nProcs);
    }

    void readFile(string filename,
			DOFVector<double>* vec0,
			DOFVector<double>* vec1,
			DOFVector<double>* vec2,
			bool writeParallel,
			int nProcs)
    {
      vector<DOFVector<double>*> vecs(0);
      
      if(vec0 || vec1 || vec2)
	vecs.push_back(vec0);
      if(vec1 || vec2)
	vecs.push_back(vec1);
      if(vec2)
	vecs.push_back(vec2);

      detail::readFile(filename, nullptr, vecs, writeParallel, nProcs);
    }

    void readFile(string filename,
			vector<DOFVector<double>*> vecs,
			bool writeParallel,
			int nProcs)
    {
      detail::readFile(filename, nullptr, vecs, writeParallel, nProcs);
    }
    
    void readFile(string filename,
		      Mesh* mesh,
		      bool writeParallel,
		      int nProcs)
    {
      vector<DOFVector<double>*> vecs(0);
      detail::readFile(filename, mesh, vecs, writeParallel, nProcs);
    }

    int readNumOfValueVectors(string filename, bool writeParallel)
    {
      FUNCNAME("Arh2Reader::readNumOfValueVectors");
      
      ifstream file;
      
      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);
#ifdef HAVE_PARALLEL_DOMAIN_AMDIS
	filename = name + "-p" + lexical_cast<string>(MPI::COMM_WORLD.Get_rank()) + "-.arh";
#endif
      }
      
      file.open(filename.c_str(), ios::in | ios::binary);
      TEST_EXIT(file.is_open())
	  ("Cannot open file %s\n", filename.c_str());

      uint32_t nValueVectors = 0;

      detail::firstRead(file);
      long pos = file.tellg();
      file.seekg(pos + 16);
      file.read(reinterpret_cast<char*>(&nValueVectors), 4);
      
      file.close();
      return nValueVectors;
    }

    int readHeaderSize(string filename, bool writeParallel)
    {
      FUNCNAME("Arh2Reader::readHeaderSize");
      
      ifstream file;
      
      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);
#ifdef HAVE_PARALLEL_DOMAIN_AMDIS
	filename = name + "-p" + lexical_cast<string>(MPI::COMM_WORLD.Get_rank()) + "-.arh";
#endif
      }
      
      file.open(filename.c_str(), ios::in | ios::binary);
      TEST_EXIT(file.is_open())
	  ("Cannot open file %s\n", filename.c_str());
    
      uint32_t headerLen = 0;
      detail::firstRead(file);
      file.read(reinterpret_cast<char*>(&headerLen), 4);

      file.close();
      return headerLen;
    }

    bool isReadable(string filename, bool writeParallel)
    {
      FUNCNAME("Arh2Reader::isReadable");
      
      ifstream file;
      
      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);
#ifdef HAVE_PARALLEL_DOMAIN_AMDIS
	filename = name + "-p" + lexical_cast<string>(MPI::COMM_WORLD.Get_rank()) + "-.arh";
#endif
      }
      
      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, ' ');
      uint8_t major = 0, minor = 0;
     
      file.read(const_cast<char*>(typeId.data()), 4);
      file.read(reinterpret_cast<char*>(&major), 1);
      file.read(reinterpret_cast<char*>(&minor), 1);
      
      file.close();
      return (typeId == "arh2" && major == MAJOR && minor <= MINOR) ? true : false;
    }
    
    void readMeta(string filename,
			    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, vecs);
    }
    
    void readMeta(string filename, vector<DOFVector<double>*> vecs)
    {
      FUNCNAME("Arh2Reader::readMeta()");
      
      Mesh* mesh = nullptr;
      for (size_t i = 0; i < vecs.size(); i++)
      {
	if(vecs[i])
	{
	  if(!mesh)
	    mesh = vecs[i]->getFeSpace()->getMesh();
	  else
	    TEST_EXIT(mesh == vecs[i]->getFeSpace()->getMesh())
	      ("The mesh of the DOFVectors should be the same for Reader because in one file there is only one mesh.\n");
	}
      }
      if(!mesh)
      {
	WARNING("You haven't specified the target.\n");
	return;
      }
      // === 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.string() + "/" + arhPrefix;
	if (nProc == 1)
	  arhFilename += ".arh";
	else
	  arhFilename += "-p" + boost::lexical_cast<string>(*it) + "-.arh";

	MSG("ARH2 file read from: %s\n", arhFilename.c_str());
	detail::read(arhFilename, mesh, vecs);
      }
    }
    
    int readMetaData(string filename,
				map<int, int> &elInRank,
				map<int, int> &elCodeSize,
				string &arhPrefix)
    {
      FUNCNAME("Arh2Reader::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;
      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 readMetaData(string filename)
    {
      FUNCNAME("Arh2Reader::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;
    }

    int readMetaFromArh(std::string filename,
			std::map<int, int> &elInRank,
			std::map<int, int> &elCodeSize){
    
      
      int nProcs = MPI::COMM_WORLD.Get_size();
      std::vector<std::set<std::pair<int, int> > > data(nProcs); 
      // data format: (rank ; (elIndex; elCodeSize) )  

      //collect data      
      for (int i = 0; i < nProcs; i++) {
	string fn = filename + "-p" + boost::lexical_cast<string>(i) + "-.arh";
	ifstream file;
	file.open(fn.c_str(), ios::in | ios::binary);
	if (!file.is_open()) {
	  cout << "Cannot open file \"" << fn << "\"\n";
	  exit(0);
	}
	cout << MPI::COMM_WORLD.Get_rank() <<" Read file \"" << fn << "\"\n";
	
	string typeId(4, ' ');
	file.read(const_cast<char*>(typeId.data()), 4);
	
	if(typeId == "arh2")
	  readMetaFromSgArh(file, typeId, i, data);
	else
	  {
	    cout << "file \"" << fn << "\"is not arh file.\n";
	    cout << "should not happen.\n";
	    exit(0);
	  }
	file.close();
      }
      //make elInRank and elCodeSize-Map
      for (int i = 0; i < nProcs; i++) {
	for (std::set<pair<int, int> >::iterator it = data[i].begin();
	     it != data[i].end(); ++it){
	  elInRank[it->first]=i;
	  elCodeSize[it->first]=it->second;
	}
      }
      return nProcs;
    }
   
    
    void readMetaFromSgArh(std::ifstream& file, string typeId, int nProc, 
			   std::vector< std::set<std::pair<int, int> > >& data)
    {
      uint8_t major = 0, minor = 0;
      uint32_t dow = 0, dim = 0;
      uint32_t headerLen = 0;
      uint32_t nMacroElements = 0;
      uint32_t nValueVectors = 0;
      uint32_t nFeSpaces = 0;
      string cps(4, ' ');    

      file.read(reinterpret_cast<char*>(&major), 1);
      file.read(reinterpret_cast<char*>(&minor), 1);
      file.read(reinterpret_cast<char*>(&headerLen), 4);
      file.read(reinterpret_cast<char*>(&dow), 4);
      file.read(reinterpret_cast<char*>(&dim), 4);
      file.read(reinterpret_cast<char*>(&nFeSpaces), 4);
      file.read(reinterpret_cast<char*>(&nValueVectors), 4);
      file.read(reinterpret_cast<char*>(&nMacroElements), 4);

      switch(minor) {
      case 1: 
	file.read(const_cast<char*>(cps.data()), 4);
	break;
      default:
	cps = "null";break;
      }
    
      vector<int> macroElIndex(nMacroElements);
      vector<int> macroElSize(nMacroElements);
      vector<vector<int> > sortedFeSpaces(nFeSpaces);
	
      // Read macro table
      for(int i = 0; i < nMacroElements; i++)
	{
	  file.read(reinterpret_cast<char*>(&macroElIndex[i]), 4);
	  uint32_t elpos = 0;
	  file.read(reinterpret_cast<char*>(&macroElSize[i]), 4);
    
	}
      // Read feSpace table
      for(int i = 0; i < nFeSpaces; i++)
	{
	  for(int j = 0; j < 4; j++)
	    {
	      uint32_t perDOFs = 0;
	      file.read(reinterpret_cast<char*>(&perDOFs), 4);
	    }
	}
      // Read value table
      for(int i = 0; i < nValueVectors; i++)
	{
	  string tmpString("");
	  uint32_t tmpInt = 0;
	  file.read(reinterpret_cast<char*>(&tmpInt), 4);
	  tmpString.resize(tmpInt, ' ');
	  file.read(const_cast<char*>(tmpString.data()), tmpInt); //
	  file.read(reinterpret_cast<char*>(&tmpInt), 4);
	  sortedFeSpaces[tmpInt].push_back(i);
	} 
  
      for (unsigned int i = 0; i < nMacroElements; i++) {
	stringstream dataStream(ios::out | ios::in | ios::binary);
	int size = macroElSize[i];
	char* buffer = new char[size];
	file.read(buffer, size);
	dataStream.write(buffer, size);
	delete[] buffer;
     
	if(cps == "zlib") {
	  stringstream tmp(ios::out | ios::in);
	  boost::iostreams::filtering_streambuf<boost::iostreams::input> in;
	  in.push(boost::iostreams::zlib_decompressor());
	  in.push(dataStream);
	  boost::iostreams::copy(in, tmp);
	  dataStream.str(tmp.str());
	}
	  
	uint32_t nStructureCodes = 0;
	uint32_t codeSize = 0;
    
	dataStream.read(reinterpret_cast<char*>(&nStructureCodes), 4);
	dataStream.read(reinterpret_cast<char*>(&codeSize), 4);

	cout << "  -> macro el " << macroElIndex[i] << "\n";
	data[nProc].insert(make_pair(macroElIndex[i], codeSize));
    
	vector<uint64_t> structureCode(nStructureCodes);
	dataStream.read(reinterpret_cast<char*>(&(structureCode[0])), 8 * nStructureCodes);

	if (nValueVectors > 0) {
	  vector<vector<double> > values(nValueVectors);
	  int valuePos = 0;
	  for(unsigned int j = 0; j < nFeSpaces; j++)
	    {
	      uint32_t nValuesPerVector = 0;
	      dataStream.read(reinterpret_cast<char*>(&nValuesPerVector), 4);
	      for(unsigned k = 0; k < sortedFeSpaces[j].size(); k++)
		{
		  values[valuePos + k].resize(nValuesPerVector);
		  dataStream.read(reinterpret_cast<char*>(&(values[valuePos + k][0])), 8 * nValuesPerVector);
		}
	      valuePos += sortedFeSpaces[j].size();
	    }
	}
      }

      cout << MPI::COMM_WORLD.Get_rank() <<" Read ko " <<"\n"; 
    }
    
    
  } // end namespace Arh2Reader
  } } // end namespace io, AMDiS