///////////////////////////////////////////////////////////////////////////////
//
// File: IMEXdirkTimeIntegrationSchemes.h
//
// For more information, please see: http://www.nektar.info
//
// The MIT License
//
// Copyright (c) 2018 Division of Applied Mathematics, Brown University (USA),
// Department of Aeronautics, Imperial College London (UK), and Scientific
// Computing and Imaging Institute, University of Utah (USA).
//
// License for the specific language governing rights and limitations under
// 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: Combined header file for all IMEX Dirk based time integration
// schemes.
//
///////////////////////////////////////////////////////////////////////////////

// Note : If adding a new integrator be sure to register the
// integrator with the Time Integration Scheme Facatory in
// SchemeInitializor.cpp.

#ifndef NEKTAR_LIB_UTILITIES_TIME_INTEGRATION_IMEX_DIRK_TIME_INTEGRATION_SCHEME
#define NEKTAR_LIB_UTILITIES_TIME_INTEGRATION_IMEX_DIRK_TIME_INTEGRATION_SCHEME

#define LUE LIB_UTILITIES_EXPORT

#include <LibUtilities/TimeIntegration/TimeIntegrationSchemeGLM.h>

namespace Nektar::LibUtilities
{

///////////////////////////////////////////////////////////////////////////////
//  IMEXDirk-(s, sigma, p), where s is the number of stages of the
//  implicit scheme, sigma is the number of stages of the explicit
//  scheme and p is the combined order of the scheme.

class IMEXdirkTimeIntegrationScheme : public TimeIntegrationSchemeGLM
{
public:
    IMEXdirkTimeIntegrationScheme(std::string variant, size_t order,
                                  std::vector<NekDouble> freeParams)
        : TimeIntegrationSchemeGLM(variant, order, freeParams)
    {
        ASSERTL1(freeParams.size() == 2,
                 "IMEX DIRK Time integration scheme invalid number "
                 "of free parameters, expected two "
                 "<implicit stages, explicit stages>, received " +
                     std::to_string(freeParams.size()));

        size_t s     = freeParams[0];
        size_t sigma = freeParams[1];

        m_integration_phases    = TimeIntegrationAlgorithmGLMVector(1);
        m_integration_phases[0] = TimeIntegrationAlgorithmGLMSharedPtr(
            new TimeIntegrationAlgorithmGLM(this));

        if (order == 1 && s == 1 && sigma == 1)
        {
            // This phase is Forward Backward Euler which has two steps.
            IMEXdirkTimeIntegrationScheme::SetupSchemeData_1_1_1(
                m_integration_phases[0]);
        }
        else
        {
            IMEXdirkTimeIntegrationScheme::SetupSchemeData(
                m_integration_phases[0], order, freeParams);
        }
    }

    ~IMEXdirkTimeIntegrationScheme() override = default;

    static TimeIntegrationSchemeSharedPtr create(
        std::string variant, size_t order, std::vector<NekDouble> freeParams)
    {
        TimeIntegrationSchemeSharedPtr p =
            MemoryManager<IMEXdirkTimeIntegrationScheme>::AllocateSharedPtr(
                variant, order, freeParams);

        return p;
    }

    static std::string className;

