/*
 * smooth.h
 *
 *  Created on: Nov 25, 2015
 *      Author: kottmanj
 */

#ifndef SRC_EXAMPLES_SMOOTH_H_
#define SRC_EXAMPLES_SMOOTH_H_

#include <iostream>
#include <string>
#include <chem/nemo.h>
#include <madness/mra/funcplot.h>
#include <cmath>
#include <chem/SCFOperators.h>
#include <madness/mra/function_common_data.h>

static double RHOMIN = 1.e-20;

double make_log(double x){
	if(x<0.0) return log(RHOMIN);
	return log(x);
}

double make_exp(double x){
	return exp(x);
}

double estimate_area(double x) {
	double thresh =FunctionDefaults<3>::get_thresh();
	if(fabs(x)<thresh) return 0.0;
	else return 1.0;
    }

static double test_1d_functor(const coord_1d &x){
	return exp(-fabs(x[0]))+0.1*sin(30.0*x[0]);
}

static double xf(const coord_3d &x){
	return x[1];
}
static double r2(const coord_3d &x){
	return x[0]*x[0]+x[1]*x[1]+x[2]*x[2];
}

static double slater_functor(const coord_3d &x){
	double rsq = r2(x);
	double r = sqrt(rsq);
	return exp(-r);
}

static double mask_factor = 5.0;
static double cutoff_radius = 5.0;
//static double mask_functor(const coord_3d &x){
//	double r = sqrt(r2(x));
//	//return 1-erf(mask_factor*r);
//	return 0.5*(1.0 - tanh(mask_factor*(r-cutoff_radius)));
//}
static double mask_functor_box(const coord_3d &x){
	double r = sqrt(r2(x));
	//return 1-erf(mask_factor*r);
	return 0.5*(1.0 - tanh(1.0*(r-15)));
}
//static double inv_mask_functor(const coord_3d &x){
//	double r = sqrt(r2(x));
//	return 0.5*(1.0 + tanh(mask_factor*(r-cutoff_radius)));
//}

struct munging_operator {
	munging_operator() : thresh_(FunctionDefaults<3>::get_thresh()) {}
	munging_operator(const double thresh) : thresh_(thresh) {}
    void operator()(const Key<3>& key,
                    real_tensor U,
                    const real_tensor& function,
                    const real_tensor& smoothed_function) const {
        ITERATOR(U,
                 double f = function(IND);
                 //double sf = smoothed_function(IND);
                 if (fabs(f)<thresh_)
                     U(IND) = 0.0;
                 else
                     U(IND) = f;
                 );
    }

    template <typename Archive>
    void serialize(Archive& ar) {}
    double thresh_;
};

struct asymptotic_density{
public:
	asymptotic_density(const double &ie):ionization_energy_(ie), center_of_charge_(0.0){}
	asymptotic_density(const double &ie, const coord_3d &coc):ionization_energy_(ie), center_of_charge_(coc){}
		double operator ()(const coord_3d &x1)const{
		coord_3d x;
		x[0] = x1[0]-center_of_charge_[0];
		x[1] = x1[1]-center_of_charge_[1];
		x[2] = x1[2]-center_of_charge_[2];
		double r = sqrt((x[0]*x[0]+x[1]*x[1]+x[2]*x[2]));
		double result= exp(-2.0*sqrt(2.0*fabs(ionization_energy_))*r);
		return result;
	}

	const double ionization_energy_;
	const coord_3d center_of_charge_;
};

struct asymptotic_slater_kernel{
public:
	asymptotic_slater_kernel(const asymptotic_density &tmp): asymptotic_rho_(tmp){}

	double operator()(const coord_3d &x)const{
		double prefactor = 1.0/M_PI;
		double r = sqrt((x[0]*x[0]+x[1]*x[1]+x[2]*x[2]));
		double kernel = prefactor*exp(2.0/3.0*2.0*sqrt(2.0*fabs(asymptotic_rho_.ionization_energy_))*r);//pow(rho,2.0/3.0);
		double rho = asymptotic_rho_(x);
 		return kernel*rho;
	}

private:
	const asymptotic_density asymptotic_rho_;
};

struct slater_kernel {
	slater_kernel() : thresh_(FunctionDefaults<3>::get_thresh()) {}
	slater_kernel(const double thresh) : thresh_(thresh) {}
    void operator()(const Key<3>& key,
                    real_tensor U,
                    const real_tensor& rho,
                    const real_tensor& pt_rho) const {
    	double safety = thresh_*0.01;
    	double safety_kernel = pow(safety,-2.0/3.0);
    	double safety_kernel_safety_pd = pre*pow(safety,1.0/3.0);
        ITERATOR(U,
                 double d = rho(IND);
                 double pd = pt_rho(IND);
                 if(d<safety and fabs(pd)<safety) U(IND)=safety_kernel_safety_pd;
                 else if(d<safety)  U(IND)= pre*safety_kernel*pd;
                 else U(IND) = pre*pow(d,-2.0/3.0)*pd;
                 );
    }

