#include "tight_binding.hpp"

#include <triqs/arrays/expressions/arithmetic.hpp>
#include <triqs/arrays/expressions/min_max.hpp>
#include <triqs/arrays/linalg/eigenelements.hpp>
#include <triqs/python_tools/converters.hpp> 
#include "grid_generator.hpp"
#include "functors.hpp"

using namespace std;
namespace triqs { namespace lattice_tools { 
 using namespace tqa;

 tight_binding::tight_binding(bravais_lattice const & bl__, map_type const & t_r) : bl_(bl__),tr(t_r) {check();}

 tight_binding::tight_binding(bravais_lattice const & bl__,boost::python::object dct) : 
  bl_(bl__), tr(triqs::python_tools::Py_to_C::convert<map_type>::invoke(dct)) {check();}

 void tight_binding::check() { 
  const size_t no = bl_.n_orbitals();
  for (map_type::const_iterator it = tr.begin(); it !=tr.end(); ++it) { 
   if (it->second.len(0) != no) throw triqs::runtime_error()<<"tight_binding construction : the first dim matrix is of size "<<  it->second.len(0) <<" instead of "<< no;
   if (it->second.len(1) != no) throw triqs::runtime_error()<<"tight_binding construction : the second dim matrix is of size "<<  it->second.len(1) <<" instead of "<< no;
  }
 }

 //------------------------------------------------------

 array_view <dcomplex,3> hopping_stack (tight_binding const & TB, array<double,2> const & k_stack) {
  result_of::Fourier<tight_binding>::type TK = Fourier(TB);  
  array<dcomplex,3> res(TB.n_bands(), TB.n_bands(), k_stack.len(1));
  for(size_t i = 0; i<k_stack.len(1); ++i) res(range(), range(), i) = TK(k_stack(range(),i));
  return res;
 }

 //------------------------------------------------------

 array_view<double,2> energies_on_bz_path(tight_binding const & TB, K_view_type K1, K_view_type K2, size_t n_pts) {

  result_of::Fourier<tight_binding>::type TK = Fourier(TB);
  const size_t norb=TB.lattice().n_orbitals();
  const size_t ndim=TB.lattice().dim();
  array<double,2> eval(norb,n_pts);
  K_type dk = (K2 - K1)/double(n_pts), k = K1;
  for (size_t i =0; i<n_pts; ++i, k += dk) {
   eval(range(),i) = linalg::eigenvalues( TK( k (range(0,ndim))), false);
  }
  return eval;
 }

 //------------------------------------------------------
 array_view<double,2> energies_on_bz_grid(tight_binding const & TB, size_t n_pts) {

  result_of::Fourier<tight_binding>::type TK = Fourier(TB);
  const size_t norb=TB.lattice().n_orbitals();
  const size_t ndim=TB.lattice().dim();
  grid_generator grid(ndim,n_pts);
  array<double,2> eval(norb,grid.size());
  for (; grid ; ++grid)  {
   eval(range(),grid.index()) = linalg::eigenvalues( TK( (*grid) (range(0,ndim))), false);
  }
  return eval;
 }

 //------------------------------------------------------

 struct bz_grid { 

	 const brillouin_zone bz;
	 const size_t nkpts, dim;
	 const size_t N_Z,N_Y,N_X;

	 size_t size () const { return (N_X * N_Y * N_Z);}

	 bz_grid( brillouin_zone const & bz_, size_t nkpts_) : 
		 bz(bz_), 
		 nkpts(nkpts_), dim (bz.lattice().dim()),
		 N_X (nkpts), N_Y (dim>1 ? nkpts : 1), N_Z  (dim>2 ? nkpts : 1) {}