    LUE static void SetupSchemeData(TimeIntegrationAlgorithmGLMSharedPtr &phase,
                                    size_t order,
                                    std::vector<NekDouble> freeParams)
    {
        ASSERTL1(freeParams.size() == 2,
                 "IMEX DIRK Time integration scheme invalid number "
                 "of free parameters, expected two "
                 "<implicit stages, explicit stages>, received " +
                     std::to_string(freeParams.size()) + ".");

        size_t s     = freeParams[0];
        size_t sigma = freeParams[1];

        phase->m_schemeType = eIMEX;
        phase->m_variant    = "DIRK";
        phase->m_order      = order;
        phase->m_name       = "IMEX_DIRK"
                              "_" +
                        std::to_string(s) + "_" + std::to_string(sigma) + "_" +
                        std::to_string(phase->m_order);

        phase->m_numsteps  = 1;
        phase->m_numstages = s + 1;

        phase->m_A = Array<OneD, Array<TwoD, NekDouble>>(2);
        phase->m_B = Array<OneD, Array<TwoD, NekDouble>>(2);

        phase->m_A[0] =
            Array<TwoD, NekDouble>(phase->m_numstages, phase->m_numstages, 0.0);
        phase->m_B[0] =
            Array<TwoD, NekDouble>(phase->m_numsteps, phase->m_numstages, 0.0);
        phase->m_A[1] =
            Array<TwoD, NekDouble>(phase->m_numstages, phase->m_numstages, 0.0);
        phase->m_B[1] =
            Array<TwoD, NekDouble>(phase->m_numsteps, phase->m_numstages, 0.0);
        phase->m_U =
            Array<TwoD, NekDouble>(phase->m_numstages, phase->m_numsteps, 1.0);
        phase->m_V =
            Array<TwoD, NekDouble>(phase->m_numsteps, phase->m_numsteps, 1.0);

        if (s == 1 && sigma == 2 && order == 1)
        {
            SetupSchemeData_1_2_1(phase);
        }
        else if (s == 1 && sigma == 2 && order == 2)
        {
            SetupSchemeData_1_2_2(phase);
        }
        else if (s == 2 && sigma == 2 && order == 2)
        {
            SetupSchemeData_2_2_2(phase);
        }
        else if (s == 2 && sigma == 3 && order == 2)
        {
            SetupSchemeData_2_3_2(phase);
        }
        else if (s == 2 && sigma == 3 && order == 3)
        {
            SetupSchemeData_2_3_3(phase);
        }
        else if (s == 3 && sigma == 4 && order == 3)
        {
            SetupSchemeData_3_4_3(phase);
        }
        else if (s == 4 && sigma == 4 && order == 3)
        {
            SetupSchemeData_4_4_3(phase);
        }
        else
        {
            ASSERTL1(false,
                     "IMEX DIRK Time integration scheme bad type "
                     "(s, sigma, order) : (" +
                         std::to_string(s) + "," + std::to_string(sigma) + "," +
                         std::to_string(order) + "). " +
                         "Allowed combinations: (1,1,1), (1,2,1), (1,2,2), "
                         "(2,2,2), (2,3,2), (2,3,3), (3,4,3), (4,4,3)");
        }

        phase->m_numMultiStepValues         = 1;
        phase->m_numMultiStepImplicitDerivs = 0;
        phase->m_numMultiStepExplicitDerivs = 0;
        phase->m_timeLevelOffset    = Array<OneD, size_t>(phase->m_numsteps);
        phase->m_timeLevelOffset[0] = 0;

        phase->CheckAndVerify();
    }

    ///////////////////////////////////////////////////////////////////////////////
    // IMEX Dirk 1 1 1 : Forward - Backward Euler IMEX
    LUE static void SetupSchemeData_1_1_1(
        TimeIntegrationAlgorithmGLMSharedPtr &phase)
    {
        phase->m_schemeType = eIMEX;
        phase->m_order      = 1;
        phase->m_name =
            std::string("IMEXdirk_1_1_" + std::to_string(phase->m_order));

        phase->m_numsteps  = 2;
        phase->m_numstages = 1;

        phase->m_A = Array<OneD, Array<TwoD, NekDouble>>(2);
        phase->m_B = Array<OneD, Array<TwoD, NekDouble>>(2);

        phase->m_A[0] =
            Array<TwoD, NekDouble>(phase->m_numstages, phase->m_numstages, 0.0);
        phase->m_B[0] =
            Array<TwoD, NekDouble>(phase->m_numsteps, phase->m_numstages, 0.0);
        phase->m_A[1] =
            Array<TwoD, NekDouble>(phase->m_numstages, phase->m_numstages, 0.0);
        phase->m_B[1] =
            Array<TwoD, NekDouble>(phase->m_numsteps, phase->m_numstages, 0.0);
        phase->m_U =
            Array<TwoD, NekDouble>(phase->m_numstages, phase->m_numsteps, 0.0);
        phase->m_V =
            Array<TwoD, NekDouble>(phase->m_numsteps, phase->m_numsteps, 0.0);

        phase->m_A[0][0][0] = 1.0;

        phase->m_B[0][0][0] = 1.0;
        phase->m_B[1][1][0] = 1.0;

        phase->m_U[0][0] = 1.0;
        phase->m_U[0][1] = 1.0;

        phase->m_V[0][0] = 1.0;
        phase->m_V[0][1] = 1.0;

        phase->m_numMultiStepValues         = 1;
        phase->m_numMultiStepImplicitDerivs = 0;
        phase->m_numMultiStepExplicitDerivs = 1;

        phase->m_timeLevelOffset    = Array<OneD, size_t>(phase->m_numsteps);
        phase->m_timeLevelOffset[0] = 0;
        phase->m_timeLevelOffset[1] = 0;

        phase->CheckAndVerify();
    }