    template <typename Archive>
    void serialize(Archive& ar) {}
    double thresh_;
    double pre = 1.0/M_PI;
};

static double analytical_slater_functor(double rho){
	double tmp;
	if(rho < 0.0) tmp = 0.0;
	else tmp = pow(rho,1.0/3.0);
	double prefactor = 1.0/M_PI;
	return prefactor*tmp;
}

struct asymptotic_slater : public FunctionFunctorInterface<double,3>{
public:
	asymptotic_slater(const double &ipot_): ipot(ipot_){}
	const double ipot;
	const double prefactor = 1.0/M_PI;

	double operator()(const coord_3d &x)const{
		const double r = sqrt((x[0]*x[0]+x[1]*x[1]+x[2]*x[2]));
		const double kernel = prefactor*exp(2.0/3.0*2.0*sqrt(2.0*fabs(ipot))*r);
		return kernel;
	}
};

struct apply_kernel_helper{

public:
	apply_kernel_helper(const double &ipot_): ipot(ipot_), asl(ipot_){}

	struct slater_kernel{
		double operator()(const double & rho)const{
			double tmp = rho;
			if(tmp<0.0) tmp=1.e-10;
			return prefactor*pow(tmp,-2.0/3.0);
		}
		const double prefactor = 1.0/M_PI;
	};



	madness::Tensor<double> slater_apply(const std::vector< madness::Tensor<double> >& t,const madness::Key<3> & key,const FunctionCommonData<double,3> cdata) const{
	madness::Tensor<double> rho = t.front();
	madness::Tensor<double> pt_rho = t.back();
        const long k = FunctionDefaults<3>::get_k();
	const Tensor<double>& quad_x = cdata.quad_x;

	// if all points of the perturbed density are below the threshold: Set the result to zero
	bool pt_zero = true;
	if(pt_rho.absmax()>safety_thresh_ptrho) pt_zero = false;

	// fast return if pt_rho is zero
	if(pt_zero){
		std::vector<long> kdims(3,k);
		Tensor<double> zero_result(kdims,true);
		return zero_result;
	}

	// if all points of the unperturbed density are below the threshold: Use asymptotical density
	bool rho_ana = true;
	//if(rho.absmax()>safety_thresh_rho) rho_ana = false;

	// analytical kernel based on asymptotical density
	Tensor<double> kernel_val(k,k,k);
	if(rho_ana){
		madness::fcube<double,3>(key, asl, quad_x,kernel_val);
	}else{
		kernel_val = copy(rho);
		slater_kernel slater_tmp;
		kernel_val.unaryop<slater_kernel>(slater_tmp);
	}

	// multiply kernel values with pt_rho values
	Tensor<double> result = kernel_val.emul(pt_rho);
	return result;

	}

	const double prefactor = 1.0/M_PI;
	const double ipot;
	asymptotic_slater asl;
	const double safety_thresh_rho = FunctionDefaults<3>::get_thresh()*0.01;
	const double safety_thresh_ptrho = FunctionDefaults<3>::get_thresh()*0.01;
};

struct slater_kernel_apply {

    const FunctionCommonData<double,3>& cdata;

    const apply_kernel_helper *xc;
    slater_kernel_apply(const apply_kernel_helper &xc_) : cdata(FunctionCommonData<double,3>::get(FunctionDefaults<3>::get_k())), xc(&xc_) {}

    madness::Tensor<double> operator()(const madness::Key<3> & key,
            const std::vector< madness::Tensor<double> >& t) const {
        madness::Tensor<double> r = xc->slater_apply(t,key,cdata);
        return r;
    }
};

struct density_mask_operator{

    density_mask_operator(const double thresh): safety(thresh){}

    madness::Tensor<double> operator()(const madness::Key<3> & key,
            const std::vector< madness::Tensor<double> >& t)const {

    	const madness::Tensor<double> &rho = t.front();
    	// set result to zero
    	const long & k = FunctionDefaults<3>::get_k();
		std::vector<long> kdims(3,k);
    	madness::Tensor<double> result(kdims,true);
    	result =3.0;
    	if(rho.absmax()<safety){
    		result = 0.0;
    	}
    	else{
    		result =1.0;
    	}
    	return result;
    }
    const double safety;
};

