/*
		This file is part of libRAMSES++
			a C++ library to access snapshot files
			generated by the simulation code RAMSES by R. Teyssier

    Copyright (C) 2008-09  Oliver Hahn, ojha@gmx.de

    This program is free software: you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation, either version 3 of the License, or
    (at your option) any later version.

    This program is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with this program.  If not, see <http://www.gnu.org/licenses/>.
*/

#ifndef __RAMSES_HYDRO_DATA_HH
#define __RAMSES_HYDRO_DATA_HH

#include <fstream>
#include <iostream>
#include <iomanip>
#include <vector>
#include <cmath>

#include "FortranUnformatted_IO.hh"
#include "RAMSES_info.hh"
#include "RAMSES_amr_data.hh"

namespace RAMSES{
namespace HYDRO{

//! internal hydro variable indices
enum hydro_var
{
	density     = 1,
	velocity_x  = 2,
	velocity_y  = 3,
	velocity_z  = 4,
	pressure    = 5,
	metallicity = 6
};

//! names of possible variables stored in a RAMSES hydro file
const char ramses_hydro_variables[][64] = {
	{"Density"},
	{"x-velocity"},
	{"y-velocity"},
	{"z-velocity"},
	{"Pressure"},
	{"Metallicity"} };


/**************************************************************************************\
\**************************************************************************************/

/*!
 * @class RAMSES::HYDRO::proto_data
 * @brief base class for all things hydro
 *
 * This class provides the base class for all cell based variables defined
 * on the AMR mesh, accessible through the tree data structure.
 * @sa RAMSES::HYDRO::data, RAMSES::HYDRO::empty_data
 */
template< typename TreeType_, typename ValueType_=double >
class proto_data{

public:
	std::vector< std::vector<ValueType_> > m_var_array;
protected:
	TreeType_& m_tree;  //!< reference to underlying AMR tree structure
	unsigned			m_cpu;			//!< the computational domain
	unsigned
		m_minlevel, 				//!< maximum refinement level to be read from file
		m_maxlevel;					//!< minimum refinement level to be read from file
	unsigned 			m_twotondim;//!< 2**ndim
	unsigned			m_ilevel; 	//!< the refinement level

	//! array holding the actual data

public:
	
	//! constructor for the base class of all hydro data objects
	/*! 
	 * @param AMRtree reference to the underlying AMR tree data structure object
	 */
	explicit proto_data( TreeType_& AMRtree )
	: m_tree(AMRtree), m_cpu( AMRtree.m_cpu ), 
	  m_minlevel( AMRtree.m_minlevel ), m_maxlevel( AMRtree.m_maxlevel ),
		m_twotondim( (unsigned)(pow(2, AMRtree.m_header.ndim)+0.5) )
	{ }

	//! access the value of the cells associated with the oct designated by the iterator
	/*!
	 * @param it the grid iterator pointing to the current oct
	 * @param ind index of the child cell of the current oct (0..7)
	 */
	inline ValueType_& cell_value( const typename TreeType_::iterator& it, int ind )
	{
		unsigned ipos   = it.get_absolute_position();
		unsigned ilevel = it.get_level();//-m_minlevel;
		return (m_var_array[ilevel])[m_twotondim*ipos+ind];
	}