    ///////////////////////////////////////////////////////////////////////////////
    // IMEX Dirk 1 2 1 : Forward - Backward Euler IMEX w/B implicit == B
    // explicit
    LUE static void SetupSchemeData_1_2_1(
        TimeIntegrationAlgorithmGLMSharedPtr &phase)
    {
        phase->m_A[0][1][1] = 1.0;
        phase->m_B[0][0][1] = 1.0;

        phase->m_A[1][1][0] = 1.0;
        phase->m_B[1][0][1] = 1.0;

        // U and V set to 1 when allocated.
    }

    ////////////////////////////////////////////////////////////////////////////
    // IMEX Dirk 1 2 2 : Implict-Explicit Midpoint IMEX
    LUE static void SetupSchemeData_1_2_2(
        TimeIntegrationAlgorithmGLMSharedPtr &phase)
    {
        phase->m_A[0][1][1] = 1.0 / 2.0;
        phase->m_B[0][0][1] = 1.0;

        phase->m_A[1][1][0] = 1.0 / 2.0;
        phase->m_B[1][0][1] = 1.0;

        // U and V set to 1 when allocated.
    }

    ////////////////////////////////////////////////////////////////////////////
    // IMEX Dirk 2 2 2 : L Stable, two stage, second order IMEX
    LUE static void SetupSchemeData_2_2_2(
        TimeIntegrationAlgorithmGLMSharedPtr &phase)
    {
        const NekDouble glambda = 1.0 - sqrt(2.0) / 2.0;
        const NekDouble gdelta  = -sqrt(2.0) / 2.0;

        phase->m_A[0][1][1] = glambda;
        phase->m_A[0][2][1] = 1.0 - glambda;
        phase->m_A[0][2][2] = glambda;

        phase->m_B[0][0][1] = 1.0 - glambda;
        phase->m_B[0][0][2] = glambda;

        phase->m_A[1][1][0] = glambda;
        phase->m_A[1][2][0] = gdelta;
        phase->m_A[1][2][1] = 1.0 - gdelta;

        phase->m_B[1][0][0] = gdelta;
        phase->m_B[1][0][1] = 1.0 - gdelta;

        // U and V set to 1 when allocated.
    }

    ////////////////////////////////////////////////////////////////////////////
    // IMEX Dirk 2 3 2 : L Stable, two stage, second order IMEX
    LUE static void SetupSchemeData_2_3_2(
        TimeIntegrationAlgorithmGLMSharedPtr &phase)
    {
        const NekDouble lambda = (2.0 - sqrt(2.0)) / 2.0;
        const NekDouble delta  = -2.0 * sqrt(2.0) / 3.0;

        phase->m_A[0][1][1] = lambda;
        phase->m_A[0][2][1] = 1.0 - lambda;
        phase->m_A[0][2][2] = lambda;

        phase->m_B[0][0][1] = 1.0 - lambda;
        phase->m_B[0][0][2] = lambda;

        phase->m_A[1][1][0] = lambda;
        phase->m_A[1][2][0] = delta;
        phase->m_A[1][2][1] = 1.0 - delta;

        phase->m_B[1][0][1] = 1.0 - lambda;
        phase->m_B[1][0][2] = lambda;

        // U and V set to 1 when allocated.
    }

    ////////////////////////////////////////////////////////////////////////////
    // IMEX Dirk 2 3 3 : L Stable, two stage, third order IMEX
    LUE static void SetupSchemeData_2_3_3(
        TimeIntegrationAlgorithmGLMSharedPtr &phase)
    {
        const NekDouble glambda = (3.0 + sqrt(3.0)) / 6.0;

        phase->m_A[0][1][1] = glambda;
        phase->m_A[0][2][1] = 1.0 - 2.0 * glambda;
        phase->m_A[0][2][2] = glambda;

        phase->m_B[0][0][1] = 0.5;
        phase->m_B[0][0][2] = 0.5;

        phase->m_A[1][1][0] = glambda;
        phase->m_A[1][2][0] = glambda - 1.0;
        phase->m_A[1][2][1] = 2.0 * (1.0 - glambda);

        phase->m_B[1][0][1] = 0.5;
        phase->m_B[1][0][2] = 0.5;

        // U and V set to 1 when allocated.
    }

