/*=========================================================================
 *
 *  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 "VNLSparseLUSolverTraits.h"
#include "itkMath.h" // itk::Math::Absolute

#include <iostream>
#include <cstdlib>

template <class TVector>
bool
VectorsEquals(const TVector & v1, const TVector & v2, const typename TVector::element_type & tolerance)
{
  if (v1.size() != v2.size())
  {
    std::cerr << "Error: v1.size() != v2.size()" << std::endl;
    return false;
  }

  for (unsigned int i = 0; i < v1.size(); ++i)
  {
    if (itk::Math::Absolute(v1(i) - v2(i)) > tolerance)
    {
      std::cerr << "Error: itk::Math::Absolute( v1(" << i << ") - v2(" << i << ") ) > " << tolerance << std::endl;
      return false;
    }
  }

  return true;
}

int
VNLSparseLUSolverTraitsTest(int, char *[])
{
  /**
   * Define an sparse LU solver traits type that operates over sparse matrices and
   * vectors of type "double"
   */
  using CoordinateType = double;
  using SolverTraits = VNLSparseLUSolverTraits<CoordinateType>;
  using MatrixType = SolverTraits::MatrixType;
  using VectorType = SolverTraits::VectorType;

  /**
   * Build the linear system to solve
   */
  constexpr unsigned int N{ 3 };
  VectorType             Bx = SolverTraits::InitializeVector(N);
  Bx.fill(0.);
  Bx[0] = 2.1;

  VectorType By = SolverTraits::InitializeVector(N);
  By.fill(0.);
  By[1] = 1.1;
  By[2] = -3.;

  VectorType Bz = SolverTraits::InitializeVector(N);
  Bz.fill(0.);
  Bz[0] = 19.4;
  Bz[1] = -4.3;

  MatrixType A = SolverTraits::InitializeSparseMatrix(N, N);
  SolverTraits::FillMatrix(A, 0, 0, 2);
  SolverTraits::FillMatrix(A, 0, 1, -1);
  SolverTraits::FillMatrix(A, 1, 0, -1);
  SolverTraits::FillMatrix(A, 1, 1, 2);
  SolverTraits::FillMatrix(A, 1, 2, -1);
  SolverTraits::FillMatrix(A, 2, 1, -1);
  SolverTraits::FillMatrix(A, 2, 2, 2);

  /**
   * Define the tolerance and expected results
   */
  constexpr CoordinateType tolerance{ 1e-9 };

  VectorType Xexpected(N);
  Xexpected(0) = 1.575;
  Xexpected(1) = 1.05;
  Xexpected(2) = 0.525;

  VectorType Yexpected(N);
  Yexpected(0) = -0.2;
  Yexpected(1) = -0.4;
  Yexpected(2) = -1.7;

  VectorType Zexpected(N);
  Zexpected(0) = 12.4;
  Zexpected(1) = 5.4;
  Zexpected(2) = 2.7;

  /**
   * Test 1: Check the result of A * X = Bx
   */
  {
    VectorType X = SolverTraits::InitializeVector(N);
    SolverTraits::Solve(A, Bx, X);
    if (!VectorsEquals(X, Xexpected, tolerance))
    {
      return EXIT_FAILURE;
    }
  }

  /**
   * Test 2: Check the result of A * X = Bx, A * Y = By
   */
  {
    VectorType X = SolverTraits::InitializeVector(N);
    VectorType Y = SolverTraits::InitializeVector(N);
    SolverTraits::Solve(A, Bx, X);
    SolverTraits::Solve(A, By, Y);
    if (!VectorsEquals(X, Xexpected, tolerance) || !VectorsEquals(Y, Yexpected, tolerance))
    {
      return EXIT_FAILURE;
    }
  }

  /**
   * Test 3: Check the result of A * X = Bx, A * Y = By, A * Z = Bz
   */
  {
    VectorType X = SolverTraits::InitializeVector(N);
    VectorType Y = SolverTraits::InitializeVector(N);
    VectorType Z = SolverTraits::InitializeVector(N);
    SolverTraits::Solve(A, Bx, X);
    SolverTraits::Solve(A, By, Y);
    SolverTraits::Solve(A, Bz, Z);
    if (!VectorsEquals(X, Xexpected, tolerance) || !VectorsEquals(Y, Yexpected, tolerance) ||
        !VectorsEquals(Z, Zexpected, tolerance))
    {
      return EXIT_FAILURE;
    }
  }

  /**
   * Test 4: Check the result of A * X = Bx (reuse the decomposed matrix for multiple back-substitutions)
   */
  {
    VectorType               X = SolverTraits::InitializeVector(N);
    SolverTraits::SolverType solver(A);

    // First back-substitution
    SolverTraits::Solve(solver, Bx, X);
    if (!VectorsEquals(X, Xexpected, tolerance))
    {
      return EXIT_FAILURE;
    }

    // Second back-substitution (reusing the already factored matrix)
    SolverTraits::Solve(solver, Bx, X);
    if (!VectorsEquals(X, Xexpected, tolerance))
    {
      return EXIT_FAILURE;
    }
  }

  /**
   * Test 5: Check the result of A * X = Bx, A * Y = By (reuse the decomposed matrix for multiple back-substitutions)
   */
  {
    VectorType               X = SolverTraits::InitializeVector(N);
    VectorType               Y = SolverTraits::InitializeVector(N);
    SolverTraits::SolverType solver(A);

    // First back-substitution
    SolverTraits::Solve(solver, Bx, X);
    SolverTraits::Solve(solver, By, Y);
    if (!VectorsEquals(X, Xexpected, tolerance) || !VectorsEquals(Y, Yexpected, tolerance))
    {
      return EXIT_FAILURE;
    }

    // Second back-substitution (reusing the already factored matrix)
    SolverTraits::Solve(solver, Bx, X);
    SolverTraits::Solve(solver, By, Y);
    if (!VectorsEquals(X, Xexpected, tolerance) || !VectorsEquals(Y, Yexpected, tolerance))
    {
      return EXIT_FAILURE;
    }
  }

  /**
   * Test 6: Check the result of A * X = Bx, A * Y = By, A * Z = Bz (reuse the decomposed matrix for multiple
   * back-substitutions)
   */
  {
    VectorType               X = SolverTraits::InitializeVector(N);
    VectorType               Y = SolverTraits::InitializeVector(N);
    VectorType               Z = SolverTraits::InitializeVector(N);
    SolverTraits::SolverType solver(A);

    // First back-substitution
    SolverTraits::Solve(solver, Bx, X);
    SolverTraits::Solve(solver, By, Y);
    SolverTraits::Solve(solver, Bz, Z);
    if (!VectorsEquals(X, Xexpected, tolerance) || !VectorsEquals(Y, Yexpected, tolerance) ||
        !VectorsEquals(Z, Zexpected, tolerance))
    {
      return EXIT_FAILURE;
    }

    // Second back-substitution (reusing the already factored matrix)
    SolverTraits::Solve(solver, Bx, X);
    SolverTraits::Solve(solver, By, Y);
    SolverTraits::Solve(solver, Bz, Z);
    if (!VectorsEquals(X, Xexpected, tolerance) || !VectorsEquals(Y, Yexpected, tolerance) ||
        !VectorsEquals(Z, Zexpected, tolerance))
    {
      return EXIT_FAILURE;
    }
  }

  return EXIT_SUCCESS;
}
