/*=========================================================================
 *
 *  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.
 *
 *=========================================================================*/

#ifndef itkFEMSolverHyperbolic_hxx
#define itkFEMSolverHyperbolic_hxx

#include "itkMath.h"

namespace itk
{
namespace fem
{

template <unsigned int VDimension>
SolverHyperbolic<VDimension>::SolverHyperbolic()
{
  this->InitializeLinearSystemWrapper();
  this->InitializeMatrixForAssembly(2);
  this->m_Beta = 0.25;
  this->m_Gamma = 0.5;
  this->m_TimeStep = 1.0;
  this->m_NumberOfIterations = 1;
}

template <unsigned int VDimension>
void
SolverHyperbolic<VDimension>::InitializeLinearSystemWrapper()
{
  // Set the maximum number of matrices and vectors that
  // we will need to store inside.
  this->m_LinearSystem->SetNumberOfMatrices(5);
  this->m_LinearSystem->SetNumberOfVectors(6);
  this->m_LinearSystem->SetNumberOfSolutions(3);
}

template <unsigned int VDimension>
void
SolverHyperbolic<VDimension>::AssembleElementMatrix(Element::Pointer e)
{
  // Copy the element stiffness matrix for faster access.
  Element::MatrixType Ke;
  e->GetStiffnessMatrix(Ke);
  Element::MatrixType Me;
  e->GetMassMatrix(Me);

  // ... same for number of DOF
  int Ne = e->GetNumberOfDegreesOfFreedom();

  // Step over all rows in element matrix
  for (int j = 0; j < Ne; ++j)
  {
    // Step over all columns in element matrix
    for (int k = 0; k < Ne; ++k)
    {
      // error checking. all GFN should be =>0 and <NGFN
      if (e->GetDegreeOfFreedom(j) >= this->GetInput()->GetNumberOfDegreesOfFreedom() ||
          e->GetDegreeOfFreedom(k) >= this->GetInput()->GetNumberOfDegreesOfFreedom())
      {
        throw FEMExceptionSolution(__FILE__, __LINE__, "Solver::AssembleElementMatrix()", "Illegal GFN!");
      }

      //
      // Here we finally update the corresponding element
      // in the master stiffness matrix. We first check if
      // element in Ke is zero, to prevent zeros from being
      // allocated in sparse matrix.
      //
      if (Math::NotExactlyEquals(Ke[j][k], Float(0.0)))
      {
        this->m_LinearSystem->AddMatrixValue(e->GetDegreeOfFreedom(j), e->GetDegreeOfFreedom(k), Ke[j][k], matrix_K);
      }
      if (Math::NotExactlyEquals(Me[j][k], Float(0.0)))
      {
        this->m_LinearSystem->AddMatrixValue(e->GetDegreeOfFreedom(j), e->GetDegreeOfFreedom(k), Me[j][k], matrix_M);
      }
    }
  }
}

template <unsigned int VDimension>
void
SolverHyperbolic<VDimension>::InitializeMatrixForAssembly(unsigned int N)
{
  this->m_LinearSystem->SetSystemOrder(N);
  this->m_LinearSystem->InitializeMatrix();
  this->m_LinearSystem->InitializeMatrix(matrix_K);
  this->m_LinearSystem->InitializeMatrix(matrix_M);
  this->m_LinearSystem->InitializeMatrix(matrix_C);
  for (unsigned int i = 0; i < N; ++i)
  {
    this->m_LinearSystem->SetMatrixValue(i, i, 1.0, matrix_C);
  }
}

template <unsigned int VDimension>
void
SolverHyperbolic<VDimension>::FinalizeMatrixAfterAssembly()
{
  // Apply the boundary conditions to the matrix
  // FIXME: this doesn't work in general
  this->ApplyBC(0, matrix_M);
  this->ApplyBC(0, matrix_K);

  // Calculate initial values of vector_a
  // M*a0=F - C*v0 - K*d0
  // FIXME: take into account the d0 and v0.
  this->m_LinearSystem->InitializeSolution(0);
  this->m_LinearSystem->InitializeSolution(solution_a);

  this->m_LinearSystem->CopyMatrix(matrix_M, 0);
  this->AssembleF();
  this->m_LinearSystem->Solve();
  this->m_LinearSystem->InitializeVector(vector_tmp);
  this->m_LinearSystem->CopySolution2Vector(0, vector_tmp);
  this->m_LinearSystem->InitializeSolution(solution_a);
  this->m_LinearSystem->CopyVector2Solution(vector_tmp, solution_a);
  this->m_LinearSystem->DestroyVector(vector_tmp);

  this->m_LinearSystem->InitializeSolution(solution_d);
  this->m_LinearSystem->InitializeSolution(solution_v);

  // Compose the lhs of system of lin. eq.
  this->m_LinearSystem->InitializeMatrix(matrix_tmp);
  this->m_LinearSystem->CopyMatrix(matrix_C, matrix_tmp);
  this->m_LinearSystem->ScaleMatrix(this->m_Gamma * this->m_TimeStep, matrix_tmp);
  this->m_LinearSystem->AddMatrixMatrix(0, matrix_tmp);

  this->m_LinearSystem->CopyMatrix(matrix_K, matrix_tmp);
  this->m_LinearSystem->ScaleMatrix(this->m_Beta * this->m_TimeStep * this->m_TimeStep, matrix_tmp);
  this->m_LinearSystem->AddMatrixMatrix(0, matrix_tmp);
  this->m_LinearSystem->DestroyMatrix(matrix_tmp);
}

template <unsigned int VDimension>
void
SolverHyperbolic<VDimension>::Solve()
{
  this->m_LinearSystem->InitializeVector(vector_tmp);
  this->m_LinearSystem->InitializeVector(vector_dhat);
  this->m_LinearSystem->InitializeVector(vector_vhat);
  this->m_LinearSystem->InitializeVector(vector_ahat);

  // We're using the Newmark method to obtain the solution
  // Assume that vectors solution_a solution_v and solution_d contain
  // solutions obtained at the previous time step.

  // Calculate the predictors
  for (unsigned int i = 0; i < this->m_LinearSystem->GetSystemOrder(); ++i)
  {
    Float d0 = this->m_LinearSystem->GetSolutionValue(i, solution_d);
    Float v0 = this->m_LinearSystem->GetSolutionValue(i, solution_v);
    Float a0 = this->m_LinearSystem->GetSolutionValue(i, solution_a);
    this->m_LinearSystem->SetVectorValue(
      i,
      -(d0 + this->m_TimeStep * v0 + 0.5 * this->m_TimeStep * this->m_TimeStep * (1.0 - 2.0 * this->m_Beta) * a0),
      vector_dhat);
    this->m_LinearSystem->SetVectorValue(i, -(v0 + this->m_TimeStep * (1.0 - this->m_Gamma) * a0), vector_vhat);
  }

  // Calculate the rhs of master equation
  this->m_LinearSystem->MultiplyMatrixVector(vector_tmp, matrix_C, vector_vhat);
  this->m_LinearSystem->AddVectorVector(0, vector_tmp);
  this->m_LinearSystem->MultiplyMatrixVector(vector_tmp, matrix_K, vector_dhat);
  this->m_LinearSystem->AddVectorVector(0, vector_tmp);

  // Solve the system of linear equations for accelerations
  this->m_LinearSystem->Solve();

  // move the solution for a to the correct vector
  this->m_LinearSystem->CopySolution2Vector(0, vector_tmp);
  this->m_LinearSystem->CopyVector2Solution(vector_tmp, solution_a);

  // Calculate displacements and velocities
  for (unsigned int i = 0; i < this->m_LinearSystem->GetSystemOrder(); ++i)
  {
    Float dhat = -this->m_LinearSystem->GetVectorValue(i, vector_dhat);
    Float vhat = -this->m_LinearSystem->GetVectorValue(i, vector_vhat);
    Float a1 = this->m_LinearSystem->GetSolutionValue(i, solution_a);

    this->m_LinearSystem->SetSolutionValue(
      i, dhat + this->m_Beta * this->m_TimeStep * this->m_TimeStep * a1, solution_d);

    this->m_LinearSystem->SetSolutionValue(i, vhat + this->m_Gamma * this->m_TimeStep * a1, solution_v);
  }

  this->m_LinearSystem->DestroyVector(vector_tmp);
  this->m_LinearSystem->DestroyVector(vector_dhat);
  this->m_LinearSystem->DestroyVector(vector_vhat);
  this->m_LinearSystem->DestroyVector(vector_ahat);
}

// ----------------------------------------------------------------------------
template <unsigned int VDimension>
void
SolverHyperbolic<VDimension>::GenerateData()
{
  // Call Solver
  this->RunSolver();
}

template <unsigned int VDimension>
void
SolverHyperbolic<VDimension>::RunSolver()
{
  this->AssembleK();  // Assemble the global stiffness matrix K
  this->DecomposeK(); // Invert K

  for (unsigned int nit = 0; nit < m_NumberOfIterations; ++nit)
  {
    this->AssembleF();
    this->Solve();
  }
}

template <unsigned int VDimension>
void
SolverHyperbolic<VDimension>::PrintSelf(std::ostream & os, Indent indent) const
{
  Superclass::PrintSelf(os, indent);
  os << indent << "Number Of Iterations: " << this->m_NumberOfIterations << std::endl;
  os << indent << "Time Step: " << this->m_TimeStep << std::endl;
  os << indent << "Beta: " << this->m_Beta << std::endl;
  os << indent << "Gamma: " << this->m_Gamma << std::endl;
}

} // end namespace fem
} // end namespace itk

#endif