#include <AcousticSolver.hpp>
#include <rang.hpp>

#include <memory>

#include <BlockPerfectGas.hpp>

typedef const double my_double;

struct AcousticSolver::ReduceMin
{
private:
  const Kokkos::View<my_double*> x_;

public:
  typedef Kokkos::View<my_double*>::non_const_value_type value_type;

  ReduceMin(const Kokkos::View<my_double*>& x) : x_ (x) {}

  typedef Kokkos::View<my_double*>::size_type size_type;
    
  KOKKOS_INLINE_FUNCTION void
  operator() (const size_type i, value_type& update) const
  {
    if (x_(i) < update) {
      update = x_(i);
    }
  }

  KOKKOS_INLINE_FUNCTION void
  join (volatile value_type& dst,
	const volatile value_type& src) const
  {
    if (src < dst) {
      dst = src;
    }
  }

  KOKKOS_INLINE_FUNCTION void
  init (value_type& dst) const
  { // The identity under max is -Inf.
    dst= Kokkos::reduction_identity<value_type>::min();
  }
};

KOKKOS_INLINE_FUNCTION
const Kokkos::View<const double*>
AcousticSolver::computeRhoCj(const Kokkos::View<const double*>& rhoj,
			     const Kokkos::View<const double*>& cj)
{
  Kokkos::View<double*> rhocj("rho_c", m_nj);
  Kokkos::parallel_for(m_nj, KOKKOS_LAMBDA(const int& j) {
      rhocj[j] = rhoj[j]*cj[j];
    });
  return rhocj;
}

KOKKOS_INLINE_FUNCTION
const Kokkos::View<const double*[2]>
AcousticSolver::computeAjr(const Kokkos::View<const double*>& rhocj,
			   const Kokkos::View<const double*[2]>& Cjr)
{
  Kokkos::View<double*[2]> Ajr("Ajr", m_nj);
  Kokkos::parallel_for(m_nj, KOKKOS_LAMBDA(const int& j) {
      for (int r=0; r<2; ++r) {
	Ajr(j,r) = rhocj(j)*Cjr(j,r)*Cjr(j,r);
      }
    });

  return Ajr;
}

KOKKOS_INLINE_FUNCTION
const Kokkos::View<const double*>
AcousticSolver::computeAr(const Kokkos::View<const double*[2]>& Ajr,
			  const Kokkos::View<const int*[2]>& node_cells,
			  const Kokkos::View<const int*[2]>& node_cell_local_node,
			  const Kokkos::View<const int*>& node_nb_cells)
{
  Kokkos::View<double*> Ar("Ar", m_nr);
  Kokkos::parallel_for(m_nr, KOKKOS_LAMBDA(const int& r) {
      double sum = 0;
      for (int j=0; j<node_nb_cells(r); ++j) {
	const int J = node_cells(r,j);
	const int R = node_cell_local_node(r,j);
  	sum += Ajr(J,R);
      }
      Ar(r) = sum;
    });

  return Ar;
}

KOKKOS_INLINE_FUNCTION
const Kokkos::View<const double*>
AcousticSolver::computeBr(const Kokkos::View<const double*[2]>& Ajr,
			  const Kokkos::View<const double*[2]>& Cjr,
			  const Kokkos::View<const double*>& uj,
			  const Kokkos::View<const double*>& pj,
			  const Kokkos::View<const int*[2]>& node_cells,
			  const Kokkos::View<const int*[2]>& node_cell_local_node,
			  const Kokkos::View<const int*>& node_nb_cells)
{
  Kokkos::View<double*> br("br", m_nr);
  Kokkos::parallel_for(m_nr, KOKKOS_LAMBDA(const int& r) {
      double sum = 0;
      for (int j=0; j<node_nb_cells(r); ++j) {
	const int J = node_cells(r,j);
	const int R = node_cell_local_node(r,j);
  	sum += Ajr(J,R)*uj(J) + Cjr(J,R)*pj[J];
      }
      br(r) = sum;
    });

  return br;
}

KOKKOS_INLINE_FUNCTION
Kokkos::View<double*>
AcousticSolver::computeUr(const Kokkos::View<const double*>& Ar,
			  const Kokkos::View<const double*>& br)
{
  const Kokkos::View<const double*> invAr = inverse(Ar);
  Kokkos::View<double*> ur("ur", m_nr);
  Kokkos::parallel_for(m_nr, KOKKOS_LAMBDA(const int& r) {
      ur[r]=br(r)*invAr(r);
    });
  ur[0]=0;
  ur[m_nr-1]=0;

  return ur;
}

