/*
 * Copyright (C) 2023
 *
 * Author: Antonin Portelli <antonin.portelli@me.com>
 * Elements based on production templates from Raoul Hodgson
 *
 * Hadrons is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 2 of the License, or
 * (at your option) any later version.
 *
 * Hadrons is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with Hadrons.  If not, see <http://www.gnu.org/licenses/>.
 *
 * See the full license in the file "LICENSE" in the top level distribution
 * directory.
 */

#pragma once

#include <Hadrons/Application.hpp>
#include <Hadrons/Modules.hpp>

namespace hadpresets
{
using namespace Grid::Hadrons;

struct RbcUkqcd
{
  // Ensemble parameters
  struct EnsembleParameters
  {
    double ml, ms, M5, scale;
    unsigned int L, T, Ls;
  };

  // M0
  inline static constexpr EnsembleParameters m0UnitaryPar{0.000678, 0.02661, 1.8, 2., 64, 128, 12};
  inline static constexpr EnsembleParameters m0LCDPar{0.0006203, 0.02661, 1.8, 2., 64, 128, 12};

  // C1M
  inline static constexpr EnsembleParameters c1mLCDPar  {0.005, 0.0362, 1.8, 2., 24, 64, 24};
  inline static constexpr EnsembleParameters c1m16LCDPar{0.005, 0.0362, 1.8, 2., 16, 64, 24};
  inline static constexpr EnsembleParameters c1m20LCDPar{0.005, 0.0362, 1.8, 2., 20, 64, 24};
  inline static constexpr EnsembleParameters c1m32LCDPar{0.005, 0.0362, 1.8, 2., 32, 64, 24};

  // Runtime deflation parameters
  struct DeflationParameters
  {
    double alpha, beta;
    unsigned int nPoly, nStop, nK, nM;
  };

  inline static constexpr DeflationParameters c1mDeflPar  {5.0e-04, 5.5, 101, 100, 110, 120};
  inline static constexpr DeflationParameters c1m16DeflPar{5.0e-03, 5.5, 101, 200, 220, 230};
  inline static constexpr DeflationParameters c1m20DeflPar{3.0e-03, 5.5, 101, 200, 220, 230};
  inline static constexpr DeflationParameters c1m32DeflPar{3.0e-06, 5.5, 101, 100, 110, 120};

  // Light solver: load deflation from disk
  static inline void addM0LightLCDSolver(Application &app, const std::string solverName,
                                         const std::string gaugeName,
                                         const std::string gaugeTransform,
                                         const std::string eigenpackPath,
                                         const double residual = 1.0e-8);
  // Light solver: deflation at runtime
  static inline void
  addLightRuntimeLCDSolver(Application &app, const RbcUkqcd::EnsembleParameters &par,
                           const RbcUkqcd::DeflationParameters &deflPar,
                           const std::string solverName, const std::string gaugeName,
                           const std::string gaugeTransform, const double residual);

  static inline void addC1MLightRuntimeLCDSolver(Application &app, const std::string solverName,
                                                 const std::string gaugeName,
                                                 const std::string gaugeTransform,
                                                 const double residual = 1.0e-8);
  static inline void addC1M16LightRuntimeLCDSolver(Application &app, const std::string solverName,
                                                   const std::string gaugeName,
                                                   const std::string gaugeTransform,
                                                   const double residual = 1.0e-8);
  static inline void addC1M20LightRuntimeLCDSolver(Application &app, const std::string solverName,
                                                   const std::string gaugeName,
                                                   const std::string gaugeTransform,
                                                   const double residual = 1.0e-8);
  static inline void addC1M32LightRuntimeLCDSolver(Application &app, const std::string solverName,
                                                   const std::string gaugeName,
                                                   const std::string gaugeTransform,
                                                   const double residual = 1.0e-8);

  // Strange solver (undeflated)
  static inline void addStrangeSolver(Application &app, const RbcUkqcd::EnsembleParameters &par,
                                      const std::string solverName, const std::string gaugeName,
                                      const std::string gaugeTransform, const double residual);

