#include <config.h>
#include <iomanip>

// Includes for the ADOL-C automatic differentiation library
// Need to come before (almost) all others.
#include <adolc/adouble.h>
#include <dune/fufem/utilities/adolcnamespaceinjections.hh>

#include <dune/common/parametertree.hh>
#include <dune/common/parallel/mpihelper.hh>

#include <dune/elasticity/assemblers/localadolcstiffness.hh>
#include <dune/elasticity/assemblers/feassembler.hh>
#include <dune/elasticity/materials/localintegralenergy.hh>
#include <dune/elasticity/materials/mooneyrivlindensity.hh>

#include <dune/functions/functionspacebases/interpolate.hh>
#include <dune/functions/functionspacebases/lagrangebasis.hh>
#include <dune/functions/functionspacebases/powerbasis.hh>

#include <dune/grid/utility/structuredgridfactory.hh>
#include <dune/grid/uggrid.hh>

// grid dimension
const int dim = 3;

//order of integration
const int order = 1;

using namespace Dune;

//differentiation method: ADOL-C
typedef adouble ValueType;

using namespace Dune;

typedef BlockVector<FieldVector<double,dim> > VectorType;
typedef BCRSMatrix<FieldMatrix<double, dim, dim> > MatrixType;

template <class Basis>
int assembleAndCompare (const Basis basis, const ParameterTree parameters, const VectorType x,
                        double expectedEnergy, double expectedGradientTwoNorm, double expectedGradientInfinityNorm, double expectedMatrixFrobeniusNorm) {
  using LocalView = typename Basis::LocalView;

  std::string mooneyrivlin_energy = parameters.get<std::string>("mooneyrivlin_energy");

  auto elasticDensity = std::make_shared<Elasticity::MooneyRivlinDensity<dim,ValueType>>(parameters);
  auto elasticEnergy = std::make_shared<Elasticity::LocalIntegralEnergy<LocalView, adouble>>(elasticDensity);
  auto localADOLCStiffness = std::make_shared<Elasticity::LocalADOLCStiffness<LocalView>>(elasticEnergy);
  Elasticity::FEAssembler<Basis,VectorType> assembler(basis, localADOLCStiffness);

  //////////////////////////////////////////////////////////////
  //  Compute the energy, assemble and compare!
  //////////////////////////////////////////////////////////////
  
  VectorType gradient;
  MatrixType hessianMatrix;
  double energy = assembler.computeEnergy(x);

  if ( std::abs(energy - expectedEnergy)/expectedEnergy > 1e-3)
  {
      std::cerr << std::setprecision(9);
      std::cerr << "In " << mooneyrivlin_energy << ": Calculated energy is " << energy << " but '" << expectedEnergy << "' was expected!" << std::endl;
      return 1;
  }

  assembler.assembleGradientAndHessian(x, gradient, hessianMatrix, true);

  double gradientTwoNorm = gradient.two_norm();
  double gradientInfinityNorm = gradient.infinity_norm();
  double matrixFrobeniusNorm = hessianMatrix.frobenius_norm();
  if ( std::abs(gradientTwoNorm - expectedGradientTwoNorm)/expectedGradientTwoNorm > 1e-4 || std::abs(gradientInfinityNorm - expectedGradientInfinityNorm)/expectedGradientInfinityNorm > 1e-8)
  {
      std::cerr << std::setprecision(9);
      std::cerr << "In " << mooneyrivlin_energy << ": Calculated gradient max norm is " << gradientInfinityNorm << " but '" << expectedGradientInfinityNorm << "' was expected!" << std::endl;
      std::cerr << "In " << mooneyrivlin_energy << ": Calculated gradient norm is " << gradientTwoNorm << " but '" << expectedGradientTwoNorm << "' was expected!" << std::endl;
      return 1;
  }

  if ( std::abs(matrixFrobeniusNorm - expectedMatrixFrobeniusNorm)/expectedMatrixFrobeniusNorm > 1e-4)
  {
      std::cerr << std::setprecision(9);
      std::cerr << "In " << mooneyrivlin_energy << ": Calculated matrix norm is " << matrixFrobeniusNorm << " but '" << expectedMatrixFrobeniusNorm << "' was expected!" << std::endl;
      return 1;
  }

  return 0;
}

