/*
 * Copyright (C) 2026
 *
 * Author: Antonin Portelli <antonin.portelli@me.com>
 * Author: Ryan Hill <Ryan.Hill@ed.ac.uk>
 * Elements based on workflows from Teseo San Jose Perez
 *
 * 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>

#include <array>
#include <stdexcept>
#include <string>
#include <string_view>

namespace hadpresets
{
using namespace Grid::Hadrons;

struct Jlqcd
{
  struct EnsembleParameters
  {
    double ml, ms, M5, scale;
    unsigned int L, T, Ls;
  };

  inline static constexpr EnsembleParameters fUd3SaPar {0.003,  0.015, 1., 2., 64, 128, 8};
  inline static constexpr EnsembleParameters mUd3SaPar {0.0042, 0.025, 1., 2., 48, 96,  8};
  inline static constexpr EnsembleParameters mUd3SbPar {0.0042, 0.018, 1., 2., 48, 96,  8};
  inline static constexpr EnsembleParameters cUd2SaPar {0.0035, 0.040, 1., 2., 32, 64,  12};
  inline static constexpr EnsembleParameters cUd2SaLPar{0.0035, 0.040, 1., 2., 48, 96,  12};
  inline static constexpr EnsembleParameters cUd3SaPar {0.007,  0.040, 1., 2., 32, 64,  12};
  inline static constexpr EnsembleParameters cUd3SbPar {0.007,  0.030, 1., 2., 32, 64,  12};
  inline static constexpr EnsembleParameters testPar   {0.010,  0.030, 1., 2.,  8, 16,   8};

  static inline const EnsembleParameters &ensemble(const std::string& name);

  static inline void addMixedPrecisionSolver(Application &app,
                                             const EnsembleParameters &par,
                                             const std::string &solverName,
                                             const std::string &gaugeName,
                                             const double mass,
                                             const std::string &twist,
                                             const std::string &boundary = "1 1 1 -1",
                                             const double residual = 1.0e-8);

  static inline void addNoFailSolver(Application &app, const EnsembleParameters &par,
                                     const std::string &solverName,
                                     const std::string &gaugeName, const double mass,
                                     const std::string &twist,
                                     const std::string &boundary = "1 1 1 -1",
                                     const double residual = 1.0e-20);

  // Light solvers (no deflation)
  static inline void addLightSolver(Application &app, const EnsembleParameters &par,
                                    const std::string &solverName, const std::string &gaugeName,
                                    const double mass, const std::string &twist,
                                    const std::string &boundary = "1 1 1 -1",
                                    const double residual = 1.0e-8);

  static inline void addLightSolver(Application &app, const EnsembleParameters &par,
                                    const std::string &solverName, const std::string &gaugeName,
                                    const std::string &twist = "0. 0. 0. 0.",
                                    const std::string &boundary = "1 1 1 -1",
                                    const double residual = 1.0e-8);

  static inline void addfUd3SaLightSolver(Application &app, const std::string &solverName,
                                          const std::string &gaugeName, const std::string &twist = "0. 0. 0. 0.",
                                          const std::string &boundary = "1 1 1 -1", const double residual = 1.0e-8);
  static inline void addmUd3SaLightSolver(Application &app, const std::string &solverName,
                                          const std::string &gaugeName, const std::string &twist = "0. 0. 0. 0.",
                                          const std::string &boundary = "1 1 1 -1", const double residual = 1.0e-8);
  static inline void addmUd3SbLightSolver(Application &app, const std::string &solverName,
                                          const std::string &gaugeName, const std::string &twist = "0. 0. 0. 0.",
                                          const std::string &boundary = "1 1 1 -1", const double residual = 1.0e-8);
  static inline void addcUd2SaLightSolver(Application &app, const std::string &solverName,
                                          const std::string &gaugeName, const std::string &twist = "0. 0. 0. 0.",
                                          const std::string &boundary = "1 1 1 -1", const double residual = 1.0e-8);
  static inline void addcUd2SaLLightSolver(Application &app, const std::string &solverName,
                                          const std::string &gaugeName, const std::string &twist = "0. 0. 0. 0.",
                                          const std::string &boundary = "1 1 1 -1", const double residual = 1.0e-8);
  static inline void addcUd3SaLightSolver(Application &app, const std::string &solverName,
                                          const std::string &gaugeName, const std::string &twist = "0. 0. 0. 0.",
                                          const std::string &boundary = "1 1 1 -1", const double residual = 1.0e-8);
  static inline void addcUd3SbLightSolver(Application &app, const std::string &solverName,
                                          const std::string &gaugeName, const std::string &twist = "0. 0. 0. 0.",
                                          const std::string &boundary = "1 1 1 -1", const double residual = 1.0e-8);

  // Strange solvers
  static inline void addStrangeSolver(Application &app, const EnsembleParameters &par,
                                      const std::string &solverName,
                                      const std::string &gaugeName, const double mass,
                                      const std::string &twist,
                                      const std::string &boundary = "1 1 1 -1",
                                      const double residual = 1.0e-8);

  static inline void addStrangeSolver(Application &app, const EnsembleParameters &par,
                                      const std::string &solverName,
                                      const std::string &gaugeName,
                                      const std::string &twist = "0. 0. 0. 0.",
                                      const std::string &boundary = "1 1 1 -1",
                                      const double residual = 1.0e-8);

  static inline void addfUd3SaStrangeSolver(Application &app, const std::string &solverName,
                                            const std::string &gaugeName, const std::string &twist = "0. 0. 0. 0.",
                                            const std::string &boundary = "1 1 1 -1", const double residual = 1.0e-8);
  static inline void addmUd3SaStrangeSolver(Application &app, const std::string &solverName,
                                            const std::string &gaugeName, const std::string &twist = "0. 0. 0. 0.",
                                            const std::string &boundary = "1 1 1 -1", const double residual = 1.0e-8);
  static inline void addmUd3SbStrangeSolver(Application &app, const std::string &solverName,
                                            const std::string &gaugeName, const std::string &twist = "0. 0. 0. 0.",
                                            const std::string &boundary = "1 1 1 -1", const double residual = 1.0e-8);
  static inline void addcUd2SaStrangeSolver(Application &app, const std::string &solverName,
                                            const std::string &gaugeName, const std::string &twist = "0. 0. 0. 0.",
                                            const std::string &boundary = "1 1 1 -1", const double residual = 1.0e-8);
  static inline void addcUd2SaLStrangeSolver(Application &app, const std::string &solverName,
                                            const std::string &gaugeName, const std::string &twist = "0. 0. 0. 0.",
                                            const std::string &boundary = "1 1 1 -1", const double residual = 1.0e-8);
  static inline void addcUd3SaStrangeSolver(Application &app, const std::string &solverName,
                                            const std::string &gaugeName, const std::string &twist = "0. 0. 0. 0.",
                                            const std::string &boundary = "1 1 1 -1", const double residual = 1.0e-8);
  static inline void addcUd3SbStrangeSolver(Application &app, const std::string &solverName,
                                            const std::string &gaugeName, const std::string &twist = "0. 0. 0. 0.",
                                            const std::string &boundary = "1 1 1 -1", const double residual = 1.0e-8);

  // Heavy solvers
  static inline void addHeavySolver(Application &app, const EnsembleParameters &par,
                                    const std::string &solverName, const std::string &gaugeName,
                                    const double mass, const std::string &twist,
                                    const std::string &boundary = "1 1 1 -1",
                                    const double residual = 1.0e-20);

  static inline void addfUd3SaHeavySolver(Application &app, const std::string &solverName, const std::string &gaugeName, 
                                          const double mass, const std::string &twist = "0. 0. 0. 0.",
                                          const std::string &boundary = "1 1 1 -1", const double residual = 1.0e-20);
  static inline void addmUd3SaHeavySolver(Application &app, const std::string &solverName, const std::string &gaugeName, 
                                          const double mass, const std::string &twist = "0. 0. 0. 0.",
                                          const std::string &boundary = "1 1 1 -1", const double residual = 1.0e-20);
  static inline void addmUd3SbHeavySolver(Application &app, const std::string &solverName, const std::string &gaugeName, 
                                          const double mass, const std::string &twist = "0. 0. 0. 0.",
                                          const std::string &boundary = "1 1 1 -1", const double residual = 1.0e-20);
  static inline void addcUd2SaHeavySolver(Application &app, const std::string &solverName, const std::string &gaugeName, 
                                          const double mass, const std::string &twist = "0. 0. 0. 0.",
                                          const std::string &boundary = "1 1 1 -1", const double residual = 1.0e-20);
  static inline void addcUd2SaLHeavySolver(Application &app, const std::string &solverName, const std::string &gaugeName, 
                                          const double mass, const std::string &twist = "0. 0. 0. 0.",
                                          const std::string &boundary = "1 1 1 -1", const double residual = 1.0e-20);
  static inline void addcUd3SaHeavySolver(Application &app, const std::string &solverName, const std::string &gaugeName, 
                                          const double mass, const std::string &twist = "0. 0. 0. 0.",
                                          const std::string &boundary = "1 1 1 -1", const double residual = 1.0e-20);
  static inline void addcUd3SbHeavySolver(Application &app, const std::string &solverName, const std::string &gaugeName, 
                                          const double mass, const std::string &twist = "0. 0. 0. 0.",
                                          const std::string &boundary = "1 1 1 -1", const double residual = 1.0e-20);

};

inline const Jlqcd::EnsembleParameters &Jlqcd::ensemble(const std::string& name)
{
  if (name == "fUd3Sa" ) return fUd3SaPar;
  if (name == "mUd3Sa" ) return mUd3SaPar;
  if (name == "mUd3Sb" ) return mUd3SbPar;
  if (name == "cUd2Sa" ) return cUd2SaPar;
  if (name == "cUd2SaL") return cUd2SaLPar;
  if (name == "cUd3Sa" ) return cUd3SaPar;
  if (name == "cUd3Sb" ) return cUd3SbPar;
  if (name == "test"   ) return testPar;

  throw std::invalid_argument("unknown JLQCD ensemble preset '" + std::string(name) + "'");
}

inline void Jlqcd::addMixedPrecisionSolver(Application &app, const EnsembleParameters &par,
                                           const std::string &solverName, const std::string &gaugeName,
                                           const double mass, const std::string &twist,
                                           const std::string &boundary, const double residual)
{
  const std::string prefix = solverName;

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

  MAction::ScaledDWF::Par actionPar;
  actionPar.gauge = gaugeName;
  actionPar.Ls = par.Ls;
  actionPar.M5 = par.M5;
  actionPar.mass = mass;
  actionPar.scale = par.scale;
  actionPar.boundary = boundary;
  actionPar.twist = twist;
  app.createModule<MAction::ScaledDWF>(prefix + "_dwf", actionPar);

  actionPar.gauge = prefix + "_gauge_fp32";
  app.createModule<MAction::ScaledDWFF>(prefix + "_dwf_fp32", actionPar);

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

inline void Jlqcd::addNoFailSolver(Application &app, const EnsembleParameters &par,
                                   const std::string &solverName, const std::string &gaugeName,
                                   const double mass, const std::string &twist,
                                   const std::string &boundary, const double residual)
{
  const std::string prefix = solverName;

  MAction::ScaledDWF::Par actionPar;
  actionPar.gauge = gaugeName;
  actionPar.Ls = par.Ls;
  actionPar.M5 = par.M5;
  actionPar.mass = mass;
  actionPar.scale = par.scale;
  actionPar.boundary = boundary;
  actionPar.twist = twist;
  app.createModule<MAction::ScaledDWF>(prefix + "_dwf", actionPar);

  MSolver::RBPrecCGNoFail::Par solverPar;
  solverPar.action = prefix + "_dwf";
  solverPar.maxIteration = 30000;
  solverPar.residual = residual;
  solverPar.guesser = "";
  app.createModule<MSolver::RBPrecCGNoFail>(solverName, solverPar);
}

// Light solvers (no deflation)
inline void Jlqcd::addLightSolver(Application &app, const EnsembleParameters &par,
                                  const std::string &solverName, const std::string &gaugeName,
                                  const double mass, const std::string &twist,
                                  const std::string &boundary, const double residual)
{
  addMixedPrecisionSolver(app, par, solverName, gaugeName, mass, twist, boundary, residual);
}

inline void Jlqcd::addLightSolver(Application &app, const EnsembleParameters &par,
                                  const std::string &solverName, const std::string &gaugeName,
                                  const std::string &twist, const std::string &boundary,
                                  const double residual)
{
  addLightSolver(app, par, solverName, gaugeName, par.ml, twist, boundary, residual);
}

inline void Jlqcd::addfUd3SaLightSolver(Application &app, const std::string &solverName,
                                        const std::string &gaugeName, const std::string &twist,
                                        const std::string &boundary, const double residual)
{
    addLightSolver(app, fUd3SaPar, solverName, gaugeName, twist, boundary, residual);
}

inline void Jlqcd::addmUd3SaLightSolver(Application &app, const std::string &solverName,
                                        const std::string &gaugeName, const std::string &twist,
                                        const std::string &boundary, const double residual)
{
    addLightSolver(app, mUd3SaPar, solverName, gaugeName, twist, boundary, residual);
}

inline void Jlqcd::addmUd3SbLightSolver(Application &app, const std::string &solverName,
                                        const std::string &gaugeName, const std::string &twist,
                                        const std::string &boundary, const double residual)
{
    addLightSolver(app, mUd3SbPar, solverName, gaugeName, twist, boundary, residual);
}

inline void Jlqcd::addcUd2SaLightSolver(Application &app, const std::string &solverName,
                                        const std::string &gaugeName, const std::string &twist,
                                        const std::string &boundary, const double residual)
{
    addLightSolver(app, cUd2SaPar, solverName, gaugeName, twist, boundary, residual);
}

inline void Jlqcd::addcUd2SaLLightSolver(Application &app, const std::string &solverName,
                                        const std::string &gaugeName, const std::string &twist,
                                        const std::string &boundary, const double residual)
{
    addLightSolver(app, cUd2SaLPar, solverName, gaugeName, twist, boundary, residual);
}

inline void Jlqcd::addcUd3SaLightSolver(Application &app, const std::string &solverName,
                                        const std::string &gaugeName, const std::string &twist,
                                        const std::string &boundary, const double residual)
{
    addLightSolver(app, cUd3SaPar, solverName, gaugeName, twist, boundary, residual);
}

inline void Jlqcd::addcUd3SbLightSolver(Application &app, const std::string &solverName,
                                        const std::string &gaugeName, const std::string &twist,
                                        const std::string &boundary, const double residual)
{
    addLightSolver(app, cUd3SbPar, solverName, gaugeName, twist, boundary, residual);
}

// Strange solvers
inline void Jlqcd::addStrangeSolver(Application &app, const EnsembleParameters &par,
                                    const std::string &solverName, const std::string &gaugeName,
                                    const double mass, const std::string &twist,
                                    const std::string &boundary, const double residual)
{
    addMixedPrecisionSolver(app, par, solverName, gaugeName, mass, twist, boundary, residual);
}

inline void Jlqcd::addStrangeSolver(Application &app, const EnsembleParameters &par,
                                    const std::string &solverName, const std::string &gaugeName,
                                    const std::string &twist, const std::string &boundary,
                                    const double residual)
{
    addStrangeSolver(app, par, solverName, gaugeName, par.ms, twist, boundary, residual);
}

inline void Jlqcd::addfUd3SaStrangeSolver(Application &app, const std::string &solverName,
                                          const std::string &gaugeName, const std::string &twist,
                                          const std::string &boundary, const double residual)
{
    addStrangeSolver(app, fUd3SaPar, solverName, gaugeName, twist, boundary, residual);
}

inline void Jlqcd::addmUd3SaStrangeSolver(Application &app, const std::string &solverName,
                                          const std::string &gaugeName, const std::string &twist,
                                          const std::string &boundary, const double residual)
{
    addStrangeSolver(app, mUd3SaPar, solverName, gaugeName, twist, boundary, residual);
}

inline void Jlqcd::addmUd3SbStrangeSolver(Application &app, const std::string &solverName,
                                          const std::string &gaugeName, const std::string &twist,
                                          const std::string &boundary, const double residual)
{
    addStrangeSolver(app, mUd3SbPar, solverName, gaugeName, twist, boundary, residual);
}

inline void Jlqcd::addcUd2SaStrangeSolver(Application &app, const std::string &solverName,
                                          const std::string &gaugeName, const std::string &twist,
                                          const std::string &boundary, const double residual)
{
    addStrangeSolver(app, cUd2SaPar, solverName, gaugeName, twist, boundary, residual);
}

inline void Jlqcd::addcUd2SaLStrangeSolver(Application &app, const std::string &solverName,
                                           const std::string &gaugeName, const std::string &twist,
                                           const std::string &boundary, const double residual)
{
    addStrangeSolver(app, cUd2SaLPar, solverName, gaugeName, twist, boundary, residual);
}

inline void Jlqcd::addcUd3SaStrangeSolver(Application &app, const std::string &solverName,
                                          const std::string &gaugeName, const std::string &twist,
                                          const std::string &boundary, const double residual)
{
    addStrangeSolver(app, cUd3SaPar, solverName, gaugeName, twist, boundary, residual);
}

inline void Jlqcd::addcUd3SbStrangeSolver(Application &app, const std::string &solverName,
                                          const std::string &gaugeName, const std::string &twist,
                                          const std::string &boundary, const double residual)
{
    addStrangeSolver(app, cUd3SbPar, solverName, gaugeName, twist, boundary, residual);
}

// Heavy solvers
inline void Jlqcd::addHeavySolver(Application &app, const EnsembleParameters &par,
                                  const std::string &solverName, const std::string &gaugeName,
                                  const double mass, const std::string &twist,
                                  const std::string &boundary, const double residual)
{
  addNoFailSolver(app, par, solverName, gaugeName, mass, twist, boundary, residual);
}

inline void Jlqcd::addfUd3SaHeavySolver(Application &app, const std::string &solverName, const std::string &gaugeName, 
                                        const double mass, const std::string &twist,
                                        const std::string &boundary, const double residual)
{
  addHeavySolver(app, fUd3SaPar, solverName, gaugeName, mass, twist, boundary, residual);
}
inline void Jlqcd::addmUd3SaHeavySolver(Application &app, const std::string &solverName, const std::string &gaugeName, 
                                        const double mass, const std::string &twist,
                                        const std::string &boundary, const double residual)
{
  addHeavySolver(app, mUd3SaPar, solverName, gaugeName, mass, twist, boundary, residual);
}
inline void Jlqcd::addmUd3SbHeavySolver(Application &app, const std::string &solverName, const std::string &gaugeName, 
                                        const double mass, const std::string &twist,
                                        const std::string &boundary, const double residual)
{
  addHeavySolver(app, mUd3SbPar, solverName, gaugeName, mass, twist, boundary, residual);
}
inline void Jlqcd::addcUd2SaHeavySolver(Application &app, const std::string &solverName, const std::string &gaugeName, 
                                        const double mass, const std::string &twist,
                                        const std::string &boundary, const double residual)
{
  addHeavySolver(app, cUd2SaPar, solverName, gaugeName, mass, twist, boundary, residual);
}
inline void Jlqcd::addcUd2SaLHeavySolver(Application &app, const std::string &solverName, const std::string &gaugeName, 
                                         const double mass, const std::string &twist,
                                         const std::string &boundary, const double residual)
{
  addHeavySolver(app, cUd2SaLPar, solverName, gaugeName, mass, twist, boundary, residual);
}
inline void Jlqcd::addcUd3SaHeavySolver(Application &app, const std::string &solverName, const std::string &gaugeName, 
                                        const double mass, const std::string &twist,
                                        const std::string &boundary, const double residual)
{
  addHeavySolver(app, cUd3SaPar, solverName, gaugeName, mass, twist, boundary, residual);
}
inline void Jlqcd::addcUd3SbHeavySolver(Application &app, const std::string &solverName, const std::string &gaugeName, 
                                        const double mass, const std::string &twist,
                                        const std::string &boundary, const double residual)
{
  addHeavySolver(app, cUd3SbPar, solverName, gaugeName, mass, twist, boundary, residual);
}

} // namespace hadpresets