/*
 * GuessFactory.h
 *
 *  Created on: Sep 27, 2018
 *      Author: kottmanj
 *
 *  Guess Factory for TDHF which can be used standalone for other purposes
 */

#ifndef SRC_APPS_CHEM_GUESSFACTORY_H_
#define SRC_APPS_CHEM_GUESSFACTORY_H_

#include <madness.h>

namespace madness {

// avoid confusion with other projects which may use similar generic names
namespace guessfactory{

// convenience macros
#define SPLITCOORD(x,y,z,r) double x=r[0]-origin[0]; double y=r[1]-origin[1];double z=r[2]-origin[2];

/// compute the centroid of a function i.e. c[xi]=<f|xi|f>/<f|f> i.e. position expectation value
coord_3d compute_centroid(const real_function_3d& f);

template<typename T, std::size_t NDIM>
std::vector<coord_3d> compute_centroids(const std::vector<Function<T,NDIM> > & vf);

/// little helper for coord (Vector<3>) and Tensor data formats
template<typename T, size_t NDIM>
Vector<T,NDIM> tensor_to_coord(const Tensor<T>& t) {
	Vector<T,NDIM> result;
	MADNESS_ASSERT(size_t(t.size()) >= NDIM);
	for (size_t i = 0; i < NDIM; ++i)
		result[i] = t[i];
	return result;
}



/// create excitation operators with unaryop (faster as explicit construction and multiplication)
/// Guess function do not need to be perfectly refined
class ExopUnaryOpStructure {
public:
	ExopUnaryOpStructure(const std::shared_ptr<FunctionFunctorInterface<double, 3> >& f) :
			exfunc(f), cdata(FunctionCommonData<double, 3>::get(FunctionDefaults<3>::get_k())) {
	}

	/// multiplies the target function by the excitation operator defined by exfunc
	void operator ()(const Key<3>& key, Tensor<double>& t) const;
	void operator ()(const Key<3>& key, Tensor<double_complex>& t) const;
    /// shared pointer to object of excitation operator
    std::shared_ptr<FunctionFunctorInterface<double,3> >  exfunc;
    FunctionCommonData<double,3> cdata;
    template <typename Archive> void serialize(Archive& ar) {}
};

/// creates a plane-wave: sin (or cos) with argument (npi/L*x)
class PlaneWaveFunctor : public FunctionFunctorInterface<double,3> {
public:
	PlaneWaveFunctor(std::vector<double> vn,std::vector<bool> vc, const coord_3d& c) : L(FunctionDefaults<3>::get_cell_width()), n(vn), cosinus(vc), origin(c) {}

	typedef double resultT;

    /// for explicit construction of this plane-wave function
	double operator ()(const coord_3d& r) const;
    /// operator for the 1D plane waves
	double operator ()(const double& x, const int& dim) const;
    /// in case this is needed at some point
    double operator()(const coord_1d & x, const int& dim)const{
    	return (*this)(x[0],dim);
    }

	std::string name(const bool& compact = false) const;

    const Tensor<double> L;
    const std::vector<double> n;
    const std::vector<bool> cosinus;
    const coord_3d origin;


};

/// GaussFunctor to let the exciation operators go to zero at the boundaries
/// totally symmetric
class GaussFunctor : public FunctionFunctorInterface<double,3> {
public:
	GaussFunctor();
	GaussFunctor(const double& width): width_(width){
		MADNESS_ASSERT(not(width<0.0));
	}
	GaussFunctor(const double& width, const coord_3d c): width_(width), center(c){
		MADNESS_ASSERT(not(width<0.0));
	}
	GaussFunctor(const double& width, const Tensor<double> c): width_(width), center(tensor_to_coord<double,3>(c)){
		MADNESS_ASSERT(not(width<0.0));
	}
	const double width_;
	const coord_3d center=coord_3d();

	/// explicit construction
	double operator ()(const coord_3d& rr) const;


};

/// Project a general 3D polynomial to the MRA Grid
class PolynomialFunctor : public FunctionFunctorInterface<double,3> {
public :
	/// simple xyz moments constructor
	PolynomialFunctor(const int& axis): input_string_(axis_to_string(axis)), data_(read_string(axis_to_string(axis))), dampf(0.0) {}
	// general polynomials or sums of polynomials
	PolynomialFunctor(const std::string input, const double& damp_width=0.0, const coord_3d& c=coord_3d()) : input_string_(input), data_(read_string(input)), dampf(damp_width), center(c) {}
	PolynomialFunctor(const std::string input,const double& damp_width, const Tensor<double>& c) : input_string_(input), data_(read_string(input)), dampf(damp_width), center(tensor_to_coord<double,3>(c)) {}

	/// construction by coordinates
	double operator ()(const coord_3d& rr) const;

	/// create the value of the polynomial according to the data in the data_ structure
	double compute_value(const coord_3d& r) const;

