/*=========================================================================
 *
 *  Copyright NumFOCUS
 *
 *  Licensed under the Apache License, Version 2.0 (the "License");
 *  you may not use this file except in compliance with the License.
 *  You may obtain a copy of the License at
 *
 *         https://www.apache.org/licenses/LICENSE-2.0.txt
 *
 *  Unless required by applicable law or agreed to in writing, software
 *  distributed under the License is distributed on an "AS IS" BASIS,
 *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 *  See the License for the specific language governing permissions and
 *  limitations under the License.
 *
 *=========================================================================*/

#include "itkFEMObject.h"
#include "itkFEMRobustSolver.h"
#include "itkFEMLoadNoisyLandmark.h"
#include "itkFEMElement2DC0LinearQuadrilateralStrain.h"
#include "itkImageToRectilinearFEMObjectFilter.h"


/**
 * RobustSolver requires a FEMObject as input.
 * In this test, we create a FEMObject manually based on a simple 2D mesh and feature points.
 *
 * In most cases, users have a mesh and feature points rather than the FEMObject as inputs.
 * Users do not need to manually convert the mesh and the feature points into a FEMOjbect.
 * We develop a FEMScatteredDataPointSetToImageFilter to facilitate this conversion.
 *
 *Example:
 * The image is 5x5, the mesh is 2x2, and four feature points represent four
 * cases: two points are on the boundary, one point is inside the element, and
 * one point is the node.
 *
 * 4   ----------------
 *   |    |   |   |   |
 * 3 |----|---|-------*
 *   |    |   |   |   |
 * 2 |----|---*---|---|
 *   |    |   |   |   |
 * 1 |----|---|---*---|
 *   |    |   |   |   |
 *    ----*-----------
 *  0     1   2   3   4
 *
 */
