// @HEADER
//
// ***********************************************************************
//
//        MueLu: A package for multigrid based preconditioning
//                  Copyright 2012 Sandia Corporation
//
// Under the terms of Contract DE-AC04-94AL85000 with Sandia Corporation,
// the U.S. Government retains certain rights in this software.
//
// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions are
// met:
//
// 1. Redistributions of source code must retain the above copyright
// notice, this list of conditions and the following disclaimer.
//
// 2. Redistributions in binary form must reproduce the above copyright
// notice, this list of conditions and the following disclaimer in the
// documentation and/or other materials provided with the distribution.
//
// 3. Neither the name of the Corporation nor the names of the
// contributors may be used to endorse or promote products derived from
// this software without specific prior written permission.
//
// THIS SOFTWARE IS PROVIDED BY SANDIA CORPORATION "AS IS" AND ANY
// EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
// PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL SANDIA CORPORATION OR THE
// CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
// EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
// PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
// PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
// LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
// NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
// SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
//
// Questions? Contact
//                    Jonathan Hu       (jhu@sandia.gov)
//                    Andrey Prokopenko (aprokop@sandia.gov)
//                    Ray Tuminaro      (rstumin@sandia.gov)
//
// ***********************************************************************
//
// @HEADER

// This example demonstrates how to clone a Hierarchy from one type of compute
// node to another.

#include <iostream>

#include <Xpetra_MultiVectorFactory.hpp>

// Galeri
#include <Galeri_XpetraParameters.hpp>
#include <Galeri_XpetraProblemFactory.hpp>
#include <Galeri_XpetraUtils.hpp>
#include <Galeri_XpetraMaps.hpp>
//

#include <MueLu.hpp>
#include <MueLu_Level.hpp>
#include <MueLu_BaseClass.hpp>
#include <MueLu_ParameterListInterpreter.hpp>  // TODO: move into MueLu.hpp
#include "MueLu_SmootherFactory.hpp"
#include "MueLu_TrilinosSmoother.hpp"
#include "MueLu_Ifpack2Smoother.hpp"
#include <MueLu_Utilities.hpp>

#include <MueLu_MutuallyExclusiveTime.hpp>

#ifdef HAVE_MUELU_BELOS
#include <BelosConfigDefs.hpp>
#include <BelosLinearProblem.hpp>
#include <BelosBlockCGSolMgr.hpp>
#include <BelosBlockGmresSolMgr.hpp>
#include <BelosXpetraAdapter.hpp>  // => This header defines Belos::XpetraOp
#include <BelosMueLuAdapter.hpp>   // => This header defines Belos::MueLuOp
#endif