	 template <class F> 
		 void foreach(F  f) { 
			 K_type k(3);
			 for (int nz = 0; nz < N_Z; nz++)
				 for (int ny = 0; ny < N_Y; ny++)
					 for (int nx = 0; nx < N_X; nx++) {
						 k(0) = nx/double(N_X);
						 k(1) = ny/double(N_Y);
						 k(2) = nz/double(N_Z);
						 f(k);
					 }
		 }
 };


#ifndef TRIQS_HAS_CPP11_LAMBDA
 void lambda_loop( result_of::Fourier<tight_binding>::type const & TK, 
		 array<dcomplex,3> & evec, array<double,2> & eval,
		 size_t & index, size_t ndim, K_type const & k) { 
	 array_view <double,1> eval_sl = eval(range(),index);
	 array_view <dcomplex,2> evec_sl =  evec(range(),range(),index);
	 boost::tie (eval_sl,evec_sl) = linalg::eigenelements( TK( k(range(0,ndim)))) ;//,  true);
	 index++;
 }  

}}
#include <boost/lambda/bind.hpp>
#include <boost/ref.hpp>
namespace triqs { namespace lattice_tools { 

	namespace BLL=boost::lambda;
	using boost::ref;
	using boost::cref;
#endif


	std::pair<array<double,1>, array<double,2> > dos(tight_binding const & TB, size_t nkpts, size_t neps) {

		// The Fourier transform of TK
		// auto TK = Fourier(TB); // C++0x .... 
		result_of::Fourier<tight_binding>::type TK = Fourier(TB);

		// loop on the BZ
		const size_t ndim=TB.lattice().dim();
		const size_t norb=TB.lattice().n_orbitals();

		bz_grid grid(brillouin_zone(TB.lattice()),nkpts);
		//grid_generator grid(ndim,nkpts);
		array<double,1> tempeval(norb);
		array<dcomplex,3> evec(norb,norb,grid.size());
		array<double,2> eval(norb,grid.size());
		size_t index=0;
#ifdef TRIQS_HAS_CPP11_LAMBDA
		grid.foreach([&] { 
				array_view <double,1> eval_sl = eval(range(),index);
				array_view <dcomplex,2> evec_sl =  evec(range(),range(),index);
				boost::tie (eval_sl,evec_sl) = linalg::eigenelements( TK( k(range(0,ndim))),  true);
				index++;
				}); 

#else

		 int N_X (nkpts), N_Y (ndim>1 ? nkpts : 1), N_Z  (ndim>2 ? nkpts : 1);
		 K_type k(3);
		 for (int nz = 0; nz < N_Z; nz++)
			 for (int ny = 0; ny < N_Y; ny++)
				 for (int nx = 0; nx < N_X; nx++) {
					 k(0) = nx/double(N_X);
					 k(1) = ny/double(N_Y);
					 k(2) = nz/double(N_Z);
					 array_view <double,1> eval_sl = eval(range(),index);
					 array_view <dcomplex,2> evec_sl =  evec(range(),range(),index);
					 boost::tie (eval_sl,evec_sl) = linalg::eigenelements( TK( k(range(0,ndim)))) ;//,  true);
					 index++;
				 }


//		grid.foreach(BLL::bind(lambda_loop,cref(TK),ref(evec),ref(eval),ref(index),ndim,BLL::_1));
#endif

/*
   for (; grid ; ++grid)  {
//cerr<<" index = "<<grid.index()<<endl;
array_view <double,1> eval_sl = eval(range(),grid.index());
array_view <dcomplex,2> evec_sl =  evec(range(),range(),grid.index());
boost::tie (eval_sl,evec_sl) = linalg::eigenelements( TK( (*grid) (range(0,ndim))),  true);
//cerr<< " point "<< *grid <<  " value "<< eval_sl<< endl; //" "<< (*grid) (range(0,ndim)) << endl;
}
 */
// define the epsilon mesh, etc.
array<double,1> epsilon(neps); 
double epsmax = tqa::max_element(eval);
double epsmin = tqa::min_element(eval);
double deps=(epsmax-epsmin)/neps;
//for (size_t i =0; i< neps; ++i) epsilon(i)= epsmin+i/(neps-1.0)*(epsmax-epsmin);
for (size_t i =0; i< neps; ++i) epsilon(i)=epsmin+(i+0.5)*deps;

// bin the eigenvalues according to their energy
// NOTE: a is defined as an integer. it is the index for the DOS.
//REPORT <<"Starting Binning ...."<<endl;
array<double,2> rho (neps,norb);rho()=0;
for(size_t l=0;l<norb;l++){
	for (size_t j=0;j<grid.size();j++){
		for (size_t k=0;k<norb;k++){
			int a=int((eval(k,j)-epsmin)/deps);
			if(a==int(neps)) a=a-1;
			rho(l,a) += real(conj(evec(l,k,j))*evec(l,k,j));
			//dos(a) +=  real(conj(evec(l,k,j))*evec(l,k,j));
		}
	}
}
//rho = rho / double(grid.size()*deps);
rho /= grid.size()*deps;
return std::make_pair( epsilon, rho);
}

//----------------------------------------------------------------------------------

std::pair<array<double,1>, array<double,1> > dos_patch(tight_binding const & TB, const array<double,2> & triangles, size_t neps, size_t ndiv) {
	// WARNING: This version only works for a single band Hamiltonian in 2 dimensions!!!!
	// triangles is an array of points defining the triangles of the patch
	// neps in the number of bins in energy
	// ndiv in the number of divisions used to divide the triangles

	//const size_t ndim=TB.lattice().dim();
	//const size_t norb=TB.lattice().n_orbitals();
	int ntri = triangles.len(0)/3;
	array<double,1> dos(neps);

	// Check consistency
	const size_t ndim=TB.lattice().dim();
	//const size_t norb=TB.lattice().n_orbitals();
	if (ndim !=2)  throw triqs::runtime_error()<<"dos_patch : dimension 2 only !";
	if (triangles.len(1) != ndim) throw triqs::runtime_error()<<"dos_patch : the second dimension of the 'triangle' array in not "<<ndim;

	// Every triangle has ndiv*ndiv k points
	size_t nk = ntri*ndiv*ndiv;
	size_t k_index = 0;
	double epsmax=-100000,epsmin=100000;
	array<dcomplex,2> thop(1,1);
	array<double,1> energ(nk), weight(nk);

	// a, b, c are the corners of the triangle
	// g the center of gravity taken from a
	array<double,1> a(ndim), b(ndim), c(ndim), g(ndim), rv(ndim);
	int pt = 0;
	double s, t;

	// The Fourier transform of TK
	// auto TK = Fourier(TB); // C++0x .... 
	result_of::Fourier<tight_binding>::type TK = Fourier(TB);

	// loop over the triangles
	for (int tri = 0; tri < ntri; tri++) {
		a = triangles(pt,range());
		pt++;
		b = triangles(pt,range());
		pt++;
		c = triangles(pt,range());
		pt++;
		g = ((a+b+c)/3.0-a)/double(ndiv);

		// the area around a k point might be different from one triangle to the other
		// so I use it to weight the sum in the dos
		double area = abs(0.5*((b(0)-a(0))*(c(1)-a(1))-(b(1)-a(1))*(c(0)-a(0)))/(ndiv*ndiv));

		for (size_t i = 0; i<ndiv; i++) {
			s = i/double(ndiv);
			for (size_t j = 0; j<ndiv-i; j++) {
				t = j/double(ndiv);
				for (size_t k = 0; k<2; k++) {

					rv = a+s*(b-a)+t*(c-a)+(k+1.0)*g;

					if(k==0 || j < ndiv-i-1) {

						energ(k_index)  = real(TK(rv)(0,0)); 
						//compute(rv);
						//energ(k_index) = real(tk_for_eval(1,1)); //tk_for_eval is Fortran array
						weight(k_index) = area;
						if (energ(k_index)>epsmax) epsmax=energ(k_index);
						if (energ(k_index)<epsmin) epsmin=energ(k_index);
						k_index++;
					}
				}
			}
		}
	}
	// check consistency
	assert(k_index == nk);

	// define the epsilon mesh, etc.
	array<double,1> epsilon(neps); 
	double deps=(epsmax-epsmin)/neps;
	for (size_t i =0; i< neps; ++i) epsilon(i)= epsmin+i/(neps-1.0)*(epsmax-epsmin);

	// bin the eigenvalues according to their energy
	int ind;
	double totalweight(0.0);
	dos() = 0.0;
	for (size_t j = 0; j < nk; j++) {
		ind=int((energ(j)-epsmin)/deps);
		if (ind == int(neps)) ind--;
		dos(ind) += weight(j);
		totalweight += weight(j);
	}
	dos /= deps;// Normalize the DOS
	return std::make_pair(epsilon, dos);
}





}}