int
itkFEMRobustSolverTest(int, char *[])
{
  constexpr unsigned int DataDimension = 2;
  constexpr unsigned int ParameterDimension = 2;

  /** Solver type alias suppot */
  using SolverType = itk::fem::RobustSolver<DataDimension>;

  /** FEMObject type alias support */
  using FEMObjectType = itk::fem::FEMObject<DataDimension>;

  /** FEM element type alias support */
  using ElementType = itk::fem::Element2DC0LinearQuadrilateralStrain;

  /** FEM node type alias support */
  using NodeType = itk::fem::Element::Node;

  /** FEM Load type alias support */
  using LoadType = itk::fem::LoadNoisyLandmark;

  /** FEM material type alias support */
  using MaterialType = itk::fem::MaterialLinearElasticity;

  /** FEM element type alias support */
  using FEMVectorType = itk::fem::Element::VectorType;

  /** FEM container type alias support */
  using LoadContainerType = FEMObjectType::LoadContainerType;
  using NodeContainerType = FEMObjectType::NodeContainerType;
  using ElementContainerType = FEMObjectType::ElementContainerType;
  using MaterialContainerType = FEMObjectType::MaterialContainerType;

  /** intepolation grid type alias support */
  using InterpolationGridType = itk::Image<itk::fem::Element::ConstPointer, ParameterDimension>;

  auto solver = SolverType::New();

  auto femObject = FEMObjectType::New();

  /** initialize material */
  MaterialContainerType * materialContainer = femObject->GetModifiableMaterialContainer();

  if (!materialContainer)
  {
    std::cerr << "Missing material container!" << std::endl;
    return EXIT_FAILURE;
  }

  materialContainer->Initialize();

  auto material = MaterialType::New();
  material->SetYoungsModulus(3000.0);
  material->SetPoissonsRatio(0.45);

  /** fix material to linear elasticity */
  femObject->AddNextMaterial(material);

  itk::FEMFactoryBase::GetFactory()->RegisterDefaultTypes();

  /** initialize nodes */
  NodeContainerType * nodeContainer = femObject->GetModifiableNodeContainer();

  if (!nodeContainer)
  {
    std::cerr << "Missing node container!" << std::endl;
    return EXIT_FAILURE;
  }

  nodeContainer->Initialize();

  unsigned int elementDimensionX = 2;
  unsigned int elementDimensionY = 2;

  FEMVectorType point(ParameterDimension);
  unsigned int  globalNumbering = 0;

  for (unsigned int j = 0; j <= elementDimensionY; ++j)
  {
    for (unsigned int i = 0; i <= elementDimensionX; ++i)
    {
      point[0] = i;
      point[1] = j;

      auto n = NodeType::New();
      n->SetCoordinates(point);
      n->SetGlobalNumber(globalNumbering);

      femObject->AddNextNode(n);

      globalNumbering++;
    }
  }

  /** initialize elements */
  ElementContainerType * elementContainer = femObject->GetModifiableElementContainer();

  if (!elementContainer)
  {
    std::cerr << "Missing element container!" << std::endl;
    return EXIT_FAILURE;
  }

  elementContainer->Initialize();

  globalNumbering = 0;

  for (unsigned int j = 0; j < elementDimensionY; ++j)
  {
    for (unsigned int i = 0; i < elementDimensionX; ++i)
    {
      unsigned int leftBottomNodeIndex = i + (elementDimensionX + 1) * j;
      unsigned int rightBottomNodeIndex = i + 1 + (elementDimensionX + 1) * j;
      unsigned int rigthUpperNodeIndex = i + 1 + (elementDimensionX + 1) * (j + 1);
      unsigned int leftUpperNodeIndex = i + (elementDimensionX + 1) * (j + 1);

      auto quadrilateral = ElementType::New();

      quadrilateral->SetNode(0, femObject->GetNode(leftBottomNodeIndex));
      quadrilateral->SetNode(1, femObject->GetNode(rightBottomNodeIndex));
      quadrilateral->SetNode(2, femObject->GetNode(rigthUpperNodeIndex));
      quadrilateral->SetNode(3, femObject->GetNode(leftUpperNodeIndex));

      quadrilateral->SetGlobalNumber(globalNumbering);
      quadrilateral->SetMaterial(static_cast<MaterialType *>(femObject->GetMaterial(0).GetPointer()));

      femObject->AddNextElement(quadrilateral);

      globalNumbering++;
    }
  }

  /** initialize loads */
  LoadContainerType * loadContainer = femObject->GetModifiableLoadContainer();

  if (!loadContainer)
  {
    std::cerr << "Missing load container!" << std::endl;
    return EXIT_FAILURE;
  }

  loadContainer->Initialize();

  FEMVectorType featurePoint0(ParameterDimension);
  FEMVectorType featurePoint1(ParameterDimension);
  FEMVectorType featurePoint2(ParameterDimension);
  FEMVectorType featurePoint3(ParameterDimension);

  FEMVectorType displacement0(DataDimension);
  FEMVectorType displacement1(DataDimension);
  FEMVectorType displacement2(DataDimension);
  FEMVectorType displacement3(DataDimension);

  // feature point is on the bottom boundary
  featurePoint0[0] = 1.0;
  featurePoint0[1] = 0.0;
  // feature point is inside an element
  featurePoint1[0] = 3.0;
  featurePoint1[1] = 1.0;
  // feature point is the node
  featurePoint2[0] = 2.0;
  featurePoint2[1] = 2.0;
  // feature point is on the right boundary
  featurePoint3[0] = 4.0;
  featurePoint3[1] = 3.0;

  // displacement associated with the feature point on the bottorm boundary
  displacement0[0] = 1.0;
  displacement0[1] = 1.0;
  // displacement associated with the feature point inside the element
  displacement1[0] = 1.0;
  displacement1[1] = 1.0;
  // displacement associated with the feature point, coincidentally being the node of the mesh
  displacement2[0] = 1.0;
  displacement2[1] = 1.0;
  // displacement associated with the feature point on the right boundary
  displacement3[0] = 1.0;
  displacement3[1] = 1.0;

  auto load0 = LoadType::New();
  load0->SetSource(featurePoint0);
  load0->SetRealDisplacement(displacement0);

  auto load1 = LoadType::New();
  load1->SetSource(featurePoint1);
  load1->SetRealDisplacement(displacement1);

  auto load2 = LoadType::New();
  load2->SetSource(featurePoint2);
  load2->SetRealDisplacement(displacement2);

  auto load3 = LoadType::New();
  load3->SetSource(featurePoint3);
  load3->SetRealDisplacement(displacement3);

  femObject->AddNextLoad(itk::fem::Load::Pointer(load0));
  femObject->AddNextLoad(itk::fem::Load::Pointer(load1));
  femObject->AddNextLoad(itk::fem::Load::Pointer(load2));
  femObject->AddNextLoad(itk::fem::Load::Pointer(load3));

  /** finialize mesh to produce global DOF */
  femObject->FinalizeMesh();

  /** set interpolation grid */

  InterpolationGridType::PointType origin;
  origin[0] = 0.0;
  origin[1] = 0.0;
  solver->SetOrigin(origin);

  InterpolationGridType::SpacingType spacing;
  spacing[0] = 1.0;
  spacing[1] = 1.0;
  solver->SetSpacing(spacing);

  InterpolationGridType::SizeType size;
  size[0] = 5;
  size[1] = 5;
  InterpolationGridType::IndexType start;
  start[0] = 0;
  start[1] = 0;
  InterpolationGridType::RegionType region;
  region.SetSize(size);
  region.SetIndex(start);
  solver->SetRegion(region);

  InterpolationGridType::DirectionType direction;
  direction[0][0] = 1.0;
  direction[0][1] = 0.0;
  direction[1][0] = 0.0;
  direction[1][1] = 1.0;
  solver->SetDirection(direction);

  // if the feature points are the grid point of the interpolation grid, set true.
  // note that since feature points come from the image, this setting is always true.
  solver->SetUseInterpolationGrid(true);

  solver->SetInput(femObject);
  solver->Update();

  int           numOfDOF = femObject->GetNumberOfDegreesOfFreedom();
  FEMVectorType solution(numOfDOF);

  bool  hasError = false;
  float groundTruthSolution[18] = { 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0,
                                    1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0 };

  for (int i = 0; i < numOfDOF; ++i)
  {
    solution[i] = solver->GetSolution(i);

    std::cout << "Solution[" << i << "]:" << solution[i] << std::endl;

    if (itk::Math::abs(groundTruthSolution[i] - solution[i]) > 0.0001)
    {
      std::cerr << "ERROR: Index " << i << ". Groundtruth " << groundTruthSolution[i] << " Solution " << solution[i]
                << std::endl;
      hasError = true;
    }
  }

  if (hasError)
  {
    std::cerr << "Test FAILED!" << std::endl;
    return EXIT_FAILURE;
  }

  std::cout << "Test PASSED!" << std::endl;
  return EXIT_SUCCESS;
}