/*
		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_PARTICLE_DATA_HH
#define __RAMSES_PARTICLE_DATA_HH

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

#include "RAMSES_info.hh"
#include "FortranUnformatted_IO.hh"
#include "data_iterators.hh"

namespace RAMSES{
namespace PART{

//! particle type identifiers returned by RAMSES_particle_data::ptype()
enum ptype_bits {
	RAMSES_DM_BIT     = 0,
	RAMSES_STAR_BIT   = 1,
	RAMSES_DEBRIS_BIT = 2
};

//! particle type identifiers, can be combined with logical or
enum ptype { 
	ptype_dm     = 1L << RAMSES_DM_BIT,			//!< selects dark matter particles
	ptype_star   = 1L << RAMSES_STAR_BIT,		//!< selects star particles
	ptype_debris = 1L << RAMSES_DEBRIS_BIT,		//!< selects debris particles
};

//... zero age of a particle is assumed if age is smaller than this number
const double zero_age      = 1.0e-8;


//... content of RAMSES particle data files is fixed, available fields
//... are described in the following two fields

		
//! possible variable names in RAMSES particle files in version 1 and 2 files
const char ramses_particle_variables_v12[][64] = {
	{"position_x"},
	{"position_y"},
	{"position_z"},
	{"velocity_x"},
	{"velocity_y"},
	{"velocity_z"},
	{"mass"},
	{"age"},
	{"metallicity"},
	{"particle_ID"},
	{"refinement_level"} };
	
//! possible variable names in RAMSES particle files in version 3 files
const char ramses_particle_variables_v3[][64] = {
	{"position_x"},
	{"position_y"},
	{"position_z"},
	{"velocity_x"},
	{"velocity_y"},
	{"velocity_z"},
	{"mass"},
	{"particle_ID"},
	{"refinement_level"},
	{"age"},
	{"metallicity"} };
	

/**************************************************************************************\
\**************************************************************************************/
	
/*! 
 * @class RAMSES::PARTICLE::data
 * @brief encapsulates particle data from a RAMSES simulation snapshot
 *
 * This class provides low-level read access to RAMSES particle files. 
 * Data from a given list of computational domains can be read and is
 *stored in internal datastructures.
 */
class data {
protected:
	const RAMSES::snapshot *m_pRAMSESsnap;	//!< pointer to RAMSES snapshot meta-information
	std::string m_fname;					//!< path/filename of the .info file
	int m_cpu;								//!< IDs of computational domains from which data is to be read.
	unsigned m_nvars;						//!< number of variables available in file
	std::vector<unsigned> m_nvar_stride; 	//!< access offset to variable locations in file
	std::vector<int> m_vardim;				//!< spatial dimension of a variable
	std::vector<std::string> m_varnames;	//!< names of the variables

	//! a hash table for variable name to internal variable index
	std::map<std::string,unsigned> m_var_name_map;
	
	//! particle file header meta data
	struct header{ 
		std::vector< int > localseed;		//!< local seed used to generate star particles
		int ncpu;							//!< number of CPUs in the simulation
		int ndim;							//!< number of spatial dimensions
		int nstar_tot;						//!< total number of star particles
		int nsink;							//!< number of sink particles
		int npart;							//!< number of particles
		int npart_dm;						//!< number of dark matter particles
		int npart_star; 					//!< number of star particles
		int npart_debris;					//!< number of debris particles
		double mstar_tot;					//!< total stellar mass
		double mstar_lost;					//!< stellar mass lost
	};
	
	struct header m_header;					//!< particle file header data
	
	
	//! generate the part_XXXXX.oYYYYY filename for a given CPU ID
	std::string gen_fname( int icpu );
	
	//! generate the part_XXXXX.oYYYYY filename from the info_XXXXX.oYYYYY filename.
	std::string rename_info2part( const std::string& info );
	
	//! Read header from particle files and determine total number of particles.
	void read_header( void );
	
public:
	