	//! access the value of the cells associated with the oct designated by the iterator
	/*!
	 * @param it the grid iterator pointing to the current oct
	 * @param ind index of the child cell of the current oct (0..7)
	 */
	inline ValueType_& operator()( const typename TreeType_::iterator& it, int ind )
	{	return cell_value(it,ind); }
	
	
	//! combines all elements of this instance with that of another using a binary operator
	/*!
	 * @param o  the other data object with which to combine the elements
	 * @param op the binary operator to be used in combining elements
	 */
	template<typename BinaryOperator_>
	void combine_with( const proto_data<TreeType_,ValueType_>& o, const BinaryOperator_& op )
	{
		if( m_minlevel != o.m_minlevel || m_maxlevel != o.m_maxlevel 
				|| m_twotondim != o.m_twotondim || &m_tree != &o.m_tree 
				|| m_var_array.size() != o.m_var_array.size() ){
			std::cerr << "this #levels=" << m_var_array.size() << ", other #levels=" << o.m_var_array.size() << std::endl;
			throw std::runtime_error("Error: trying to combine incompatible mesh data.");
		}
		
		for( unsigned ilvl=0; ilvl<m_var_array.size(); ++ilvl )
		{
			if( m_var_array[ilvl].size() != o.m_var_array[ilvl].size() ){
				std::cerr << "ilvl=" << ilvl << ", this size=" << m_var_array[ilvl].size() << ", other size=" << o.m_var_array[ilvl].size() << std::endl;
				throw std::runtime_error("Error: trying to combine incompatible mesh data.");
			}
			
			for( unsigned i=0; i<m_var_array[ilvl].size(); ++i )
				m_var_array[ilvl][i] = op( m_var_array[ilvl][i], o.m_var_array[ilvl][i] );
		}
		
	}
};

/**************************************************************************************\
\**************************************************************************************/

/*!
 * @class RAMSES::HYDRO::empty_data
 * @brief encapsulates additional (non-RAMSES) hydro data
 *
 * This class provides a wrapper for additional hydro variables created
 * during post-processing. They are stored compatibly with the normal
 * RAMSES hydro variables and thus can be accessed via the tree.
 * @sa RAMSES::HYDRO::data, RAMSES::HYDRO::proto_data
 */
template< typename TreeType_, typename ValueType_=double >
class empty_data : public proto_data<TreeType_,ValueType_>{
	
public:

	//! constructor for an empty hydro variable defined on the AMR mesh
	/*!
	 * @param AMRtree the underlying hierarchical tree data structure
	 * @param v the value with which the variable should be initialized, default is a (double) zero.
	 */
	explicit empty_data( TreeType_& AMRtree, ValueType_ v=(ValueType_)0.0 )
	: proto_data<TreeType_,ValueType_>(AMRtree)
	{
		this->m_var_array.assign( this->m_maxlevel+1, std::vector<ValueType_>() );
		
		for( unsigned ilvl = this->m_minlevel; ilvl<=this->m_maxlevel; ++ilvl ){
			typename TreeType_::iterator grid_it = this->m_tree.begin( ilvl );
			while( grid_it != this->m_tree.end(ilvl) ){

				for( unsigned j=0;j<this->m_twotondim;++j )
					this->m_var_array[ilvl].push_back( v );

				++grid_it;
			}
		}
	}