KOKKOS_INLINE_FUNCTION
Kokkos::View<double*[2]>
AcousticSolver::computeFjr(const Kokkos::View<const double*[2]>& Ajr,
			   const Kokkos::View<const double*>& ur,
			   const Kokkos::View<const double*[2]>& Cjr,
			   const Kokkos::View<const double*>& uj,
			   const Kokkos::View<const double*>& pj,
			   const Kokkos::View<const int*[2]>& cell_nodes)
{
  Kokkos::View<double*[2]> Fjr("Fjr", m_nj);

  Kokkos::parallel_for(m_nj, KOKKOS_LAMBDA(const int& j) {
      for (int r=0; r<2; ++r) {
	Fjr(j,r) = Ajr(j,r)*(uj(j)-ur(cell_nodes(j,r)))+Cjr(j,r)*pj(j);
      }
    });

  return Fjr;
}

KOKKOS_INLINE_FUNCTION
double AcousticSolver::
acoustic_dt(const Kokkos::View<const double*>& Vj,
	    const Kokkos::View<const double*>& cj) const
{
  double dt = std::numeric_limits<double>::max();

  Kokkos::View<double*> Vj_cj("Vj_cj", m_nj);

  Kokkos::parallel_for(m_nj, KOKKOS_LAMBDA(const int& j){
      Vj_cj[j] = Vj[j]/cj[j];
    });

  Kokkos::parallel_reduce(m_nj, ReduceMin(Vj_cj), dt);

  return dt;
}

KOKKOS_INLINE_FUNCTION
const Kokkos::View<const double*>
AcousticSolver::inverse(const Kokkos::View<const double*>& x) const
{
  Kokkos::View<double*> inv_x("inv_x", x.size());
  Kokkos::parallel_for(x.size(), KOKKOS_LAMBDA(const int& r) {
      inv_x(r) = 1./x(r);
    });

  return inv_x;
}

KOKKOS_INLINE_FUNCTION
void AcousticSolver::computeExplicitFluxes(const Kokkos::View<const double*>& xr,
					   const Kokkos::View<const double*>& xj,
					   const Kokkos::View<const double*>& rhoj,
					   const Kokkos::View<const double*>& uj,
					   const Kokkos::View<const double*>& pj,
					   const Kokkos::View<const double*>& cj,
					   const Kokkos::View<const double*>& Vj,
					   const Kokkos::View<const double*[2]>& Cjr,
					   const Kokkos::View<const int*[2]>& cell_nodes,
					   const Kokkos::View<const int*[2]>& node_cells,
					   const Kokkos::View<const int*>& node_nb_cells,
					   const Kokkos::View<const int*[2]>& node_cell_local_node,
					   Kokkos::View<double*>& ur,
					   Kokkos::View<double*[2]>& Fjr)
{
  const Kokkos::View<const double*> rhocj  = computeRhoCj(rhoj, cj);
  const Kokkos::View<const double*[2]> Ajr = computeAjr(rhocj, Cjr);

  const Kokkos::View<const double*> Ar = computeAr(Ajr, node_cells, node_cell_local_node, node_nb_cells);
  const Kokkos::View<const double*> br = computeBr(Ajr, Cjr, uj, pj,
						   node_cells, node_cell_local_node, node_nb_cells);

  ur  = computeUr(Ar, br);
  Fjr = computeFjr(Ajr, ur, Cjr, uj, pj, cell_nodes);
}