int main(int argc, char* argv[]) {
#include <MueLu_UseDefaultTypes.hpp>
#include <MueLu_UseShortNames.hpp>

  using Teuchos::RCP;  // reference count pointers
  using Teuchos::rcp;
  using Teuchos::TimeMonitor;

  // =========================================================================
  // MPI initialization using Teuchos
  // =========================================================================
  Teuchos::GlobalMPISession mpiSession(&argc, &argv, NULL);
  RCP<const Teuchos::Comm<int> > comm = Teuchos::DefaultComm<int>::getComm();

  // =========================================================================
  // Convenient definitions
  // =========================================================================
  SC zero = Teuchos::ScalarTraits<SC>::zero(), one = Teuchos::ScalarTraits<SC>::one();

  // Instead of checking each time for rank, create a rank 0 stream
  RCP<Teuchos::FancyOStream> fancy = Teuchos::fancyOStream(Teuchos::rcpFromRef(std::cout));
  Teuchos::FancyOStream& fancyout  = *fancy;
  fancyout.setOutputToRootOnly(0);

  // =========================================================================
  // Parameters initialization
  // =========================================================================
  Teuchos::CommandLineProcessor clp(false);

  GO nx = 100, ny = 100, nz = 100;
  Galeri::Xpetra::Parameters<GO> matrixParameters(clp, nx, ny, nz, "Laplace2D");  // manage parameters of the test case
  Xpetra::Parameters xpetraParameters(clp);                                       // manage parameters of Xpetra

  std::string xmlFileName = "scalingTest.xml";
  clp.setOption("xml", &xmlFileName, "read parameters from a file. Otherwise, this example uses by default 'scalingTest.xml'");
  int amgAsPrecond = 1;
  clp.setOption("precond", &amgAsPrecond, "apply multigrid as preconditioner");
  int amgAsSolver = 0;
  clp.setOption("fixPoint", &amgAsSolver, "apply multigrid as solver");
  bool printTimings = true;
  clp.setOption("timings", "notimings", &printTimings, "print timings to screen");
  int writeMatricesOPT = -2;
  clp.setOption("write", &writeMatricesOPT, "write matrices to file (-1 means all; i>=0 means level i)");
  double tol = 1e-12;
  clp.setOption("tol", &tol, "solver convergence tolerance");
  std::string krylovMethod = "cg";
  clp.setOption("krylov", &krylovMethod, "outer Krylov method");
  std::string optSmooType = "cheby";
  clp.setOption("smooType", &optSmooType, "smoother type ('l1-sgs', 'sgs 'or 'cheby')");
  int optSweeps = 2;
  clp.setOption("sweeps", &optSweeps, "sweeps to be used in SGS (or Chebyshev degree)");

  switch (clp.parse(argc, argv)) {
    case Teuchos::CommandLineProcessor::PARSE_HELP_PRINTED: return EXIT_SUCCESS;
    case Teuchos::CommandLineProcessor::PARSE_ERROR:
    case Teuchos::CommandLineProcessor::PARSE_UNRECOGNIZED_OPTION: return EXIT_FAILURE;
    case Teuchos::CommandLineProcessor::PARSE_SUCCESSFUL: break;
  }

  fancyout << "========================================================\n"
           << xpetraParameters << matrixParameters;

  // =========================================================================
  // Problem construction
  // =========================================================================
  RCP<TimeMonitor> globalTimeMonitor = rcp(new TimeMonitor(*TimeMonitor::getNewTimer("ScalingTest: S - Global Time"))), tm;

  comm->barrier();
  tm = rcp(new TimeMonitor(*TimeMonitor::getNewTimer("ScalingTest: 1 - Matrix Build")));

  RCP<const Map> map;
  RCP<MultiVector> coordinates;

  // Retrieve matrix parameters (they may have been changed on the command line), and pass them to Galeri.
  // Galeri will attempt to create a square-as-possible distribution of subdomains di, e.g.,
  //                                 d1  d2  d3
  //                                 d4  d5  d6
  //                                 d7  d8  d9
  //                                 d10 d11 d12
  // A perfect distribution is only possible when the #processors is a perfect square.
  // This *will* result in "strip" distribution if the #processors is a prime number or if the factors are very different in
  // size. For example, np=14 will give a 7-by-2 distribution.
  // If you don't want Galeri to do this, specify mx or my on the galeriList.
  Teuchos::ParameterList pl = matrixParameters.GetParameterList();
  Teuchos::ParameterList galeriList;
  galeriList.set("nx", pl.get("nx", nx));
  galeriList.set("ny", pl.get("ny", ny));
  galeriList.set("nz", pl.get("nz", nz));
  // galeriList.set("mx", comm->getSize());
  // galeriList.set("my", 1);

  // Create map and coordinates
  // In the future, we hope to be able to first create a Galeri problem, and then request map and coordinates from it
  // At the moment, however, things are fragile as we hope that the Problem uses same map and coordinates inside
  if (matrixParameters.GetMatrixType() == "Laplace1D") {
    map         = Galeri::Xpetra::CreateMap<LO, GO, Node>(xpetraParameters.GetLib(), "Cartesian1D", comm, galeriList);
    coordinates = Galeri::Xpetra::Utils::CreateCartesianCoordinates<SC, LO, GO, Map, MultiVector>("1D", map, matrixParameters.GetParameterList());
  } else if (matrixParameters.GetMatrixType() == "Laplace2D" || matrixParameters.GetMatrixType() == "Star2D" || matrixParameters.GetMatrixType() == "Elasticity2D") {
    map         = Galeri::Xpetra::CreateMap<LO, GO, Node>(xpetraParameters.GetLib(), "Cartesian2D", comm, galeriList);
    coordinates = Galeri::Xpetra::Utils::CreateCartesianCoordinates<SC, LO, GO, Map, MultiVector>("2D", map, matrixParameters.GetParameterList());
  } else if (matrixParameters.GetMatrixType() == "Laplace3D" || matrixParameters.GetMatrixType() == "Elasticity3D") {
    map         = Galeri::Xpetra::CreateMap<LO, GO, Node>(xpetraParameters.GetLib(), "Cartesian3D", comm, galeriList);
    coordinates = Galeri::Xpetra::Utils::CreateCartesianCoordinates<SC, LO, GO, Map, MultiVector>("3D", map, matrixParameters.GetParameterList());
  }
  // Expand map to do multiple DOF per node for block problems
  if (matrixParameters.GetMatrixType() == "Elasticity2D")
    map = Xpetra::MapFactory<LO, GO, Node>::Build(map, 2);
  if (matrixParameters.GetMatrixType() == "Elasticity3D")
    map = Xpetra::MapFactory<LO, GO, Node>::Build(map, 3);

  if (comm->getRank() == 0) {
    GO mx = galeriList.get("mx", -1);
    GO my = galeriList.get("my", -1);
    GO mz = galeriList.get("mz", -1);
    fancyout << "Processor subdomains in x direction: " << mx << std::endl
             << "Processor subdomains in y direction: " << my << std::endl
             << "Processor subdomains in z direction: " << mz << std::endl
             << "========================================================" << std::endl;
  }

  Teuchos::ParameterList matrixParams = matrixParameters.GetParameterList();
  matrixParams.set("mx", galeriList.get("mx", -1));
  matrixParams.set("my", galeriList.get("my", -1));
  matrixParams.set("mz", galeriList.get("mz", -1));
  if (matrixParameters.GetMatrixType() == "Elasticity2D" || matrixParameters.GetMatrixType() == "Elasticity3D") {
    // Our default test case for elasticity: all boundaries of a square/cube have Neumann b.c. except left which has Dirichlet
    matrixParams.set("right boundary", "Neumann");
    matrixParams.set("bottom boundary", "Neumann");
    matrixParams.set("top boundary", "Neumann");
    matrixParams.set("front boundary", "Neumann");
    matrixParams.set("back boundary", "Neumann");
  }

  RCP<Galeri::Xpetra::Problem<Map, CrsMatrixWrap, MultiVector> > Pr =
      Galeri::Xpetra::BuildProblem<SC, LO, GO, Map, CrsMatrixWrap, MultiVector>(matrixParameters.GetMatrixType(), map, matrixParams);
  RCP<Matrix> A = Pr->BuildMatrix();

  RCP<MultiVector> nullspace = MultiVectorFactory::Build(map, 1);
  if (matrixParameters.GetMatrixType() == "Elasticity2D" ||
      matrixParameters.GetMatrixType() == "Elasticity3D") {
    nullspace = Pr->BuildNullspace();
    A->SetFixedBlockSize((matrixParameters.GetMatrixType() == "Elasticity2D") ? 2 : 3);

  } else {
    nullspace->putScalar(one);
  }

  fancyout << "Galeri complete.\n========================================================" << std::endl;

  // =========================================================================
  // Preconditioner construction
  // =========================================================================

  // Multigrid Hierarchy
  RCP<Hierarchy> H = rcp(new Hierarchy(A));
  H->setDefaultVerbLevel(Teuchos::VERB_HIGH);
  FactoryManager M;

  // Smoothers
  std::string ifpackType;
  Teuchos::ParameterList ifpackList;
  ifpackList.set("relaxation: sweeps", (LO)optSweeps);
  ifpackList.set("relaxation: damping factor", (SC)1.0);
  if (optSmooType == "sgs") {
    ifpackType = "RELAXATION";
    ifpackList.set("relaxation: type", "Symmetric Gauss-Seidel");
  } else if (optSmooType == "l1-sgs") {
    ifpackType = "RELAXATION";
    ifpackList.set("relaxation: type", "Symmetric Gauss-Seidel");
    ifpackList.set("relaxation: use l1", true);
  } else if (optSmooType == "cheby") {
    ifpackType = "CHEBYSHEV";
    ifpackList.set("chebyshev: degree", (LO)optSweeps);

    if (matrixParameters.GetMatrixType() == "Laplace1D") {
      ifpackList.set("chebyshev: ratio eigenvalue", (SC)3);
    } else if (matrixParameters.GetMatrixType() == "Laplace2D") {
      ifpackList.set("chebyshev: ratio eigenvalue", (SC)7);
    } else if (matrixParameters.GetMatrixType() == "Laplace3D") {
      ifpackList.set("chebyshev: ratio eigenvalue", (SC)20);
    }
  }
  RCP<SmootherPrototype> smootherPrototype = rcp(new Ifpack2Smoother(ifpackType, ifpackList));
  M.SetFactory("Smoother", rcp(new SmootherFactory(smootherPrototype)));

  // create coarsest smoother
  RCP<SmootherPrototype> coarsestSmooProto;
  Teuchos::ParameterList coarsestSmooList;
  coarsestSmooProto                     = rcp(new Ifpack2Smoother("RILUK", coarsestSmooList));
  RCP<SmootherFactory> coarsestSmooFact = rcp(new SmootherFactory(coarsestSmooProto, Teuchos::null));
  M.SetFactory("CoarseSolver", coarsestSmooFact);

  int startLevel   = 0;
  int optMaxLevels = 10;
  H->Setup(M, startLevel, optMaxLevels);

  // Print out the hierarchy stats. We should not need this line, but for some reason the
  // print out in the hierarchy construction does not work.
  H->print(fancyout);

  // =========================================================================
  // System solution (Ax = b)
  // =========================================================================

  RCP<MultiVector> X = MultiVectorFactory::Build(map, 1);
  RCP<MultiVector> B = MultiVectorFactory::Build(map, 1);

  {
    // we set seed for reproducibility
    X->setSeed(846930886);
    X->randomize();
    A->apply(*X, *B, Teuchos::NO_TRANS, one, zero);

    Teuchos::Array<ST::magnitudeType> norms(1);
    B->norm2(norms);
    B->scale(1.0 / norms[0]);
    X->putScalar(zero);
  }

  if (amgAsSolver) {
    H->IsPreconditioner(false);
    H->Iterate(*B, *X, 25);

  } else if (amgAsPrecond) {
#ifdef HAVE_MUELU_BELOS
    // Operator and Multivector type that will be used with Belos
    typedef MultiVector MV;
    typedef Belos::OperatorT<MV> OP;
    H->IsPreconditioner(true);

    // Define Operator and Preconditioner
    Teuchos::RCP<OP> belosOp   = Teuchos::rcp(new Belos::XpetraOp<SC, LO, GO, NO>(A));  // Turns a Xpetra::Matrix object into a Belos operator
    Teuchos::RCP<OP> belosPrec = Teuchos::rcp(new Belos::MueLuOp<SC, LO, GO, NO>(H));   // Turns a MueLu::Hierarchy object into a Belos operator

    // Construct a Belos LinearProblem object
    RCP<Belos::LinearProblem<SC, MV, OP> > belosProblem = rcp(new Belos::LinearProblem<SC, MV, OP>(belosOp, X, B));
    belosProblem->setLeftPrec(belosPrec);

    bool set = belosProblem->setProblem();
    if (set == false) {
      fancyout << "\nERROR:  Belos::LinearProblem failed to set up correctly!" << std::endl;
      return EXIT_FAILURE;
    }

    // Belos parameter list
    int maxIts = 2000;
    Teuchos::ParameterList belosList;
    belosList.set("Maximum Iterations", maxIts);  // Maximum number of iterations allowed
    belosList.set("Convergence Tolerance", tol);  // Relative convergence tolerance requested
    belosList.set("Verbosity", Belos::Errors + Belos::Warnings + Belos::StatusTestDetails);
    belosList.set("Output Frequency", 1);
    belosList.set("Output Style", Belos::Brief);

    // Create an iterative solver manager
    RCP<Belos::SolverManager<SC, MV, OP> > solver;
    if (krylovMethod == "cg") {
      solver = rcp(new Belos::BlockCGSolMgr<SC, MV, OP>(belosProblem, rcp(&belosList, false)));
    } else if (krylovMethod == "gmres") {
      solver = rcp(new Belos::BlockGmresSolMgr<SC, MV, OP>(belosProblem, rcp(&belosList, false)));
    } else {
      TEUCHOS_TEST_FOR_EXCEPTION(true, MueLu::Exceptions::RuntimeError, "Invalid Krylov method.  Options are \"cg\" or \" gmres\".");
    }

    // Perform solve
    Belos::ReturnType ret = Belos::Unconverged;
    try {
      ret = solver->solve();

      // Get the number of iterations for this solve.
      fancyout << "Number of iterations performed for this solve: " << solver->getNumIters() << std::endl;

    } catch (...) {
      fancyout << std::endl
               << "ERROR:  Belos threw an error! " << std::endl;
    }

    // Check convergence
    if (ret != Belos::Converged)
      fancyout << std::endl
               << "ERROR:  Belos did not converge! " << std::endl;
    else
      fancyout << std::endl
               << "SUCCESS:  Belos converged!" << std::endl;

    // Clone the preconditioner to ThrustGPU node type
    typedef KokkosClassic::ThrustGPUNode NO2;
    typedef MueLu::Hierarchy<SC, LO, GO, NO2> Hierarchy2;
    typedef Xpetra::MultiVector<SC, LO, GO, NO2> MV2;
    typedef Belos::OperatorT<MV2> OP2;

    ParameterList plClone;
    plClone.set<LocalOrdinal>("Verbose", 1);
    RCP<NO2> node           = rcp(new NO2(plClone));
    RCP<Hierarchy2> clonedH = H->clone<NO2>(node);

    // Clone A, X, B to new node type
    RCP<Xpetra::Matrix<SC, LO, GO, NO2> > clonedA = Xpetra::clone(*A, node);
    RCP<MV2> clonedX                              = Xpetra::clone(*X, node);
    clonedX->putScalar(zero);
    RCP<MV2> clonedB = Xpetra::clone(*B, node);
    clonedH->IsPreconditioner(true);

    // Define Operator and Preconditioner
    Teuchos::RCP<OP2> belosOp2   = Teuchos::rcp(new Belos::XpetraOp<SC, LO, GO, NO2>(clonedA));  // Turns a Xpetra::Matrix object into a Belos operator
    Teuchos::RCP<OP2> belosPrec2 = Teuchos::rcp(new Belos::MueLuOp<SC, LO, GO, NO2>(clonedH));   // Turns a MueLu::Hierarchy object into a Belos operator

    // Construct a Belos LinearProblem object
    RCP<Belos::LinearProblem<SC, MV2, OP2> > belosProblem2 = rcp(new Belos::LinearProblem<SC, MV2, OP2>(belosOp2, clonedX, clonedB));
    belosProblem2->setLeftPrec(belosPrec2);

    bool set2 = belosProblem2->setProblem();
    if (set2 == false) {
      fancyout << "\nERROR:  Belos::LinearProblem for new node failed to set up correctly!" << std::endl;
      return EXIT_FAILURE;
    }
    // Create an iterative solver manager
    RCP<Belos::SolverManager<SC, MV2, OP2> > solver2;
    if (krylovMethod == "cg") {
      solver2 = rcp(new Belos::BlockCGSolMgr<SC, MV2, OP2>(belosProblem2, rcp(&belosList, false)));
    } else if (krylovMethod == "gmres") {
      solver2 = rcp(new Belos::BlockGmresSolMgr<SC, MV2, OP2>(belosProblem2, rcp(&belosList, false)));
    } else {
      TEUCHOS_TEST_FOR_EXCEPTION(true, MueLu::Exceptions::RuntimeError, "Invalid Krylov method.  Options are \"cg\" or \" gmres\".");
    }
    // Perform solve
    Belos::ReturnType ret2 = Belos::Unconverged;
    try {
      ret2 = solver2->solve();
      // Get the number of iterations for this solve.
      fancyout << "Number of iterations performed for this solve: " << solver2->getNumIters() << std::endl;
    } catch (...) {
      fancyout << std::endl
               << "ERROR:  Belos threw an error! " << std::endl;
    }
    // Check convergence
    if (ret2 != Belos::Converged)
      fancyout << std::endl
               << "ERROR:  Belos did not converge! " << std::endl;
    else
      fancyout << std::endl
               << "SUCCESS:  Belos converged!" << std::endl;

    // Determine if example passed
    RCP<KokkosClassic::DefaultNode::DefaultNodeType> defaultNode =
        rcp(new Tpetra::KokkosClassic::DefaultNode::DefaultNodeType(pl));
    RCP<MV> clonedXcpu = Xpetra::clone(*clonedX, defaultNode);
    clonedXcpu->update(1.0, *X, -1.0);
    Scalar norm;
    clonedXcpu->norm2(Teuchos::arrayView(&norm, 1));
    std::cout << "\nNorm of serial node soln - ThrustGPU node soln = "
              << norm << std::endl;

    bool passed = false;
    if (norm <= Scalar(1e-10))
      passed = true;
    if (passed)
      std::cout << "Example Passed!" << std::endl;
    else
      std::cout << "Example Failed!" << std::endl;
  }
#endif  // ifdef HAVE_MUELU_BELOS

  return 0;
}  // main