struct merging_operator{

    merging_operator(const double c): cutoff(c){}

    madness::Tensor<double> operator()(const madness::Key<3> & key,
            const std::vector< madness::Tensor<double> >& t)const {

    	const madness::Tensor<double> &f = t.front();
    	const madness::Tensor<double> &sf = t.back();
    	if(f.absmax()>cutoff){
    		return f;
    	}
    	else{
    		return sf;
    	}
    }
    const double cutoff;
};


template<std::size_t NDIM>
static double make_radius(const Vector<double,NDIM> &x){
	double r2=0.0;
	for(double xi:x) r2+=xi*xi;
	return sqrt(r2);
}
template<std::size_t NDIM>
static double mask_munging(const double x){
	if(fabs(x)<0.1*FunctionDefaults<NDIM>::get_thresh()) return 0.0;
	else return x;
}
template<std::size_t NDIM>
static double unitfunctor(const Vector<double,NDIM> &x){
	return 1.0;
}

template<std::size_t NDIM>
static double munge(const double x){
	if(fabs(x)<FunctionDefaults<NDIM>::get_thresh()*0.001) return 1.e-20;
	else return x;
}

template<typename T, std::size_t NDIM>
class smooth{
public:
	smooth(World&world):
		world(world){
	}
	smooth(World&world, const double &box_mask_factor, const double box_mask_cutoff):
		world(world){
	}

	struct mask_functor{
	public:
		mask_functor(const double &factor, const double &cutoff): factor_(factor), cutoff_radius_(cutoff), shift_(0.0){}
		mask_functor(const double &factor, const double &cutoff, const Vector<double,NDIM> shift): factor_(factor), cutoff_radius_(cutoff), shift_(shift){std::cout << "shift is " << shift_ << std::endl;}
		double operator()(const Vector<double,NDIM> &x)const{
			const Vector<double,3> shifted_x = x-shift_;
			const double r = make_radius<NDIM>(shifted_x);
			return (1.0-tanh(r));
		}
	private:
		const double factor_;
		const double cutoff_radius_;
		const Vector<double,NDIM> shift_;
	};
	struct inv_mask_functor{
	public:
		inv_mask_functor(const double &factor, const double &cutoff): factor_(factor), cutoff_radius_(cutoff), shift_(0.0){}
		inv_mask_functor(const double &factor, const double &cutoff, const Vector<double,NDIM> shift): factor_(factor), cutoff_radius_(cutoff), shift_(shift){}
		double operator()(const Vector<double,NDIM> &x)const{
			const Vector<double,3> shifted_x = x-shift_;
			const double r = make_radius<NDIM>(shifted_x);
			return 0.5*(1.0 + tanh(factor_*(r-cutoff_radius_)));
		}
	private:
		const double factor_;
		const double cutoff_radius_;
		const Vector<double,NDIM> shift_;
	};



	Function<T,NDIM> smooth_sigma(const Function<T,NDIM> &density){
		MADNESS_ASSERT(NDIM==3);
		std::cout << "Make Sigma\n";
		real_function_3d sigma = make_sigma(density);
		make_plots(density,"density");
		make_plots(sigma,"sigma");

		const double cutoff = get_density_thresh();

		std::cout << "Make Smooth Density\n";
		real_function_3d smooth_rho = cut_oscillations(density,cutoff);

		std::cout << "Make Smooth Sigma\n";
		real_function_3d smooth_sigma = cut_oscillations(sigma,cutoff);
		make_plots(smooth_sigma,"smooth_sigma");
		sigma.print_size("sigma");
		smooth_sigma.print_size("smooth_sigma");
		make_plots(smooth_rho,"smooth_density");


		return sigma;
	}

	Function<T,NDIM> smooth_density_from_orbitals(const std::vector<Function<T,NDIM>> &orbitals)const{
		const double thresh = FunctionDefaults<NDIM>::get_thresh();
		Function<T,NDIM> density = FunctionFactory<T,NDIM>(world);
		for(size_t i=0;i<orbitals.size();i++){
			Function<T,NDIM> orb = orbitals[i];
			orb.refine();
			Function<T,NDIM> boxes = cut_oscillations(orb,1000.0,1);
			make_plots(boxes,"orbital_sceleton");
			Function<T,NDIM> sf = cut_oscillations(orbitals[i],thresh,1);
			make_plots(sf,"smooth_orbital_1_"+stringify(i));
			sf = cut_oscillations(sf,thresh*0.1,1);
			make_plots(sf,"smooth_orbital_2_"+stringify(i));
			sf = cut_oscillations(sf,thresh*0.01,1);
			Function<T,NDIM> sf_squared = sf*sf;
			density += sf_squared;
			make_plots(sf,"smooth_orbital_3_"+stringify(i));
			make_plots(sf_squared,"smooth_squared_orbital_"+stringify(i));
		}
		make_plots(density,"smooth_density");
		return density;
	}