	//! write the new hydro variable to a RAMSES compatible output file
	/*!
	 * writes the hydro variable to a RAMSES compatible output file named
	 * after the convention
	 *  (path)/(basename)_(DOMAIN).out(snap_num)
	 *
	 * @param path the path where to store the files
	 * @param basename the filename base string to prepend to the domain number
	 * @param snap_num the number of the snapshot (default is zero).
	 */
	void save( std::string path, std::string basename, unsigned snap_num=0 )
	{
		char fullname[256];
		sprintf(fullname,"%s/%s_%05d.out%05d",path.c_str(),basename.c_str(), snap_num, this->m_tree.m_cpu );
		std::ofstream ofs( fullname, std::ios::binary|std::ios::trunc );
		
		
		typename TreeType_::iterator it;
		unsigned ncpu = this->m_tree.m_header.ncpu;
		
		//std::cerr << "ncpu = " << ncpu << std::endl;

		for( unsigned ilvl = 0; ilvl<=this->m_maxlevel; ++ilvl ){
			std::vector< std::vector<ValueType_> > 
				temp1 (ncpu, std::vector<ValueType_>() ),
				temp2 (ncpu, std::vector<ValueType_>() ),
				temp3 (ncpu, std::vector<ValueType_>() ),
				temp4 (ncpu, std::vector<ValueType_>() ),
				temp5 (ncpu, std::vector<ValueType_>() ),
				temp6 (ncpu, std::vector<ValueType_>() ),
				temp7 (ncpu, std::vector<ValueType_>() ),
				temp8 (ncpu, std::vector<ValueType_>() );
			
			it = this->m_tree.begin(ilvl);
			
			while( it!= this->m_tree.end(ilvl) )
			{
				temp1[ it.get_domain()-1 ].push_back( (*this)(it,0) );
				temp2[ it.get_domain()-1 ].push_back( (*this)(it,1) );
				temp3[ it.get_domain()-1 ].push_back( (*this)(it,2) );
				temp4[ it.get_domain()-1 ].push_back( (*this)(it,3) );
				temp5[ it.get_domain()-1 ].push_back( (*this)(it,4) );
				temp6[ it.get_domain()-1 ].push_back( (*this)(it,5) );
				temp7[ it.get_domain()-1 ].push_back( (*this)(it,6) );
				temp8[ it.get_domain()-1 ].push_back( (*this)(it,7) );
				++it;
			}
			
			
			for( unsigned icpu = 0; icpu<ncpu; ++icpu ){
				unsigned nn;
				
				nn = temp1[icpu].size() * sizeof( ValueType_ );
				
				if( nn > 0 )
				{
					ofs.write( (char*)&nn, sizeof(unsigned) );
					ofs.write( (char*)&temp1[icpu][0], nn );
					ofs.write( (char*)&nn, sizeof(unsigned) );
					
					ofs.write( (char*)&nn, sizeof(unsigned) );
					ofs.write( (char*)&temp2[icpu][0], nn );
					ofs.write( (char*)&nn, sizeof(unsigned) );
					
					ofs.write( (char*)&nn, sizeof(unsigned) );
					ofs.write( (char*)&temp3[icpu][0], nn );
					ofs.write( (char*)&nn, sizeof(unsigned) );
					
					ofs.write( (char*)&nn, sizeof(unsigned) );
					ofs.write( (char*)&temp4[icpu][0], nn );
					ofs.write( (char*)&nn, sizeof(unsigned) );
					
					ofs.write( (char*)&nn, sizeof(unsigned) );
					ofs.write( (char*)&temp5[icpu][0], nn );
					ofs.write( (char*)&nn, sizeof(unsigned) );
					
					ofs.write( (char*)&nn, sizeof(unsigned) );
					ofs.write( (char*)&temp6[icpu][0], nn );
					ofs.write( (char*)&nn, sizeof(unsigned) );
					
					ofs.write( (char*)&nn, sizeof(unsigned) );
					ofs.write( (char*)&temp7[icpu][0], nn );
					ofs.write( (char*)&nn, sizeof(unsigned) );
					
					ofs.write( (char*)&nn, sizeof(unsigned) );
					ofs.write( (char*)&temp8[icpu][0], nn );
					ofs.write( (char*)&nn, sizeof(unsigned) );
				}
			}
		}
	}
	
	//! reads an additional from a RAMSES compatible (single var) output file
	/*!
	 * @param basename the base string used for the variable
	 */
	void read( std::string basename )
	{
		char fullname[256];
		sprintf(fullname,"%s_%05d.out%05d",basename.c_str(), this->m_tree.m_header.nout[0], this->m_tree.m_cpu );
		std::ifstream ifs( fullname, std::ios::binary );

		for( unsigned ilvl = this->m_minlevel; ilvl<=this->m_maxlevel; ++ilvl ){
			unsigned nn = this->m_var_array[ilvl].size() * sizeof(ValueType_);
			unsigned nnfile;
			ifs.read( (char*)&nnfile, sizeof(unsigned) );
			if( nn != nnfile ){
				std::cerr << "Error: dimension mismatch between AMR tree and file data!" << std::endl;
				std::cerr << "       found " << nnfile << ", expected " << nn << " in file \'" << fullname << "\'" << std::endl;
				return;
			}
			ifs.read( (char*)&this->m_var_array[ilvl][0], nn );
			ifs.read( (char*)&nn, sizeof(unsigned) );
		}
		ifs.close();
	}

};


/**************************************************************************************\
\**************************************************************************************/

/*!
 * @class data
 * @brief encapsulates hydro data from a RAMSES simulation snapshot
 *
 * This class provides low-level read access to RAMSES hydro_XXXXX.out files.
 * Data from a given list of computational domains can be read and is
 * stored in internal datastructures.
 * Access to cell position and threaded tree structure of the cell is provided
 * through the member functions of class RAMSES_amr_level.
 * @sa RAMSES_amr_level
 */
template< typename TreeType_, typename Real_=double >
class data : public proto_data<TreeType_,Real_>{

public:
	struct header{
		unsigned ncpu;		//!< number of CPUs in simulation
		unsigned nvar;		//!< number of hydrodynamic variables
		unsigned ndim;		//!< number of spatial dimensions
		unsigned nlevelmax;	//!< maximum allowed refinement level
		unsigned nboundary;	//!< number of boundary regions
		double gamma;	//!< adiabatic exponent
	};


