/*
 * 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;
  };

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

  // Solvers (LCD: Local Coherence Deflation)
  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);
  static inline void addM0StrangeSolver(Application &app, const std::string solverName,
                                        const std::string gaugeName,
                                        const std::string gaugeTransform,
                                        const double residual = 1.0e-8);
};

// Implementations /////////////////////////////////////////////////////////////////////////////////
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);
}

void RbcUkqcd::addM0StrangeSolver(Application &app, 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 = RbcUkqcd::m0LCDPar.Ls;
  actionPar.M5 = RbcUkqcd::m0LCDPar.M5;
  actionPar.mass = RbcUkqcd::m0LCDPar.ms;
  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);

  // Batched 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);
}

} // namespace hadpresets