AcousticSolver::AcousticSolver(const long int& nj)
  : m_nj(nj),
    m_nr(nj+1)
{
  Kokkos::View<double*> xj("xj",m_nj);
  Kokkos::View<double*> rhoj("rhoj",m_nj);

  Kokkos::View<double*> uj("uj",m_nj);

  Kokkos::View<double*> Ej("Ej",m_nj);
  Kokkos::View<double*> ej("ej",m_nj);
  Kokkos::View<double*> pj("pj",m_nj);
  Kokkos::View<double*> Vj("Vj",m_nj);
  Kokkos::View<double*> gammaj("gammaj",m_nj);
  Kokkos::View<double*> cj("cj",m_nj);
  Kokkos::View<double*> mj("mj",m_nj);

  Kokkos::View<double*>  xr("xr", m_nr);

  const double delta_x = 1./m_nj;

  Kokkos::parallel_for(m_nr, KOKKOS_LAMBDA(const int& r){
      xr[r] = r*delta_x;
    });

  Kokkos::View<int*[2]> cell_nodes("cell_nodes",m_nj,2);
  Kokkos::View<int*[2]> node_cells("node_cells",m_nr,2);
  Kokkos::View<int*[2]> node_cell_local_node("node_cells",m_nr,2);
  Kokkos::View<int*> node_nb_cells("node_cells",m_nr);

  Kokkos::parallel_for(m_nr, KOKKOS_LAMBDA(const int& r){
      node_nb_cells(r) = 2;
    });
  node_nb_cells(0) = 1;
  node_nb_cells(m_nr-1) = 1;

  node_cells(0,0) = 0;
  Kokkos::parallel_for(m_nr-2, KOKKOS_LAMBDA(const int& r){
      node_cells(r+1,0) = r;
      node_cells(r+1,1) = r+1;
    });
  node_cells(m_nr-1,0) =m_nj-1;

  Kokkos::parallel_for(nj, KOKKOS_LAMBDA(const int& j){
      cell_nodes(j,0) = j;
      cell_nodes(j,1) = j+1;
    });

  Kokkos::parallel_for(m_nr, KOKKOS_LAMBDA(const int& r){
      for (int J=0; J<node_nb_cells(r); ++J) {
	int j = node_cells(r,J);
	for (int R=0; R<2; ++R) {
	  if (cell_nodes(j,R) == r) {
	    node_cell_local_node(r,J) = R;
	  }
	}
      }
    });

  Kokkos::parallel_for(nj, KOKKOS_LAMBDA(const int& j){
      xj[j] = 0.5*(xr[cell_nodes(j,0)]+xr[cell_nodes(j,1)]);
    });

  Kokkos::parallel_for(nj, KOKKOS_LAMBDA(const int& j){
      Vj[j] = xr[cell_nodes(j,1)]-xr[cell_nodes(j,0)];
    });

  Kokkos::parallel_for(nj, KOKKOS_LAMBDA(const int& j){
    if (xj[j]<0.5) {
      rhoj[j]=1;
    } else {
      rhoj[j]=0.125;
    }
  });

  Kokkos::parallel_for(nj, KOKKOS_LAMBDA(const int& j){
    if (xj[j]<0.5) {
      pj[j]=1;
    } else {
      pj[j]=0.1;
    }
  });

  Kokkos::parallel_for(nj, KOKKOS_LAMBDA(const int& j){
      gammaj[j] = 1.4;
    });

  BlockPerfectGas block_eos(rhoj, ej, pj, gammaj, cj);

  block_eos.updateEandCFromRhoP();

  Kokkos::parallel_for(nj, KOKKOS_LAMBDA(const int& j){
      Ej[j] = ej[j]+0.5*uj[j]*uj[j];
  });

  Kokkos::parallel_for(nj, KOKKOS_LAMBDA(const int& j){
    mj[j] = rhoj[j] * Vj[j];
  });

  const Kokkos::View<const double*> inv_mj=inverse(mj);

  const double tmax=0.2;
  double t=0;

  int itermax=std::numeric_limits<int>::max();
  int iteration=0;

  Kokkos::View<double*[2]> Cjr("Cjr",m_nj);
  Kokkos::parallel_for(nj, KOKKOS_LAMBDA(const int& j) {
      Cjr(j,0)=-1;
      Cjr(j,1)= 1;
    });

  while((t<tmax) and (iteration<itermax)) {
    double dt = 0.4*acoustic_dt(Vj, cj);
    if (t+dt<tmax) {
      t+=dt;
    } else {
      dt=tmax-t;
      t=tmax;
    }
    
    Kokkos::View<double*> ur;
    Kokkos::View<double*[2]> Fjr;
  
    computeExplicitFluxes(xr, xj,
			  rhoj, uj, pj, cj, Vj, Cjr,
			  cell_nodes, node_cells, node_nb_cells, node_cell_local_node,
			  ur, Fjr);

    Kokkos::parallel_for(nj, KOKKOS_LAMBDA(const int& j){
	double momentum_fluxes = 0;
	double energy_fluxes = 0;
	for (int R=0; R<2; ++R) {
	  const int r=cell_nodes(j,R);
	  momentum_fluxes += Fjr(j,R);
	  energy_fluxes += Fjr(j,R)*ur[r];
	}
	uj[j] -= dt*inv_mj[j]*(momentum_fluxes);
	Ej[j] -= dt*inv_mj[j]*(energy_fluxes);
      });

    Kokkos::parallel_for(nj, KOKKOS_LAMBDA(const int& j) {
	ej[j] = Ej[j] - 0.5 * uj[j]*uj[j];
      });

    Kokkos::parallel_for(m_nr, KOKKOS_LAMBDA(const int& r){
	xr[r] += dt*ur[r];
      });

    Kokkos::parallel_for(nj, KOKKOS_LAMBDA(const int& j){
      xj[j] = 0.5*(xr[j]+xr[j+1]);
      Vj[j] = xr[j+1]-xr[j];
      });

    Kokkos::parallel_for(nj, KOKKOS_LAMBDA(const int& j){
	rhoj[j] = mj[j]/Vj[j];
      });

    block_eos.updatePandCFromRhoE();    
    
    ++iteration;
  }

  {
    std::ofstream fout("rho");
    for (int j=0; j<nj; ++j) {
      fout << xj[j] << ' ' << rhoj[j] << '\n';
    }
  }

  std::cout << "* " << rang::style::underline << "Final time" << rang::style::reset
	    << ":  " << rang::fgB::green << t << rang::fg::reset << " (" << iteration << " iterations)\n";

}