	std::string 	m_fname;		//!< the file name	
	struct header	m_header;	 	//!< header meta data

	const unsigned m_nvars; //!< number of variables stored in file

	std::vector<std::string> m_varnames;	//!< names of the variables stored in file
	std::map<std::string,unsigned> m_var_name_map; //!< a hash table for variable name to internal variable index


protected:

	//! generates a hydro_XXXX filename for specified cpu
	std::string gen_fname( int icpu );

	//! generate hydro_XXXXX filename from info filename
	std::string rename_info2hydro( const std::string& info );

	//! generate hydro_XXXXX filename from amr filename
	std::string rename_amr2hydro( const std::string& info );

	//! read header data containing meta information
	void read_header( void );

	//! get internal index for given variable string identifier
	/*!
	 * @param varname the string identifier of the hydro variable
	 * @return internal variable index
	 */
	int get_var_idx( const std::string& varname )
	{
		int ivar;
		std::map<std::string,unsigned>::iterator mit;
		if( (mit=m_var_name_map.find(varname)) != m_var_name_map.end() )
			ivar = (*mit).second;
		else
			throw std::runtime_error("RAMSES::HYDRO::data::get_var_idx :"\
				" Error, cannot find variable named \'"+varname+"\'");
		return ivar;
	}

	//! perform read operation of one hydro variable (internal use)
	/*!
	 * users should always call the read( std::string ) member function
	 * and read variables through their string identifiers
	 * @param var the index of the hydro variable
	 */
	void read( unsigned var );

public:

	//! constructor for hydro data
	/*!
	 * @param AMRtree reference to the underlying tree object
	 */
	explicit data( TreeType_& AMRtree )
	: proto_data<TreeType_,Real_>( AMRtree ),
	  m_fname( rename_amr2hydro(AMRtree.m_fname) ),
		m_nvars( 6 )
	{
		read_header();

		if( this->m_cpu > m_header.ncpu || this->m_cpu < 1 )
			throw std::runtime_error("RAMSES::HYDRO::data : expect to read from out of range CPU.");

		if( this->m_minlevel < 0 || this->m_maxlevel >= m_header.nlevelmax )
			throw std::runtime_error("RAMSES::HYDRO::data : requested level is invalid.");

		//m_twotondim = (unsigned)(pow(2,m_header.ndim)+0.5);

		for( unsigned i=0; i<m_nvars; ++i ){
                std::string tmp(ramses_hydro_variables[i]);
				m_var_name_map.insert( std::pair<std::string,unsigned>( tmp, i+1 ) );
				m_varnames.push_back( tmp );
			}
	}

	//! retrieve the names of variables available in the hydro data source
	/*! Variables have an internal index but are accessed through their name, given as a string
	 *  hydro data files may contain different variables depending on the kind of simulation run.
	 * @param names An output iterator to which std::string designating available variables are sent
	 * @return Final position of the output iterator
	 */
	template< typename _OutputIterator >
	_OutputIterator get_var_names( _OutputIterator names )
	{
		std::vector<std::string>::iterator it( m_varnames.begin() );
		while( it != m_varnames.end() ){
			*names = *it;
			++it; ++names;
		}
		return names;
	}