	// Function gets linearized if it falls below certain threshold
	Function<T,NDIM> cut_oscillations(const Function<T,NDIM> &f,const double & cutoff, const size_t order)const{
		// project the function to linear scaling functions
		Function<T,NDIM> lf = project(f,order);
		lf = project(lf,FunctionDefaults<NDIM>::get_k());
		// merge
		const Function<T,NDIM> sf = merge_functions(f,lf,cutoff);
		return sf;
	}

	Function<T,NDIM> make_sigma(const Function<T,NDIM> &density)const{
		MADNESS_ASSERT(NDIM==3);
		std::vector < std::shared_ptr<real_derivative_3d> > gradop;
		gradop = gradient_operator<double, 3>(world);
		real_function_3d sigma = real_factory_3d(world);
		for(auto d:gradop){
			density.autorefine();
			real_function_3d drho = (*d)(density);
			drho.autorefine();
			sigma += drho*drho;
		}
		return sigma;
	}

	Function<T,NDIM> smooth_density(const Function<T,NDIM> &density, const std::string & msg ="function"){
		make_plots(density,"density");
		// if thresh is 1.e-x then density thresh is 1.e-2x
		const double density_thresh = get_density_thresh();
		std::cout << "Density thresh is " << density_thresh << std::endl;
		Function<T,NDIM> ln_density = copy(density);
		ln_density.unaryop(make_log);
		make_plots(ln_density,"ln_density");
		Function<T,NDIM> linearized_ln_density = linearize(ln_density);
		linearized_ln_density = gaussian_smoothing(linearized_ln_density,0.5,"linearized_ln_density");
		Function<T,NDIM> smooth_density = copy(linearized_ln_density);
		smooth_density.refine();
		smooth_density.unaryop(make_exp);
		Function<T,NDIM> munged_density = merge_functions(density,smooth_density,get_density_thresh());
		return munged_density;
	}

	Function<T,NDIM> linearize(const Function<T,NDIM> &ln_rho)const{
		Function<T,NDIM> linear_ln_rho = project(ln_rho,2);
		make_plots(linear_ln_rho,"linear_ln_rho");
		Function<T,NDIM> result = project(linear_ln_rho,FunctionDefaults<NDIM>::get_k());
		make_plots(result,"linear_ln_rho_backprojected");
		return result;
	}

	Function<T,NDIM> munge_density(const Function<T,NDIM> &density)const{
		Function<T,NDIM> munged_density = copy(density);
		munged_density.unaryop(munge<NDIM>);
		return munged_density;
	}

	Function<T,NDIM> make_density_mask(const Function<T,NDIM> &density)const{
		vecfuncT density_wrapper;
		density.reconstruct();
		density_wrapper.push_back(density);
		density_mask_operator op_tmp(get_density_thresh());
	    real_function_3d density_mask=multiop_values<double, density_mask_operator, 3>(op_tmp, density_wrapper);
	    make_plots(density_mask,"density_mask");
	    return density_mask;
	}

	double get_density_thresh()const{
		const double &thresh = FunctionDefaults<NDIM>::get_thresh();
		double density_thresh = thresh;
		for(size_t i=0;i<10;i++){
			double di = (double) i;
			if(thresh>pow(10.0,-di)) break;
			else density_thresh = pow(10.0,-2.0*di);
		}
		return density_thresh;
	}

	Function<T,NDIM> make_mask(const Function<T,NDIM> &samplef, const double &thresh)const{
		vecfuncT sample_wrapper;
		samplef.reconstruct();
		sample_wrapper.push_back(samplef);
		density_mask_operator op_tmp(thresh);
	    real_function_3d mask=multiop_values<double, density_mask_operator, 3>(op_tmp, sample_wrapper);
	    // The mask is constant on each box so make shure this will be the case (avoid flutuations)
	    real_function_3d const_mask = project(mask,1);
	    mask = project(const_mask,FunctionDefaults<NDIM>::get_k());
	    make_plots(mask,"mask");
	    return mask;
	}

	Function<T,NDIM> make_inv_mask(const Function<T,NDIM> &mask)const{
		Function<T,NDIM> unit = FunctionFactory<T,NDIM>(world).f(unitfunctor);
		real_function_3d inv = unit-mask;
	    real_function_3d const_inv = project(inv,1);
	    inv = project(const_inv,FunctionDefaults<NDIM>::get_k());
		return inv;
	}