    ////////////////////////////////////////////////////////////////////////////
    // IMEX Dirk 3 4 3 : L Stable, three stage, third order IMEX
    LUE static void SetupSchemeData_3_4_3(
        TimeIntegrationAlgorithmGLMSharedPtr &phase)
    {
        constexpr NekDouble lambda = 0.4358665215;

        phase->m_A[0][1][1] = lambda;
        phase->m_A[0][2][1] = 0.5 * (1.0 - lambda);
        phase->m_A[0][2][2] = lambda;
        phase->m_A[0][3][1] =
            0.25 * (-6.0 * lambda * lambda + 16.0 * lambda - 1.0);
        phase->m_A[0][3][2] =
            0.25 * (6.0 * lambda * lambda - 20.0 * lambda + 5.0);
        phase->m_A[0][3][3] = lambda;

        phase->m_B[0][0][1] =
            0.25 * (-6.0 * lambda * lambda + 16.0 * lambda - 1.0);
        phase->m_B[0][0][2] =
            0.25 * (6.0 * lambda * lambda - 20.0 * lambda + 5.0);
        phase->m_B[0][0][3] = lambda;

        phase->m_A[1][1][0] = 0.4358665215;
        phase->m_A[1][2][0] = 0.3212788860;
        phase->m_A[1][2][1] = 0.3966543747;
        phase->m_A[1][3][0] = -0.105858296;
        phase->m_A[1][3][1] = 0.5529291479;
        phase->m_A[1][3][2] = 0.5529291479;

        phase->m_B[1][0][1] =
            0.25 * (-6.0 * lambda * lambda + 16.0 * lambda - 1.0);
        phase->m_B[1][0][2] =
            0.25 * (6.0 * lambda * lambda - 20.0 * lambda + 5.0);
        phase->m_B[1][0][3] = lambda;

        // U and V set to 1 when allocated.
    }

    ////////////////////////////////////////////////////////////////////////////
    // IMEX Dirk 4 4 3 : L Stable, four stage, third order IMEX
    LUE static void SetupSchemeData_4_4_3(
        TimeIntegrationAlgorithmGLMSharedPtr &phase)
    {
        phase->m_A[0][1][1] = 1.0 / 2.0;
        phase->m_A[0][2][1] = 1.0 / 6.0;
        phase->m_A[0][2][2] = 1.0 / 2.0;
        phase->m_A[0][3][1] = -1.0 / 2.0;
        phase->m_A[0][3][2] = 1.0 / 2.0;
        phase->m_A[0][3][3] = 1.0 / 2.0;
        phase->m_A[0][4][1] = 3.0 / 2.0;
        phase->m_A[0][4][2] = -3.0 / 2.0;
        phase->m_A[0][4][3] = 1.0 / 2.0;
        phase->m_A[0][4][4] = 1.0 / 2.0;

        phase->m_B[0][0][1] = 3.0 / 2.0;
        phase->m_B[0][0][2] = -3.0 / 2.0;
        phase->m_B[0][0][3] = 1.0 / 2.0;
        phase->m_B[0][0][4] = 1.0 / 2.0;

        phase->m_A[1][1][0] = 1.0 / 2.0;
        phase->m_A[1][2][0] = 11.0 / 18.0;
        phase->m_A[1][2][1] = 1.0 / 18.0;
        phase->m_A[1][3][0] = 5.0 / 6.0;
        phase->m_A[1][3][1] = -5.0 / 6.0;
        phase->m_A[1][3][2] = 1.0 / 2.0;
        phase->m_A[1][4][0] = 1.0 / 4.0;
        phase->m_A[1][4][1] = 7.0 / 4.0;
        phase->m_A[1][4][2] = 3.0 / 4.0;
        phase->m_A[1][4][3] = -7.0 / 4.0;

        phase->m_B[1][0][0] = 1.0 / 4.0;
        phase->m_B[1][0][1] = 7.0 / 4.0;
        phase->m_B[1][0][2] = 3.0 / 4.0;
        phase->m_B[1][0][3] = -7.0 / 4.0;

        // U and V set to 1 when allocated.
    }

protected:
    LUE std::string v_GetFullName() const override
    {
        return m_integration_phases.back()->m_name;
    }

    LUE std::string v_GetName() const override
    {
        return std::string("IMEX");
    }

    LUE NekDouble v_GetTimeStability() const override
    {
        return 1.0;
    }

}; // end class IMEXdirkTimeIntegrationScheme

} // namespace Nektar::LibUtilities

#endif