	//! perform read operation of one hydro variable
	/*!
	 * @param varname the string identifier of the hydro variable
	 */
	void read( std::string varname )
	{	this->read( get_var_idx( varname ) );  }

};


template< typename TreeType_, typename Real_ >
void data<TreeType_,Real_>::read_header( void )
{
	FortranUnformatted ff( gen_fname( this->m_cpu ) );

	//-- read header data --//
	ff.read( m_header.ncpu );
	ff.read( m_header.nvar );
	ff.read( m_header.ndim );
	ff.read( m_header.nlevelmax );
	ff.read( m_header.nboundary );
	ff.read( m_header.gamma );
}


template< typename TreeType_, typename Real_ >
void data<TreeType_,Real_>::read( unsigned var )
{
	this->m_var_array.clear();
	//int twotondim = (int)(pow(2,m_header.ndim)+0.5);

	FortranUnformatted ff( gen_fname( this->m_cpu ) );

	//.. skip header entries ..//
	ff.skip_n_from_start( 6 ); //.. skip header

	if( var < 1 || var > m_header.nvar )
		throw std::runtime_error("RAMSES::HYDRO::data::read : requested variable is invalid in file \'"+m_fname+"\'.");

	this->m_var_array.clear();

	for( unsigned ilvl = 0; ilvl<=this->m_maxlevel; ++ilvl ){

		this->m_var_array.push_back( std::vector<Real_>() );

		for( unsigned icpu = 0; icpu<m_header.ncpu+m_header.nboundary; ++icpu ){

				unsigned file_ilevel, file_ncache;
				ff.read(file_ilevel);
				ff.read(file_ncache);

				if( file_ncache == 0 )
						continue;

			if( ilvl >= this->m_minlevel ){
				if( file_ilevel != ilvl+1 )
					throw std::runtime_error("RAMSES::HYDRO::data::read : corrupted file " \
						 "or file seek failure in file \'"+m_fname+"\'.");


				std::vector<float> tmp;
				for( unsigned i=0; i<this->m_twotondim; ++i )
				{
					ff.skip_n( var-1 );
					ff.read<double>( std::back_inserter(tmp) );
					ff.skip_n( m_header.nvar-var );
				}
				//.. reorder array to increase data locality..//
				this->m_var_array.reserve( tmp.size() );
				for( unsigned i=0; i<file_ncache; ++i ){
					for( unsigned j=0; j<this->m_twotondim; ++j ){
						(this->m_var_array.back()).push_back( tmp[i+j*file_ncache] );
					}
				}
			}else{
				ff.skip_n( this->m_twotondim*m_header.nvar );
			}
		}
	}
}


template< typename TreeType_, typename Real_ >
std::string data<TreeType_,Real_>::gen_fname( int icpu )
{
	std::string fname;
	char ext[32];
	fname = m_fname;
	fname.erase(fname.rfind('.')+1);
	sprintf(ext,"out%05d",icpu);
	fname.append(std::string(ext));
	return fname;
}


template< typename TreeType_, typename Real_ >
std::string data<TreeType_,Real_>::rename_info2hydro( const std::string& info )
{
	std::string amr;
	unsigned ii = info.rfind("info");
	amr = info.substr(0,ii)+"hydro" + info.substr(ii+4, 6) + ".out00001";
	return amr;
}


template< typename TreeType_, typename Real_ >
std::string data<TreeType_,Real_>::rename_amr2hydro( const std::string& info )
{
	std::string amr;
	unsigned ii = info.rfind("amr");
	amr = info.substr(0,ii)+"hydro" + info.substr(ii+3, 6) + ".out00001";
	return amr;
}

/**************************************************************************************\
\**************************************************************************************/


/*! 
 * @class RAMSES::HYDRO::multi_domain_data
 * @brief encapsulates hydro data from multiple computational domain
 *
 * This class provides high-level data access to bundled domain data, Bundling
 * of domain data is useful when analysis of snapshots is performed in parallel
 * but on a number of cores different than that used for the RAMSES simulation.
 */
template< typename TreeType_, typename DataType_, typename ValueType_=double >
class multi_domain_data
{
	public:
	