int main (int argc, char *argv[])
{

  Dune::MPIHelper& mpiHelper = MPIHelper::instance(argc, argv);

  // ///////////////////////////////////////
  //    Create the grid
  // ///////////////////////////////////////
  using GridType = UGGrid<dim>;
  auto grid = StructuredGridFactory<GridType>::createCubeGrid({0,0,0}, {1,1,1}, {3,3,3});
  int refine = 3;
  while(refine > 0) {
    for (auto&& e : elements(grid->leafGridView())){
      bool refine = false;
      for (int i = 0; i < e.geometry().corners(); i++) {
          refine = refine || (e.geometry().corner(i)[0] > 0.99);
      }
      grid->mark(refine ? 1 : 0, e);
    }
    grid->adapt();
    refine--;
  }
  grid->loadBalance();

  using GridView = GridType::LeafGridView;
  GridView gridView = grid->leafGridView();

  /////////////////////////////////////////////////////////////////////////
  //  Create a power basis
  /////////////////////////////////////////////////////////////////////////

  using namespace Functions::BasisFactory;
  auto basis = makeBasis(
    gridView,
    power<dim>(
      lagrange<order>(),
      blockedInterleaved()
  ));

  using PowerBasis = decltype(basis);

  //////////////////////////////////////////////////////////////
  //  Prepare the iterate x where we want to assemble
  //////////////////////////////////////////////////////////////
  VectorType x(basis.size());
  Functions::interpolate(basis, x, [](FieldVector<double,dim> x){
    FieldVector<double,dim> y = {0.5*x[0]*x[0], 0.5*x[1] + 0.2*x[0], 4*std::abs(std::sin(x[2]))};
    return y;
  });

  ParameterTree parameters;
  parameters["mooneyrivlin_10"] = "1.67e+6";
  parameters["mooneyrivlin_01"] = "1.94e+6";
  parameters["mooneyrivlin_20"] = "2.42e+6";
  parameters["mooneyrivlin_02"] = "6.52e+6";
  parameters["mooneyrivlin_11"] = "-7.34e+6";
  parameters["mooneyrivlin_30"] = "0";
  parameters["mooneyrivlin_03"] = "0";
  parameters["mooneyrivlin_21"] = "0";
  parameters["mooneyrivlin_12"] = "0";
  parameters["mooneyrivlin_k"] = "75e+6";

  parameters["mooneyrivlin_energy"] = "square";
  double expectedEnergy = 167636683;
  double expectedGradientTwoNorm = 2.10685704e+09;
  double expectedGradientInfinityNorm = 589695526;
  double expectedMatrixFrobeniusNorm = 1.61941167e+11;
  int testSquare = assembleAndCompare(basis, parameters, x, expectedEnergy, expectedGradientTwoNorm, expectedGradientInfinityNorm, expectedMatrixFrobeniusNorm);

  parameters["mooneyrivlin_energy"] = "log";
  expectedEnergy = 181180273;
  expectedGradientTwoNorm = 2.29399362e+09;
  expectedGradientInfinityNorm = 648551554;
  expectedMatrixFrobeniusNorm = 1.67152642e+11;
  int testLog = assembleAndCompare(basis, parameters, x, expectedEnergy, expectedGradientTwoNorm, expectedGradientInfinityNorm, expectedMatrixFrobeniusNorm);

  parameters["mooneyrivlin_energy"] = "ciarlet";
  parameters["mooneyrivlin_a"] = "2.42e+6";
  parameters["mooneyrivlin_b"] = "6.52e+6";
  parameters["mooneyrivlin_c"] = "-7.34e+6";
  expectedEnergy = 83069427.2;
  expectedGradientTwoNorm = 163210957;
  expectedGradientInfinityNorm = 49284369.2;
  expectedMatrixFrobeniusNorm = 5.92126562e+09;
  int testCiarlet = assembleAndCompare(basis, parameters, x, expectedEnergy, expectedGradientTwoNorm, expectedGradientInfinityNorm, expectedMatrixFrobeniusNorm);

  return testCiarlet + testLog + testSquare;
}