	//! constructor for a RAMSES particle data interface
	/*! The constructor checks the validity of the requested computational
	 *  domains and reads the particle header data to determine the total
	 *  number of particles.
   	 * @param snap reference to the underlying snapshot object
	 * @param cpu the domain number
     */
	data( const RAMSES::snapshot& snap, int cpu )
		: m_pRAMSESsnap( &snap ), 
		  m_fname( rename_info2part(snap.m_filename) ), 
		  m_cpu( cpu )
    { 
		m_header.npart_dm     = 0;
		m_header.npart_star   = 0;
		m_header.npart_debris = 0,
			
		read_header();
		
		//... compute offsets of variable arrays in file and access hash map ...//
		m_nvar_stride.push_back(0);
		
		
		switch( m_pRAMSESsnap->m_version )
		{
			case RAMSES::version1:
			case RAMSES::version2:
		
				if( m_header.nstar_tot > 0 )
				{
					//... file contains star particles ...//
					m_nvars = 11;
					for( unsigned i=0; i<m_nvars; ++i ){
						m_nvar_stride.push_back(m_nvar_stride[i]+1);
						m_var_name_map.insert( std::pair<std::string,unsigned>( ramses_particle_variables_v12[i], i ) );
						m_varnames.push_back( ramses_particle_variables_v12[i] );
					}
				}
				else
				{
					//... file contains only DM particles, then some fields are not available
					m_nvars = 9;
					for( unsigned i=0,j=0; i<m_nvars; ++i,++j ){
						if( i==7 ) j=9;
						m_nvar_stride.push_back(m_nvar_stride[i]+1);
						m_var_name_map.insert( std::pair<std::string,unsigned>( ramses_particle_variables_v12[j], i ) );
						m_varnames.push_back( ramses_particle_variables_v12[j] );
					}
				}
				
				break;
				
			case RAMSES::version3:
			
				if( m_header.nstar_tot > 0 )
					//... file contains star particles ...//
					m_nvars = 11;
				else
					//... file contains only DM particles, then last fields are not available
					m_nvars = 9;

				for( unsigned i=0; i<m_nvars; ++i )
				{
					m_nvar_stride.push_back(m_nvar_stride[i]+1);
					m_var_name_map.insert( std::pair<std::string,unsigned>( ramses_particle_variables_v3[i], i ) );
					m_varnames.push_back( ramses_particle_variables_v3[i] );
				}
			
				break;
				
			
			default:
				throw std::runtime_error("RAMSES::PART::data: cannot handle data for this RAMSES version.");
		}
		
		//... check if supplied range of domains is valid ...//
		if( m_cpu < 1 || m_cpu > m_header.ncpu )
			throw std::runtime_error("RAMSES::PART::data: attempt to read from out of range CPU.");
	}
	
	
	//=== implementation of interface derived compatible to multi variable data source skeleton ===//
	//=== see file README.devel, section 1, for details on the interface skeleton               ===//
	
	//! retrieve the names of variables available in the particle data source
	/*! Variables have an internal index but are accessed through their name, given as a string
	 *  Particle 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;
	}
	
protected:

	//! return the internal index for a variable of specified name
	/*! throws a runtime_error exception if variable name does not exist
	 * @param varname string designating the variable
	 * @return internal index of variable
	 */
	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::PART::data::get_var_idx :"\
				" Error, cannot find variable named \'"+varname+"\'");
		return ivar;
	}
	
	
	//! get spatial dimensionality of a variable
	int get_var_dim( int ivar )
	{	return m_vardim[ivar]; }
	
	