		unsigned m_ndomains;				//!< number of domains bundled
		RAMSES::snapshot &m_rsnap;			//!< reference to the underlying snapshot object
		std::vector<DataType_*> m_data;		//!< vector of bundled data objects
		std::vector<TreeType_*> m_ptrees;	//!< vector of bundled tree objects
		
	public:
	
		//! constructor for bundled multi-domain hydro variables
		/*!
		 * @param rsnap reference to the underlying snapshot object
		 * @param ptrees vector of trees to be bundled
		 */
		multi_domain_data( RAMSES::snapshot& rsnap, std::vector<TreeType_*> ptrees )
		: m_ndomains( ptrees.size() ), m_rsnap(rsnap), m_ptrees( ptrees )
		{ 
			for( unsigned idom=0; idom<m_ndomains; ++idom )
				m_data.push_back( new DataType_(*m_ptrees[idom]) );
			
		}
		
		
		//! combines all elements of this instance with that of another using a binary operator
		/*!
	 	 * @param o  the other multi-domain data object with which to combine the elements
	 	 * @param op the C++ binary operator to be used in combining elements
	 	 */
		template<typename BinaryOperator, typename OtherDataType_>
		void combine_with( const multi_domain_data<TreeType_,OtherDataType_,ValueType_>& o, const BinaryOperator& op )
		{
			if( m_data.size() != o.m_data.size() )
				throw std::runtime_error("Error: trying to combine incompatible multi_domain_data.");
			
			for( unsigned i=0; i<m_data.size(); ++i ) //<OtherDataType_,BinaryOperator>
				m_data[i]->combine_with( *o.m_data[i], op );
		}
		
		
		//! bracket operator to access bundled multi-domain data
		/*!
		 * @param idomain the domain index (1 based)
		 * @param it tree iterator pointing to the grid in the tree to be accessed
		 * @param CellIndex index of the cell in the grid
		 */
		ValueType_& operator() ( unsigned idomain, const typename TreeType_::iterator& it, char CellIndex )
		{
			if( idomain >= m_data.size() ){
				std::cerr << "out of range domain: " << idomain << " >= " << m_data.size() << std::endl;
				idomain=0;
			}
			
			return m_data[idomain]->cell_value(it,CellIndex);
		}
		
		
		//! bracket operator to access bundled multi-domain data (const)
		/*!
		 * @param idomain the domain index (1 based)
		 * @param it tree const_iterator pointing to the grid in the tree to be accessed
		 * @param CellIndex index of the cell in the grid
		 */
		ValueType_& operator() ( unsigned idomain, const typename TreeType_::const_iterator& it, char CellIndex )
		{
			if( idomain >= m_data.size() ){
				std::cerr << "out of range domain: " << idomain << " >= " << m_data.size() << std::endl;
				idomain=0;
			}
			return m_data[idomain]->cell_value(it,CellIndex);
		}
		
		//! destructor for bundled multi-domain data
		~multi_domain_data()
		{
			for( unsigned idom=0; idom<m_data.size(); ++idom )
				delete m_data[idom];
		}
		
		
		//! bundled read functions, reads in multi-domain data
		void get_var( std::string var_name )
		{
			//m_data.clear();
			if( m_ndomains != m_data.size() ){
				std::cerr << "Error: Internal consistency check failed in multi_domain_data::get_var.";
				return;
			}
			
			for( unsigned idom=0; idom<m_ndomains; ++idom )
			{
				//m_data.push_back( new DataType_(*m_ptrees[idom]) );
				m_data[idom]->read(var_name);
			}
		}

		//! TBD
		unsigned size( void ){ return m_data.size(); }
		
		//! TBD
		unsigned size( unsigned idom ) { return m_data.at(idom)->size(); }
};

/**************************************************************************************\
\**************************************************************************************/


} // namespace HYDRO
} // namespace RAMSES

#endif //__RAMSES_HYDRO_DATA_HH
