#include <unicorn/ErrorEstimate.h>
#include <unicorn/MeshQuality.h>

using namespace dolfin::unicorn;
using namespace dolfin;

//-----------------------------------------------------------------------------
ErrorEstimate::ErrorEstimate(Mesh& mesh, Form* Lres, Form* Lgradphi) :
  mesh(mesh), Lres(Lres), Lgradphi(Lgradphi),
  e_indx(mesh.numCells()), assembler(mesh)
{
}
//-----------------------------------------------------------------------------
void ErrorEstimate::ComputeError(real& error)
{
  real sum = 0.0;
  real* eindarr = e_indx.vec().array();

  for (CellIterator cell(mesh); !cell.end(); ++cell)
  {
    int id = (*cell).index();
    real ei = eindarr[id];
    sum += ei;
  }
  cout << "ERROR = " << sum << endl << endl << endl;

  error = sum;
}
//-----------------------------------------------------------------------------
void ErrorEstimate::ComputeErrorIndicator(real t, real k, real T)
{
  int N = mesh.numCells();

  int nsd = mesh.topology().dim();

  // Assemble strong residual at time t
  assembler.assemble(resx, *Lres);

  // Assemble gradient of phi at time t
  assembler.assemble(gradphix, *Lgradphi);

  real* gradphiarr = gradphix.vec().array();
  real* resarr = resx.vec().array();
  real* eindarr = e_indx.vec().array();

  // Compute error indicator using end-point quadrature in time
  for (CellIterator c(mesh); !c.end(); ++c)
    {
      Cell& cell = *c;

      real Si = 0.0;
      int id = cell.index();
      real h = MeshQuality::myDiameter(cell);

      for(int ind = 0; ind < nsd*nsd; ind++)
	Si += sqr(gradphiarr[id + ind * N]);
      
      Si = sqrt(Si * cell.volume());

      real norm_R = 0.0;
      for (int ind = 0; ind < nsd; ind++)
	norm_R += sqr(resarr[id + ind * N]);
      norm_R = sqrt(norm_R * cell.volume());

      real ei = 1.0 / T * k * h * Si * norm_R;

      eindarr[id] += ei;
    }

}
//-----------------------------------------------------------------------------
void ErrorEstimate::ComputeLargestIndicators(std::vector<int>& cells,
					 real percentage)
{
  int N = e_indx.size();
  int M = (int)((real)N * percentage * 0.01);

  cout << "Computing largest indicators" << endl;
  cout << "percentage: " << percentage << endl;
  cout << "N: " << N << endl;
  cout << "M: " << M << endl;

  cells.resize(M);

  std::vector<std::pair<int, real> > indicators(N);

  real* eindarr = e_indx.vec().array();

  for (CellIterator cell(mesh); !cell.end(); ++cell)
    {
      int id = (*cell).index();
      std::pair<int, real> p;
      p.first = id;
      p.second = eindarr[id];

      indicators[id] = p;
    }

  less_pair comp;

  std::sort(indicators.begin(), indicators.end(), comp);

  for(int i = 0; i < M; i++)
    {
      std::pair<int, real> p = indicators[N - 1 - i];
      cells[i] = p.first;

      //cout << "cell id: " << p.first << endl;
      //cout << "error: " << p.second << endl;
    }
}
//-----------------------------------------------------------------------------
void ErrorEstimate::AdaptiveRefinement(std::vector<int>& cells, int iter)
{
  cout << "Adaptive refinement" << endl;
  cout << "cells before: " << mesh.numCells() << endl;
  cout << "vertices before: " << mesh.numVertices() << endl;
  Mesh refined_mesh(mesh);
  
  MeshFunction<bool> cell_refinement_marker(refined_mesh);
  cell_refinement_marker.init(refined_mesh.topology().dim());

  int M = cells.size();

  for (CellIterator c(refined_mesh); !c.end(); ++c)
    {
      cell_refinement_marker.set(c->index(), false);
    }

  for(int i = 0; i < M; i++)
    {
      cell_refinement_marker.set(cells[i], true);
    }

  refined_mesh.refine(cell_refinement_marker);

  std::stringstream is;
  is << iter << "_";

  File meshfile("refined" + is.str() + ".xml");
  meshfile << refined_mesh;
  File meshfile2("refined.xml");
  meshfile2 << refined_mesh;

  //  Function zero(refined_mesh, element);
  //  File vtkfile("refined.pvd");
  //  vtkfile << zero;

//   std::stringstream filename, ss;
//   ss << iter;
//   filename << "refined" << ss.str() << ".pvd" << std::ends;

//   Function zero(refined_mesh, element);
//   File vtkfile("refined" + is.str() + ".pvd");
//   vtkfile << zero;

  cout << "cells after: " << refined_mesh.numCells() << endl;
  cout << "vertices after: " << refined_mesh.numVertices() << endl;
}
//-----------------------------------------------------------------------------