  static inline void addM0StrangeSolver(Application &app, const std::string solverName,
                                        const std::string gaugeName,
                                        const std::string gaugeTransform,
                                        const double residual = 1.0e-8);
  static inline void addC1MStrangeSolver(Application &app, const std::string solverName,
                                         const std::string gaugeName,
                                         const std::string gaugeTransform,
                                         const double residual = 1.0e-8);
  static inline void addC1M16StrangeSolver(Application &app, const std::string solverName,
                                           const std::string gaugeName,
                                           const std::string gaugeTransform,
                                           const double residual = 1.0e-8);
  static inline void addC1M20StrangeSolver(Application &app, const std::string solverName,
                                           const std::string gaugeName,
                                           const std::string gaugeTransform,
                                           const double residual = 1.0e-8);
  static inline void addC1M32StrangeSolver(Application &app, const std::string solverName,
                                           const std::string gaugeName,
                                           const std::string gaugeTransform,
                                           const double residual = 1.0e-8);
};

// Implementations /////////////////////////////////////////////////////////////////////////////////

// Light M0 (load deflation from disk)
void RbcUkqcd::addM0LightLCDSolver(Application &app, const std::string solverName,
                                   const std::string gaugeName, const std::string gaugeTransform,
                                   const std::string eigenpackPath, const double residual)
{
  const std::string prefix = solverName;

  // Gauge field FP32 cast
  MUtilities::GaugeSinglePrecisionCast::Par gaugeCastPar;

  gaugeCastPar.field = gaugeName;
  app.createModule<MUtilities::GaugeSinglePrecisionCast>(prefix + "_gauge_fp32", gaugeCastPar);

  // Scaled DWF action + FP32 version
  MAction::ScaledDWF::Par actionPar;

  actionPar.gauge = gaugeName;
  actionPar.Ls = RbcUkqcd::m0LCDPar.Ls;
  actionPar.M5 = RbcUkqcd::m0LCDPar.M5;
  actionPar.mass = RbcUkqcd::m0LCDPar.ml;
  actionPar.scale = RbcUkqcd::m0LCDPar.scale;
  actionPar.boundary = "1 1 1 1";
  actionPar.twist = "0. 0. 0. 0.";
  app.createModule<MAction::ScaledDWF>(prefix + "_dwf", actionPar);
  actionPar.gauge = prefix + "_gauge_fp32";
  app.createModule<MAction::ScaledDWFF>(prefix + "_dwf_fp32", actionPar);

  // Compressed eigenpack
  MIO::LoadCoarseFermionEigenPack250F::Par epPar;

  epPar.filestem = eigenpackPath;
  epPar.multiFile = true;
  epPar.redBlack = true;
  epPar.sizeFine = 250;
  epPar.sizeCoarse = 2000;
  epPar.Ls = 12;
  epPar.blockSize = "4 4 4 4 12";
  epPar.orthogonalise = false;
  epPar.gaugeXform = gaugeTransform;
  app.createModule<MIO::LoadCoarseFermionEigenPack250F>(prefix + "_epack", epPar);

  // Inner guesser
  MGuesser::CoarseDeflation250F::Par iguessPar;

  iguessPar.eigenPack = prefix + "_epack";
  iguessPar.size = 2000;
  app.createModule<MGuesser::CoarseDeflation250F>(prefix + "_iguesser", iguessPar);

  // Batched mixed-precision red-black preconditionned CG
  MSolver::MixedPrecisionRBPrecCGBatched::Par solverPar;

  solverPar.innerAction = prefix + "_dwf_fp32";
  solverPar.outerAction = prefix + "_dwf";
  solverPar.maxInnerIteration = 400;
  solverPar.maxOuterIteration = 100;
  solverPar.maxPatchupIteration = 1000;
  solverPar.residual = residual;
  solverPar.updateResidual = true;
  solverPar.innerGuesser = prefix + "_iguesser";
  solverPar.outerGuesser = "";
  app.createModule<MSolver::MixedPrecisionRBPrecCGBatched>(solverName, solverPar);
}

// Light M0
void RbcUkqcd::addM0LightLCDSolver(Application &app, const std::string solverName,
                                   const std::string gaugeName, const std::string gaugeTransform,
                                   const std::string eigenpackPath, const double residual)
{
  RbcUkqcd::addLightLCDSolver(app, RbcUkqcd::m0LCDPar, solverName, gaugeName, gaugeTransform,
                              eigenpackPath, residual);
}

// Deflation at runtime
void RbcUkqcd::addLightRuntimeLCDSolver(Application &app, const RbcUkqcd::EnsembleParameters &par,
                                        const RbcUkqcd::DeflationParameters &deflPar,
                                        const std::string solverName, const std::string gaugeName,
                                        const std::string gaugeTransform, const double residual)
{
  const std::string prefix = solverName;

  // Gauge field FP32 cast
  MUtilities::GaugeSinglePrecisionCast::Par gaugeCastPar;

  gaugeCastPar.field = gaugeName;
  app.createModule<MUtilities::GaugeSinglePrecisionCast>(prefix + "_gauge_fp32", gaugeCastPar);

  // Scaled DWF action + FP32 version
  MAction::ScaledDWF::Par actionPar;

  actionPar.gauge = gaugeName;
  actionPar.Ls = par.Ls;
  actionPar.M5 = par.M5;
  actionPar.mass = par.ml;
  actionPar.scale = par.scale;
  actionPar.boundary = "1 1 1 -1";
  actionPar.twist = "0. 0. 0. 0.";
  app.createModule<MAction::ScaledDWF>(prefix + "_dwf", actionPar);
  actionPar.gauge = prefix + "_gauge_fp32";
  app.createModule<MAction::ScaledDWFF>(prefix + "_dwf_fp32", actionPar);

  // Guesser
  MFermion::Operators::Par opPar;
  opPar.action = prefix + "_dwf";
  app.createModule<MFermion::Operators>(prefix + "_dwf_op", opPar);

  MSolver::FermionImplicitlyRestartedLanczosIo32::Par lanPar;
  lanPar.op = prefix + "_dwf_op_schur";
  lanPar.multiFile = false;
  lanPar.redBlack = true;
  lanPar.lanczosParams.Cheby.alpha = deflPar.alpha;
  lanPar.lanczosParams.Cheby.beta = deflPar.beta;
  lanPar.lanczosParams.Cheby.Npoly = deflPar.nPoly;
  lanPar.lanczosParams.Nstop = deflPar.nStop;
  lanPar.lanczosParams.Nk = deflPar.nK;
  lanPar.lanczosParams.Nm = deflPar.nM;
  lanPar.lanczosParams.resid = 3e-6;
  lanPar.lanczosParams.MaxIt = 10000;
  lanPar.lanczosParams.betastp = 0;
  lanPar.lanczosParams.MinRes = 0;
  lanPar.output = "";
  app.createModule<MSolver::FermionImplicitlyRestartedLanczosIo32>(prefix + "_epack", lanPar);

  MGuesser::ExactDeflation::Par guessPar;
  guessPar.eigenPack = prefix + "_epack";
  guessPar.size = deflPar.nStop;
  app.createModule<MGuesser::ExactDeflation>(prefix + "_defl", guessPar);

  // Mixed-precision red-black preconditionned CG
  MSolver::MixedPrecisionRBPrecCG::Par solverPar;

  solverPar.innerAction = prefix + "_dwf_fp32";
  solverPar.outerAction = prefix + "_dwf";
  solverPar.maxInnerIteration = 30000;
  solverPar.maxOuterIteration = 100;
  solverPar.residual = residual;
  solverPar.innerGuesser = "";
  solverPar.outerGuesser = prefix + "_defl";
  app.createModule<MSolver::MixedPrecisionRBPrecCG>(solverName, solverPar);
}

// Light C1M
void RbcUkqcd::addC1MLightRuntimeLCDSolver(Application &app, const std::string solverName,
                                           const std::string gaugeName,
                                           const std::string gaugeTransform, const double residual)
{
  RbcUkqcd::addLightRuntimeLCDSolver(app, RbcUkqcd::c1mLCDPar, RbcUkqcd::c1mDeflPar, solverName,
                                     gaugeName, gaugeTransform, residual);
}

// Light C1M16
void RbcUkqcd::addC1M16LightRuntimeLCDSolver(Application &app, const std::string solverName,
                                           const std::string gaugeName,
                                           const std::string gaugeTransform, const double residual)
{
  RbcUkqcd::addLightRuntimeLCDSolver(app, RbcUkqcd::c1m16LCDPar, RbcUkqcd::c1m16DeflPar, solverName,
                                     gaugeName, gaugeTransform, residual);
}

// Light C1M20
void RbcUkqcd::addC1M20LightRuntimeLCDSolver(Application &app, const std::string solverName,
                                           const std::string gaugeName,
                                           const std::string gaugeTransform, const double residual)
{
  RbcUkqcd::addLightRuntimeLCDSolver(app, RbcUkqcd::c1m20LCDPar, RbcUkqcd::c1m20DeflPar, solverName,
                                     gaugeName, gaugeTransform, residual);
}

// Light C1M32
void RbcUkqcd::addC1M32LightRuntimeLCDSolver(Application &app, const std::string solverName,
                                           const std::string gaugeName,
                                           const std::string gaugeTransform, const double residual)
{
  RbcUkqcd::addLightRuntimeLCDSolver(app, RbcUkqcd::c1m32LCDPar, RbcUkqcd::c1m32DeflPar, solverName,
                                     gaugeName, gaugeTransform, residual);
}

// Strange 
void RbcUkqcd::addStrangeSolver(Application &app, const RbcUkqcd::EnsembleParameters &par,
                                const std::string solverName, const std::string gaugeName,
                                const std::string gaugeTransform, const double residual)
{
  const std::string prefix = solverName;

  // Gauge field FP32 cast
  MUtilities::GaugeSinglePrecisionCast::Par gaugeCastPar;

  gaugeCastPar.field = gaugeName;
  app.createModule<MUtilities::GaugeSinglePrecisionCast>(prefix + "_gauge_fp32", gaugeCastPar);

  // Scaled DWF action + FP32 version
  MAction::ScaledDWF::Par actionPar;

  actionPar.gauge = gaugeName;
  actionPar.Ls = par.Ls;
  actionPar.M5 = par.M5;
  actionPar.mass = par.ms;
  actionPar.scale = par.scale;
  actionPar.boundary = "1 1 1 1";
  actionPar.twist = "0. 0. 0. 0.";
  app.createModule<MAction::ScaledDWF>(prefix + "_dwf", actionPar);
  actionPar.gauge = prefix + "_gauge_fp32";
  app.createModule<MAction::ScaledDWFF>(prefix + "_dwf_fp32", actionPar);

  // Mixed-precision red-black preconditionned CG
  MSolver::MixedPrecisionRBPrecCG::Par solverPar;

  solverPar.innerAction = prefix + "_dwf_fp32";
  solverPar.outerAction = prefix + "_dwf";
  solverPar.maxInnerIteration = 30000;
  solverPar.maxOuterIteration = 100;
  solverPar.residual = residual;
  solverPar.innerGuesser = "";
  solverPar.outerGuesser = "";
  app.createModule<MSolver::MixedPrecisionRBPrecCG>(solverName, solverPar);
}

void RbcUkqcd::addM0StrangeSolver(Application &app, const std::string solverName,
                                  const std::string gaugeName, const std::string gaugeTransform,
                                  const double residual)
{
  RbcUkqcd::addStrangeSolver(app, RbcUkqcd::m0LCDPar, solverName, gaugeName, gaugeTransform,
                             residual);
}

void RbcUkqcd::addC1MStrangeSolver(Application &app, const std::string solverName,
                                   const std::string gaugeName, const std::string gaugeTransform,
                                   const double residual)
{
  RbcUkqcd::addStrangeSolver(app, RbcUkqcd::c1mLCDPar, solverName, gaugeName, gaugeTransform,
                             residual);
}

void RbcUkqcd::addC1M16StrangeSolver(Application &app, const std::string solverName,
                                     const std::string gaugeName, const std::string gaugeTransform,
                                     const double residual)
{
  RbcUkqcd::addStrangeSolver(app, RbcUkqcd::c1m16LCDPar, solverName, gaugeName, gaugeTransform,
                             residual);
}

void RbcUkqcd::addC1M20StrangeSolver(Application &app, const std::string solverName,
                                     const std::string gaugeName, const std::string gaugeTransform,
                                     const double residual)
{
  RbcUkqcd::addStrangeSolver(app, RbcUkqcd::c1m20LCDPar, solverName, gaugeName, gaugeTransform,
                             residual);
}

void RbcUkqcd::addC1M32StrangeSolver(Application &app, const std::string solverName,
                                     const std::string gaugeName, const std::string gaugeTransform,
                                     const double residual)
{
  RbcUkqcd::addStrangeSolver(app, RbcUkqcd::c1m32LCDPar, solverName, gaugeName, gaugeTransform,
                             residual);
}

} // namespace hadpresets