	Function<T,NDIM> make_explog(const Function<T,NDIM> &f, const std::string & msg="function")const{
		// make logf
		Function<T,NDIM> logf = copy(f);
		logf.unaryop(make_log);

		// smooth logf
		Function<T,NDIM> smoothed_logf = gaussian_smoothing(logf,0.5);

		// reconstruct f via f=exp^(log(f));
		Function<T,NDIM> explogf = copy(smoothed_logf);
		explogf.unaryop(make_exp);

		// apply the box_mask to eliminate weird boundary effects
		//explogf = box_mask_*explogf;

		// make some plots
		make_plots(f,msg);
		make_plots(smoothed_logf,"smoothed_log"+msg);
		make_plots(logf,"log"+msg);
		make_plots(explogf,"explog_"+msg);
		return explogf;
	}

	Function<T,NDIM> gaussian_smoothing(const Function<T,NDIM> &f, const double &eps=0.04,const std::string &name="nooutput")const{
		if(name!="noputput") output("\nSmoothing " + name + " with eps " + stringify(eps));
		double exponent = 1.0/(2.0*eps);
		Tensor<double> coeffs(1), exponents(1);
	    exponents(0L) =  exponent;
	    coeffs(0L) = gauss_norm(exponent);
	    SeparatedConvolution<T,NDIM> op(world, coeffs, exponents);
	    Function<T,NDIM> smoothed_f = apply(op,f);
	    make_plots(smoothed_f,"smoothed_"+name);
	    return smoothed_f;

	}

	Function<T,NDIM> merge_functions(const Function<T,NDIM> &f, const Function<T,NDIM> &sf,const Function<T,NDIM> &mask,const std::string &name="f")const{
		Function<T,NDIM> inv_mask = make_inv_mask(mask);
		Function<T,NDIM> part1 = mask*f;
		Function<T,NDIM> part2 = inv_mask*sf;
		Function<T,NDIM> result = part1 + part2;
		make_plots(mask,"mask_"+name);
		make_plots(inv_mask,"inv_mask_"+name);
		make_plots(part1,"masked_"+name);
		make_plots(part2,"masked_s"+name);
		return result;
	}

	Function<T,NDIM> merge_functions(const Function<T,NDIM>&f, const Function<T,NDIM>&sf, const double &thresh)const{
		MADNESS_ASSERT(NDIM==3);
		vecfuncT wrapper;
		wrapper.push_back(f);
		wrapper.push_back(sf);
		refine_to_common_level(world,wrapper);
		merging_operator op_tmp(thresh);
		real_function_3d merged_f=multiop_values<double, merging_operator, 3>(op_tmp,wrapper);
		return merged_f;
	}

	void make_plots(const Function<T,NDIM> &f,const std::string &name="function")const{
		double width = FunctionDefaults<3>::get_cell_min_width()/2.0 - 1.e-3;
		plot_plane(world,f,name);
		coord_3d start(0.0); start[0]=-width;
		coord_3d end(0.0); end[0]=width;
		plot_line(("line_"+name).c_str(),1000,start,end,f);
	}


	void test_1d()const{
		FunctionDefaults<1>::set_thresh(1.e-5);
		FunctionDefaults<1>::set_k(8);
		FunctionDefaults<1>::set_cubic_cell(-25.0,25.0);
		real_function_1d f = real_factory_1d(world).f(test_1d_functor);
		Tensor<double> coeffs(1), exponents(1);
	    exponents(0L) =  1.0/(2.0*0.04);
	    coeffs(0L) = pow((1.0/(2.0*0.04))/M_PI,0.5);
	    SeparatedConvolution<double,1> op(world, coeffs, exponents);
	    real_function_1d sf = apply(op,f);
	    double diff = (f-sf).norm2();
	    output("||f - smoothed_f||_1D="+stringify(diff));
	    coord_1d start; start[0]=-10.0;
	    coord_1d end; end[0]=10.0;
	    plot_line("smoothed_1d",1000,start,end,sf);
	    plot_line("noise_1d",1000,start,end,f);
	}

	double gauss_norm(const double &exponent)const{
		return pow(exponent/M_PI,0.5*3.0);
	}

	void output(const std::string &msg)const{
		if(world.rank()==0) std::cout << msg << std::endl;
	}

private:
	World &world;
	const Function<T,NDIM> box_mask_;
	Function<T,NDIM> mask_;
	Function<T,NDIM> inv_mask_;
};



#endif /* SRC_EXAMPLES_SMOOTH_H_ */