public:
	
	//! get spatial dimensionality of a variable
	int get_var_dim( const std::string& varname )
	{ return m_vardim[get_var_idx(varname)]; }
	
	
	//! retrieve data of specified variable from data source using a mask access pattern
	/*! The data for the specified variable, fixing a spatial dimension, is sent to an
	 *  output iterator. For each read operation, the mask iterator is checked
	 *  and the data is discarded if it points to 'false'
	 * @param varname name of the variable to be retrieved
	 * @param mask input iterator supplying the mask access pattern
	 * @param val output iterator to which the retrieved data is sent
	 * @return final position of the mask iterator
	 */
	template< typename _Basetype, typename _InputIterator, typename _OutputIterator >
	_InputIterator get_var( const std::string& varname, _InputIterator mask, _OutputIterator val )
	{ 
		unsigned ivar = get_var_idx( varname );
		FortranUnformatted ff( gen_fname(m_cpu) );
		//.. skip header and particle position entries ..//
		ff.skip_n_from_start( 8+m_nvar_stride[ivar] );
		mask=ff.read<_Basetype, _InputIterator, _OutputIterator>(mask,val);

		return mask;
	}
	
	
	//! retrieve data of specified variable from data source
	/*! The data for the specified variable, fixing a spatial dimension, is sent to an
	 *  output iterator. 
	 * @param varname name of the variable to be retrieved
	 * @param val output iterator to which the retrieved data is sent
	 * @return final position of the mask iterator
	 */
	template< typename _Basetype, typename _OutputIterator >
	_OutputIterator get_var( const std::string& varname, _OutputIterator val )
	{ 
		unsigned ivar = get_var_idx( varname );
		FortranUnformatted ff( gen_fname(m_cpu) );
		//.. skip header and particle position entries ..//
		ff.skip_n_from_start( 8+m_nvar_stride[ivar] );
		ff.read<_Basetype, _OutputIterator>(val);
		
		return val;
	}
	
	
	
	//=== the following member functions are simply copied from the skeleton ===//
	
		
	//! generate bit mask for specified particle type
	/*! returns a particle type identifier for the given particle index
	 *  this is determined from the age and ID of the particle. Return 
	 *  value is RAMSES_PTYPE_DM for dark matter particles (age=0,ID>0),
	 *  RAMSES_PTYPE_STAR for star particles (age>0,ID>0) and 
	 *  RAMSES PTYPE_DEBRIS for debris type (SN remnant) particles 
	 *  (age>0,ID=0).
	 * @param ptype particle type, i.e. RAMSES_PTYPE_DM, RAMSES_PTYPE_STAR or RAMSES_PTYPE_DEBRIS
	 * @param age_first
	 * @param age_last
	 * @param ids_first
	 * @param mask
	 * @return particle mask iterator
	 */
	template<typename _InputIterator1, typename _InputIterator2, typename _OutputIterator>
	_OutputIterator mask_particle_type( int ptype, 
										const _InputIterator1& age_first, const _InputIterator1& age_last,
										const _InputIterator2& ids_first, _OutputIterator mask )
	{
		_InputIterator1 age_it=age_first;
		_InputIterator2 ids_it=ids_first;
	
		switch( ptype )
		{
			case ptype_dm:	
				while( age_it != age_last )
				{
					if( fabs(*age_it) <= zero_age )
						*mask = true;
					else
						*mask = false;
					++mask;	++age_it;
				}
				break;
			
			case ptype_star:	
				while( age_it != age_last )
				{
					if( fabs(*age_it) > zero_age && (*ids_it) > 0)
						*mask = true;
					else
						*mask = false;
					++mask;	++age_it;	++ids_it;
				}
				break;
								
			case ptype_debris:	
				while( age_it != age_last )
				{
					if( fabs(*age_it) > zero_age && (*ids_it) == 0)
						*mask = true;
					else
						*mask = false;
					++mask;	++age_it;	++ids_it;
				}
				break;
								
			default: throw std::runtime_error("get_particle_type_mask: Error, invalid particle type.");
		}
		
		return mask;
	}
	
	
	
	
	//! generate bit mask for a specified bounding region
	/*! Checks whether data points lie within a given region. The _BoundingRegion class object
	 *  must return true or false for a call of 'region(x,y,z)' if the point (x,y,z)
	 *  does or does not lie in the desired region.
	 * @param region A region object used for querying whether points lie in a desired subspace.
	 * @param x_first Input iterator pointing on first x-coordinate
	 * @param x_last Input iterator pointing behind last x-coordinate
	 * @param y_first Input iterator pointing on first y-coordinate
	 * @param z_first Input iterator pointing on first z-coordinate
	 * @param mask Output iterator to which true/false values are written
	 * @return Final state of mask iterator
	 */
	template<typename _InputIterator, typename _BoundingRegion, typename _OutputIterator>
	_OutputIterator mask_region( _BoundingRegion region,
								 const _InputIterator& x_first, const _InputIterator& x_last,
								 const _InputIterator& y_first, const _InputIterator& z_first,
								 _OutputIterator mask )
	{		
		_InputIterator xit=x_first, yit=y_first, zit=z_first;
		
		while( xit != x_last )
		{
			*mask = region( *xit, *yit, *zit );
			++mask; ++xit; ++yit; ++zit;
		}
		
		return mask;
	}
	
};

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

#define R_CHECK_BIT(var,pos) ((var) & (1<<(pos)))

template< typename Real_, typename IDType_ >
inline bool is_of_type( Real_ age, IDType_ id, int ptype )
{
	//if( ptype==ptype_star )
	if( R_CHECK_BIT( ptype, RAMSES_STAR_BIT ) )
		if( fabs(age)>zero_age && id > 0 )
			return true;
		//else return false;

	//if( ptype==ptype_dm )
	if( R_CHECK_BIT( ptype, RAMSES_DM_BIT ) )
		if( fabs(age)<=zero_age )
			return true;
		//else return false;
			
	//if( ptype==ptype_debris )
	if( R_CHECK_BIT( ptype, RAMSES_DEBRIS_BIT ) )
		if( fabs(age) > zero_age && id == 0 )
			return true;
	return false;
}

#undef R_CHECK_BIT

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

inline std::string data::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;
}

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