	/// convert a given axis to the appropriate input string
	std::string axis_to_string(const int& axis)const{
		std::string result;
		if(axis==0) result="x 1.0";
		else if(axis==1) result="y 1.0";
		else if (axis==2) result="z 1.0";
		else MADNESS_EXCEPTION("polynomial functor only defined up to 3 dimensions",1);
		return result;
	}

protected:
	const std::string input_string_;
	/// The data for the construction of the polynomial chain
	/// every entry of data_ is vector containing the threee exponents and the coefficient of a monomial dx^ay^bz^c , data_[i] = (a,b,c,d)
	const std::vector<std::vector<double>> data_;
	/// damping function
	GaussFunctor dampf;
	coord_3d center=coord_3d();
public:
	std::vector<std::vector<double> > read_string(const std::string string) const;
	void test();
	std::vector<std::vector<double> > give_data(){return data_;}
};

/// instead of x,y,z use sin(x), sin(y), sin(z)
/// shows the same transformation behaviour but does not grow unbounded with larger x,y,z values
class PolynomialTrigonometricsFunctor : public PolynomialFunctor {
public:
	/// c++11 constructor inheritance
	using PolynomialFunctor::PolynomialFunctor;
	// overload
	/// create the value of the polynomial according to the data in the data_ structure
	/// instead of x,y,z use sin(x), sin(y), sin(z)
	double compute_value(const coord_3d& r) const;
};



/// excite a vector of functions with a specific excitation operator
/// @param[in/out] vf the function which gets excited, exop*f on return
/// @param[in] exop_input , the excitation operator defined by a string (see the polynomial_functor class for details)
/// @return exop*vf i.e. result[i]=exop*vf[i]
vector_real_function_3d apply_polynomial_exop(vector_real_function_3d& vf, const std::string& exop_input, std::vector<coord_3d> centers = std::vector<coord_3d>(), const bool& fence = false);
/// convenience wrapper
real_function_3d apply_polynomial_exop(real_function_3d& f, const std::string& exop_input, coord_3d center = coord_3d(), const bool& fence = false);

/// excite a vector of functions with a specific excitation operator
/// @param[in/out] vf the function which gets excited, exop*f on return
/// @param[in] exop_input, the excitation operator defined by a string (see the polynomial_functor class for details)
/// @param[in] the centers of the vf functions, if none were given they are recomputed
/// @return exop*vf i.e. result[i]=exop*vf[i]
//template<typename T, std::size_t NDIM>
//std::vector<Function<T,NDIM> > apply_trigonometric_exop(std::vector<Function<T,NDIM> >& vf,
//		const std::string& exop_input, std::vector<coord_3d> centers = std::vector<coord_3d>(), const bool& fence = false);
/// excite a vector of functions with a specific excitation operator
/// @param[in/out] vf the function which gets excited, exop*f on return
/// @param[in] exop_input, the excitation operator defined by a string (see the polynomial_functor class for details)
/// @param[in] the centers of the vf functions, if none were given they are recomputed
/// @return exop*vf i.e. result[i]=exop*vf[i]
template<typename T, std::size_t NDIM>
std::vector<Function<T,NDIM> > apply_trigonometric_exop(std::vector<Function<T,NDIM> >& vf,
		const std::string& exop_input, std::vector<coord_3d> centers, const bool& fence) {
	if (vf.empty()) return vf;

	//recompute centers if necessary
	if (centers.empty()) centers = compute_centroids(vf);

	ExopUnaryOpStructure exop(std::make_shared<PolynomialTrigonometricsFunctor>(PolynomialTrigonometricsFunctor(exop_input)));
	for (auto& f : vf) f.unaryop(exop, false);
	if (fence) vf.front().world().gop.fence();
	return vf;
}


/// convenience wrapper
template<typename T, std::size_t NDIM>
Function<T,NDIM> apply_trigonometric_exop(Function<T,NDIM>& f, const std::string& exop_input, coord_3d center, const bool& fence) {
	std::vector<Function<T,NDIM> > vf(1, f);
	std::vector<coord_3d> centers(1, center);
	return apply_trigonometric_exop(vf, exop_input, centers, fence).front();
}
//template<typename T, std::size_t NDIM>
//std::vector<Function<T,NDIM> > apply_trigonometric_exop(Function<T,NDIM>& f, const std::string& exop_input, coord_3d center = coord_3d(), const bool& fence = false);


/// Makes an excitation operator string based on predefined keywords
std::vector<std::string> make_predefined_exop_strings(const std::string what);
/// Makes an automated excitation operator string for the excitation operators needed to create virtuals from the reference orbitals
std::vector<std::string> make_auto_polynom_strings(const size_t order);


} /* namespace guessfactory */
} /* namespace madness */

#endif /* SRC_APPS_CHEM_GUESSFACTORY_H_ */
