///////////////////////////////////////////////////////////////////////////////
//
// File: DriverArnoldi.cpp
//
// For more information, please see: http://www.nektar.info
//
// The MIT License
//
// Copyright (c) 2006 Division of Applied Mathematics, Brown University (USA),
// Department of Aeronautics, Imperial College London (UK), and Scientific
// Computing and Imaging Institute, University of Utah (USA).
//
// Permission is hereby granted, free of charge, to any person obtaining a
// copy of this software and associated documentation files (the "Software"),
// to deal in the Software without restriction, including without limitation
// the rights to use, copy, modify, merge, publish, distribute, sublicense,
// and/or sell copies of the Software, and to permit persons to whom the
// Software is furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included
// in all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
// OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
// THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
// FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
// DEALINGS IN THE SOFTWARE.
//
// Description: Base Driver class for the stability solver
//
///////////////////////////////////////////////////////////////////////////////

#include <iomanip>

#include <SolverUtils/DriverArnoldi.h>

namespace Nektar::SolverUtils
{

/**
 * Constructor
 */
DriverArnoldi::DriverArnoldi(
    const LibUtilities::SessionReaderSharedPtr pSession,
    const SpatialDomains::MeshGraphSharedPtr pGraph)
    : Driver(pSession, pGraph)
{
    m_session->LoadParameter("IO_InfoSteps", m_infosteps, 1);
}

/**
 * Arnoldi driver initialisation
 */
void DriverArnoldi::v_InitObject(std::ostream &out)
{
    Driver::v_InitObject(out);
    m_session->MatchSolverInfo("SolverType", "VelocityCorrectionScheme",
                               m_timeSteppingAlgorithm, false);

    if (m_timeSteppingAlgorithm)
    {
        m_period = m_session->GetParameter("TimeStep") *
                   m_session->GetParameter("NumSteps");
        m_nfields = m_equ[0]->UpdateFields().size() - 1;
    }
    else
    {
        m_period  = 1.0;
        m_nfields = m_equ[0]->UpdateFields().size();
    }

    if (m_session->DefinesSolverInfo("ModeType") &&
        (boost::iequals(m_session->GetSolverInfo("ModeType"), "SingleMode") ||
         boost::iequals(m_session->GetSolverInfo("ModeType"), "HalfMode")))
    {
        for (int i = 0; i < m_nfields; ++i)
        {
            m_equ[0]->UpdateFields()[i]->SetWaveSpace(true);
        }
    }
    m_negatedOp = m_equ[0]->NegatedOp();

    m_session->LoadParameter("kdim", m_kdim, 16);
    m_session->LoadParameter("nvec", m_nvec, 2);
    m_session->LoadParameter("nits", m_nits, 500);
    m_session->LoadParameter("evtol", m_evtol, 1e-06);

    ASSERTL0(m_kdim >= m_nvec, "nvec cannot be larger than kdim.");

    m_session->LoadParameter("realShift", m_realShift, 0.0);
    m_equ[0]->SetLambda(m_realShift);

    m_session->LoadParameter("imagShift", m_imagShift, 0.0);

    // The imaginary shift is applied at system assembly
    // Only if using HOMOGENEOUS expansion and ModeType set to SingleMode
    if (m_imagShift != 0.0)
    {
        if (!m_session->DefinesSolverInfo("HOMOGENEOUS") &&
            !m_session->DefinesSolverInfo("ModeType"))
        {
            NEKERROR(ErrorUtil::efatal,
                     "Imaginary shift only supported with HOMOGENEOUS "
                     "expansion and ModeType set to SingleMode");
        }
        else if (!boost::iequals(m_session->GetSolverInfo("ModeType"),
                                 "SingleMode"))
        {
            NEKERROR(ErrorUtil::efatal,
                     "Imaginary shift only supported with HOMOGENEOUS "
                     "expansion and ModeType set to SingleMode");
        }
    }
    MaskInit();
}

/**
 *
 */
void DriverArnoldi::v_Execute([[maybe_unused]] std::ostream &out)
{
    ASSERTL0(false, "Specific version of Arnoldi driver not implemented");
}

/**
 *
 */
void DriverArnoldi::ArnoldiSummary(std::ostream &out)
{
    if (m_comm->GetRank() == 0)
    {
        if (m_session->DefinesSolverInfo("ModeType") &&
            boost::iequals(m_session->GetSolverInfo("ModeType"), "SingleMode"))
        {
            out << "\tSingle Fourier mode    : true " << std::endl;
            ASSERTL0(m_session->DefinesSolverInfo("Homogeneous"),
                     "Expected a homogeneous expansion to be defined "
                     "with single mode");
        }
        else
        {
            out << "\tSingle Fourier mode    : false " << std::endl;
        }
        if (m_session->DefinesSolverInfo("BetaZero"))
        {
            out << "\tBeta set to Zero       : true (overrides LHom)"
                << std::endl;
        }
        else
        {
            out << "\tBeta set to Zero       : false " << std::endl;
        }

        if (m_timeSteppingAlgorithm)
        {
            out << "\tEvolution operator     : "
                << m_session->GetSolverInfo("EvolutionOperator") << std::endl;
        }
        else
        {
            out << "\tShift (Real,Imag)      : " << m_realShift << ","
                << m_imagShift << std::endl;
        }
        out << "\tKrylov-space dimension : " << m_kdim << std::endl;
        out << "\tNumber of vectors      : " << m_nvec << std::endl;
        out << "\tMax iterations         : " << m_nits << std::endl;
        out << "\tEigenvalue tolerance   : " << m_evtol << std::endl;
        out << "======================================================"
            << std::endl;
    }
}

/**
 * Copy Arnoldi array to field variables which depend from
 * either the m_fields or m_forces
 */
void DriverArnoldi::CopyArnoldiArrayToField(Array<OneD, NekDouble> &array)
{

    Array<OneD, MultiRegions::ExpListSharedPtr> &fields =
        m_equ[0]->UpdateFields();
    int nq = fields[0]->GetNcoeffs();

    for (int k = 0; k < m_nfields; ++k)
    {
        Vmath::Vcopy(nq, &array[k * nq], 1, &fields[k]->UpdateCoeffs()[0], 1);
        fields[k]->SetPhysState(false);
    }
}

/**
 * Copy field variables which depend from either the m_fields
 * or m_forces array the Arnoldi array
 */
void DriverArnoldi::CopyFieldToArnoldiArray(Array<OneD, NekDouble> &array)
{

    Array<OneD, MultiRegions::ExpListSharedPtr> fields;

    if (m_EvolutionOperator == eAdaptiveSFD)
    {
        // This matters for the Adaptive SFD method because
        // m_equ[1] is the nonlinear problem with non
        // homogeneous BCs.
        fields = m_equ[0]->UpdateFields();
    }
    else
    {
        fields = m_equ[m_nequ - 1]->UpdateFields();
    }

    for (int k = 0; k < m_nfields; ++k)
    {
        int nq = fields[0]->GetNcoeffs();
        Vmath::Vcopy(nq, &fields[k]->GetCoeffs()[0], 1, &array[k * nq], 1);
        fields[k]->SetPhysState(false);
    }
}

/**
 * Initialisation for the transient growth
 */
void DriverArnoldi::CopyFwdToAdj()
{
    Array<OneD, MultiRegions::ExpListSharedPtr> fields;

    ASSERTL0(m_timeSteppingAlgorithm,
             "Transient Growth non available for Coupled Solver");

    fields = m_equ[0]->UpdateFields();
    int nq = fields[0]->GetNcoeffs();

    for (int k = 0; k < m_nfields; ++k)
    {
        Vmath::Vcopy(nq, &fields[k]->GetCoeffs()[0], 1,
                     &m_equ[1]->UpdateFields()[k]->UpdateCoeffs()[0], 1);
    }
}

/**
 *
 */
void DriverArnoldi::WriteFld(std::string file,
                             std::vector<Array<OneD, NekDouble>> coeffs)
{

    std::vector<std::string> variables(m_nfields);

    ASSERTL1(coeffs.size() >= m_nfields, "coeffs is not of the correct length");
    for (int i = 0; i < m_nfields; ++i)
    {
        variables[i] = m_equ[0]->GetVariable(i);
    }

    m_equ[0]->WriteFld(file, m_equ[0]->UpdateFields()[0], coeffs, variables);
}

/**
 *
 */
void DriverArnoldi::WriteFld(std::string file, Array<OneD, NekDouble> coeffs)
{

    std::vector<std::string> variables(m_nfields);
    std::vector<Array<OneD, NekDouble>> fieldcoeffs(m_nfields);

    int ncoeffs = m_equ[0]->UpdateFields()[0]->GetNcoeffs();
    ASSERTL1(coeffs.size() >= ncoeffs * m_nfields,
             "coeffs is not of sufficient size");

    for (int i = 0; i < m_nfields; ++i)
    {
        variables[i]   = m_equ[0]->GetVariable(i);
        fieldcoeffs[i] = coeffs + i * ncoeffs;
    }

    m_equ[0]->WriteFld(file, m_equ[0]->UpdateFields()[0], fieldcoeffs,
                       variables);
}

/**
 *
 */
void DriverArnoldi::WriteEvs(std::ostream &evlout, const int i,
                             const NekDouble re_ev, const NekDouble im_ev,
                             NekDouble resid, bool DumpInverse)
{
    if (m_timeSteppingAlgorithm)
    {
        NekDouble abs_ev = hypot(re_ev, im_ev);
        NekDouble ang_ev = atan2(im_ev, re_ev);

        evlout << "EV: " << std::setw(2) << i << std::setw(12) << abs_ev
               << std::setw(12) << ang_ev << std::setw(12)
               << log(abs_ev) / m_period << std::setw(12) << ang_ev / m_period;

        if (resid != NekConstants::kNekUnsetDouble)
        {
            evlout << std::setw(12) << resid;
        }
        evlout << std::endl;
    }
    else
    {
        NekDouble invmag = 1.0 / (re_ev * re_ev + im_ev * im_ev);
        NekDouble sign;
        if (m_negatedOp)
        {
            sign = -1.0;
        }
        else
        {
            sign = 1.0;
        }

        evlout << "EV: " << std::setw(2) << i << std::setw(14) << sign * re_ev
               << std::setw(14) << sign * im_ev;

        if (DumpInverse)
        {
            evlout << std::setw(14) << sign * re_ev * invmag + m_realShift
                   << std::setw(14) << sign * im_ev * invmag + m_imagShift;
        }

        if (resid != NekConstants::kNekUnsetDouble)
        {
            evlout << std::setw(12) << resid;
        }
        evlout << std::endl;
    }
}

/**
 *
 */
void DriverArnoldi::GetMaskInfo(
    std::vector<std::vector<LibUtilities::EquationSharedPtr>> &selectedDomains,
    std::set<int> &unselectedVariables)
{
    selectedDomains.clear();
    std::string domain("SelectEVCalcDomain0");
    std::string condition("C0");
    for (size_t i = 0; i < 10; ++i)
    {
        domain[domain.size() - 1] = '0' + i;
        if (!m_session->DefinesFunction(domain))
        {
            break;
        }
        for (size_t j = 0; j < 10; ++j)
        {
            condition[condition.size() - 1] = '0' + j;
            if (!m_session->DefinesFunction(domain, condition))
            {
                break;
            }
            if (j == 0)
            {
                selectedDomains.push_back(
                    std::vector<LibUtilities::EquationSharedPtr>());
            }
            selectedDomains[selectedDomains.size() - 1].push_back(
                m_session->GetFunction(domain, condition));
        }
    }
    unselectedVariables.clear();
    std::string funName("SelectEVCalcVariables");
    std::vector<std::string> variables = m_session->GetVariables();
    for (size_t v = 0; v < m_nfields; ++v)
    {
        if (!m_session->DefinesFunction(funName, variables[v]))
        {
            unselectedVariables.insert(v);
        }
    }
    if (unselectedVariables.size() == m_nfields)
    {
        unselectedVariables.clear();
    }
}

/**
 *
 */
void DriverArnoldi::MaskInit()
{
    std::vector<std::vector<LibUtilities::EquationSharedPtr>> selectedDomains;
    std::set<int> unselectedVariables;
    GetMaskInfo(selectedDomains, unselectedVariables);
    if (selectedDomains.size() == 0 && unselectedVariables.size() == 0)
    {
        m_useMask = false;
        return;
    }
    m_useMask                            = true;
    MultiRegions::ExpListSharedPtr field = m_equ[0]->UpdateFields()[0];
    int ncoef                            = field->GetNcoeffs();
    int nphys                            = field->GetNpoints();
    m_maskCoeffs = Array<OneD, double>(ncoef * m_nfields, 1.);
    m_maskPhys   = Array<OneD, double>(nphys * m_nfields, 1.);
    for (size_t i = 0; i < field->GetExpSize(); ++i)
    {
        LocalRegions::ExpansionSharedPtr exp   = field->GetExp(i);
        SpatialDomains::GeometrySharedPtr geom = exp->GetGeom();
        int nv                                 = geom->GetNumVerts();
        NekDouble gc[3]                        = {0., 0., 0.};
        NekDouble gct[3]                       = {0., 0., 0.};
        for (size_t j = 0; j < nv; ++j)
        {
            SpatialDomains::PointGeomSharedPtr vertex = geom->GetVertex(j);
            vertex->GetCoords(gct[0], gct[1], gct[2]);
            gc[0] += gct[0] / NekDouble(nv);
            gc[1] += gct[1] / NekDouble(nv);
            gc[2] += gct[2] / NekDouble(nv);
        }
        int unmask = 1;
        for (size_t m = 0; m < selectedDomains.size(); ++m)
        {
            unmask = 0;
            for (size_t n = 0; n < selectedDomains[m].size(); ++n)
            {
                if (selectedDomains[m][n]->Evaluate(gc[0], gc[1], gc[2]) <= 0.)
                {
                    unmask = 0;
                    break;
                }
                else
                {
                    unmask = 1;
                }
            }
            if (unmask == 1)
            {
                break;
            }
        }
        for (int j = 0; j < m_nfields; ++j)
        {
            if (unmask == 0 || unselectedVariables.count(j))
            {
                Vmath::Fill(
                    exp->GetNcoeffs(), 0.,
                    &m_maskCoeffs[field->GetCoeff_Offset(i) + j * ncoef], 1);
                Vmath::Fill(exp->GetTotPoints(), 0.,
                            &m_maskPhys[field->GetPhys_Offset(i) + j * nphys],
                            1);
            }
        }
    }
}

} // namespace Nektar::SolverUtils