inline void data::read_header( void )
{
	FortranUnformatted ff( gen_fname(m_cpu) );	
	
	//-- read header data --//
	ff.read( m_header.ncpu );
	ff.read( m_header.ndim );
	ff.read( m_header.npart );
	ff.read<int>( std::back_inserter(m_header.localseed) );
	ff.read( m_header.nstar_tot );
	ff.read( m_header.mstar_tot );
	ff.read( m_header.mstar_lost );
	ff.read( m_header.nsink );
	
}

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

inline std::string data::rename_info2part( const std::string& info )
{
	std::string amr;
	unsigned ii = info.rfind("info");
	amr = info.substr(0,ii)+"part" + info.substr(ii+4, 6) + ".out00001";
	return amr;
}

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

/*! 
 * @class RAMSES::PART::multi_domain_data
 * @brief encapsulates particle 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 ValueType_=double >
class multi_domain_data
{
	protected:
		unsigned m_ndomains;							//!< number of domains bundled
		RAMSES::snapshot &m_rsnap;						//!< reference to the underlying snapshot object
		std::vector< std::vector<ValueType_> > m_data;	//!< vector of bundled data objects
		std::vector<TreeType_*> m_ptrees;				//!< vector of bundled tree objects
		
	public:
	
		//! constructor for bundled multi-domain particle data
		/*!
		 * @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 )
		{ }
		
		
		//! access a particle in a specific domain
		/*!
		 * @param idomain the domain ID
		 * @param ind the particle index
		 */
		ValueType_& operator() ( unsigned idomain, unsigned ind )
		{
			return m_data[idomain][ind];
		}
		
		//! access the vector of particle of a specific domain
		/*!
		 * @param idomain the domain ID
		 */
		std::vector<ValueType_>& operator() ( unsigned idomain )
		{
			return m_data[idomain];
		}
		
		//! bundled read functions, reads in multi-domain data
		void get_var( std::string var_name )
		{
			m_data.clear();
			for( unsigned idom=0; idom<m_ndomains; ++idom )
			{
				data local_data( m_rsnap, m_ptrees[idom]->m_cpu );
				m_data.push_back( std::vector<ValueType_>() );
				local_data.get_var<ValueType_>(var_name, std::back_inserter(m_data.back()) );
			}
		}
		
		//! get the pointer to the vector of particles of a specific domain
		std::vector<ValueType_>* get_vpointer( unsigned idomain )
		{
			return &m_data[idomain];
		}
	
		//! TBD
		unsigned size( void ){ return m_data.size(); }
		
		//! TBD
		unsigned size( unsigned idom ) { return m_data.at(idom).size(); }
};


}// namespace PARTICLE

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

namespace GEOM{

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

//! TBD
class bounding_sphere{
	protected:
		double 
			xc,	//!< TBD
			yc,	//!< TBD
			zc,	//!< TBD
			r,	//!< TBD
			r2;	//!< TBD
			
	public:
		//! TBD
		bounding_sphere( double _xc, double _yc, double _zc, double _r )
		: xc(_xc), yc(_yc), zc(_zc), r(_r), r2(r*r)
		{ }
		
		//! TBD
		inline bool operator()( double x, double y, double z )
		{
			double dx(x-xc),dy(y-yc),dz(z-zc);
			return dx*dx+dy*dy+dz*dz < r2;
		}
};

/**************************************************************************************\
\**************************************************************************************/
//! TBD
class bounding_box{
	protected:
		double 
			xc,		//!< TBD
			yc,		//!< TBD
			zc,		//!< TBD
			Lx,		//!< TBD
			Ly,		//!< TBD
			Lz,		//!< TBD
			Lx2,	//!< TBD
			Ly2,	//!< TBD
			Lz2;	//!< TBD
	public:
		//! TBD
		bounding_box( double _xc, double _yc, double _zc, double _Lx, double _Ly, double _Lz )
		: xc(_xc), yc(_yc), zc(_zc), Lx(_Lx), Ly(_Ly), Lz(_Lz), Lx2(0.5*Lx), Ly2(0.5*Ly), Lz2(0.5*Lz)
		{ }
		
		//! TBD
		inline bool operator()( double x, double y, double z )
		{
			double dx(x-xc),dy(y-yc),dz(z-zc);
			return fabs(dx)<Lx2 && fabs(dy)<Ly2 && fabs(dz)<Lz2; 
		}
};

/**************************************************************************************\
\**************************************************************************************/
//! TBD
class bounding_cube:public bounding_box{
	public:
		//! TBD
		bounding_cube( double _xc, double _yc, double _zc, double _L )
		: bounding_box( _xc, _yc, _zc, _L, _L, _L )
		{ }
};

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

}// namespace GEOMETRY

}// namespace RAMSES

#endif //__RAMSES_PARTICLE_DATA
