// ---------------------------------------------------------------------
//
// Copyright (c) 2014 - 2020 by the IBAMR developers
// All rights reserved.
//
// This file is part of IBAMR.
//
// IBAMR is free software and is distributed under the 3-clause BSD
// license. The full text of the license can be found in the file
// COPYRIGHT at the top level directory of IBAMR.
//
// ---------------------------------------------------------------------

/////////////////////////////// INCLUDES /////////////////////////////////////

#include "ibtk/CartGridFunction.h"
#include "ibtk/ibtk_utilities.h"
#include "ibtk/muParserCartGridFunction.h"

#include "Box.h"
#include "CartesianGridGeometry.h"
#include "CartesianPatchGeometry.h"
#include "CellData.h"
#include "CellIndex.h"
#include "CellIterator.h"
#include "EdgeData.h"
#include "EdgeIndex.h"
#include "EdgeIterator.h"
#include "FaceData.h"
#include "FaceIndex.h"
#include "FaceIterator.h"
#include "Index.h"
#include "IntVector.h"
#include "NodeData.h"
#include "NodeIndex.h"
#include "NodeIterator.h"
#include "Patch.h"
#include "PatchData.h"
#include "SideData.h"
#include "SideIndex.h"
#include "SideIterator.h"
#include "tbox/Array.h"
#include "tbox/Database.h"
#include "tbox/Pointer.h"
#include "tbox/Utilities.h"

#include "muParser.h"
#include "muParserError.h"

#include <algorithm>
#include <map>
#include <memory>
#include <ostream>
#include <string>
#include <utility>
#include <vector>

#include "ibtk/namespaces.h" // IWYU pragma: keep

namespace SAMRAI
{
namespace hier
{
template <int DIM>
class PatchLevel;
} // namespace hier
} // namespace SAMRAI

namespace SAMRAI
{
namespace hier
{
template <int DIM>
class Variable;
} // namespace hier
} // namespace SAMRAI

/////////////////////////////// NAMESPACE ////////////////////////////////////

namespace IBTK
{
/////////////////////////////// STATIC ///////////////////////////////////////

/////////////////////////////// PUBLIC ///////////////////////////////////////

muParserCartGridFunction::muParserCartGridFunction(std::string object_name,
                                                   Pointer<Database> input_db,
                                                   Pointer<CartesianGridGeometry<NDIM> > grid_geom)
    : CartGridFunction(std::move(object_name)), d_grid_geom(grid_geom)
{
#if !defined(NDEBUG)
    TBOX_ASSERT(!d_object_name.empty());
    TBOX_ASSERT(input_db);
#endif
    // Read in user-provided constants.
    Array<std::string> db_key_names = input_db->getAllKeys();
    for (int k = 0; k < db_key_names.size(); ++k)
    {
        const std::string& name = db_key_names[k];
        if (input_db->isDouble(name))
        {
            d_constants[name] = input_db->getDouble(name);
        }
        else if (input_db->isFloat(name))
        {
            d_constants[name] = input_db->getFloat(name);
        }
        else if (input_db->isInteger(name))
        {
            d_constants[name] = input_db->getInteger(name);
        }
    }

    // Initialize the parser(s) with data read in from the input database.
    if (input_db->isString("function"))
    {
        // Assume scalar-valued function.
        d_function_strings.push_back(input_db->getString("function"));
        d_parsers.resize(1);
        try
        {
            d_parsers.back().SetExpr(d_function_strings.back());
        }
        catch (mu::ParserError& e)
        {
            TBOX_ERROR("muParserCartGridFunction::setDataOnPatch():\n"
                       << "  error: " << e.GetMsg() << "\n"
                       << "  in:    " << e.GetExpr() << "\n");
        }
        catch (...)
        {
            TBOX_ERROR("muParserCartGridFunction::setDataOnPatch():\n"
                       << "  unrecognized exception generated by muParser library.\n");
        }
    }
    else
    {
        // Assume vector-valued function.
        int d = 0;
        std::string key_name = "function_0";
        while (input_db->isString(key_name))
        {
            d_function_strings.push_back(input_db->getString(key_name));
            d_parsers.resize(d_parsers.size() + 1);
            try
            {
                d_parsers.back().SetExpr(d_function_strings.back());
            }
            catch (mu::ParserError& e)
            {
                TBOX_ERROR("muParserCartGridFunction::setDataOnPatch():\n"
                           << "  error: " << e.GetMsg() << "\n"
                           << "  in:    " << e.GetExpr() << "\n");
            }
            catch (...)
            {
                TBOX_ERROR("muParserCartGridFunction::setDataOnPatch():\n"
                           << "  unrecognized exception generated by muParser library.\n");
            }
            ++d;
            key_name = "function_" + std::to_string(d);
        }
    }

    // Ensure that we've read in at least one function.
    if (d_function_strings.empty())
    {
        TBOX_ERROR("muParserCartGridFunction::muParserCartGridFunction():\n"
                   << "  no function keys found in input database.\n"
                   << "  note that function specifications are assumed to be strings." << std::endl);
    }

    // Define the default and user-provided constants.
    const double pi = 3.1415926535897932384626433832795;
    const double* const x_lower = grid_geom->getXLower();
    const double* const x_upper = grid_geom->getXUpper();
    for (auto& parser : d_parsers)
    {
        // Various names for pi.
        parser.DefineConst("pi", pi);
        parser.DefineConst("Pi", pi);
        parser.DefineConst("PI", pi);

        // The extents of the domain.
        for (unsigned int d = 0; d < NDIM; ++d)
        {
            const std::string postfix = std::to_string(d);

            parser.DefineConst("X_LOWER" + postfix, x_lower[d]);
            parser.DefineConst("X_lower" + postfix, x_lower[d]);
            parser.DefineConst("x_lower" + postfix, x_lower[d]);
            parser.DefineConst("x_LOWER" + postfix, x_lower[d]);
            parser.DefineConst("X_Lower" + postfix, x_lower[d]);
            parser.DefineConst("X_lower" + postfix, x_lower[d]);
            parser.DefineConst("XLower" + postfix, x_lower[d]);
            parser.DefineConst("Xlower" + postfix, x_lower[d]);
            parser.DefineConst("x_Lower" + postfix, x_lower[d]);
            parser.DefineConst("x_lower" + postfix, x_lower[d]);
            parser.DefineConst("xLower" + postfix, x_lower[d]);
            parser.DefineConst("xlower" + postfix, x_lower[d]);

            parser.DefineConst("X_LOWER_" + postfix, x_lower[d]);
            parser.DefineConst("X_lower_" + postfix, x_lower[d]);
            parser.DefineConst("x_lower_" + postfix, x_lower[d]);
            parser.DefineConst("x_LOWER_" + postfix, x_lower[d]);
            parser.DefineConst("X_Lower_" + postfix, x_lower[d]);
            parser.DefineConst("X_lower_" + postfix, x_lower[d]);
            parser.DefineConst("XLower_" + postfix, x_lower[d]);
            parser.DefineConst("Xlower_" + postfix, x_lower[d]);
            parser.DefineConst("x_Lower_" + postfix, x_lower[d]);
            parser.DefineConst("x_lower_" + postfix, x_lower[d]);
            parser.DefineConst("xLower_" + postfix, x_lower[d]);
            parser.DefineConst("xlower_" + postfix, x_lower[d]);

            parser.DefineConst("X_UPPER" + postfix, x_upper[d]);
            parser.DefineConst("X_upper" + postfix, x_upper[d]);
            parser.DefineConst("x_upper" + postfix, x_upper[d]);
            parser.DefineConst("x_UPPER" + postfix, x_upper[d]);
            parser.DefineConst("X_Upper" + postfix, x_upper[d]);
            parser.DefineConst("X_upper" + postfix, x_upper[d]);
            parser.DefineConst("XUpper" + postfix, x_upper[d]);
            parser.DefineConst("Xupper" + postfix, x_upper[d]);
            parser.DefineConst("x_Upper" + postfix, x_upper[d]);
            parser.DefineConst("x_upper" + postfix, x_upper[d]);
            parser.DefineConst("xUpper" + postfix, x_upper[d]);
            parser.DefineConst("xupper" + postfix, x_upper[d]);

            parser.DefineConst("X_UPPER_" + postfix, x_upper[d]);
            parser.DefineConst("X_upper_" + postfix, x_upper[d]);
            parser.DefineConst("x_upper_" + postfix, x_upper[d]);
            parser.DefineConst("x_UPPER_" + postfix, x_upper[d]);
            parser.DefineConst("X_Upper_" + postfix, x_upper[d]);
            parser.DefineConst("X_upper_" + postfix, x_upper[d]);
            parser.DefineConst("XUpper_" + postfix, x_upper[d]);
            parser.DefineConst("Xupper_" + postfix, x_upper[d]);
            parser.DefineConst("x_Upper_" + postfix, x_upper[d]);
            parser.DefineConst("x_upper_" + postfix, x_upper[d]);
            parser.DefineConst("xUpper_" + postfix, x_upper[d]);
            parser.DefineConst("xupper_" + postfix, x_upper[d]);
        }

        // User-provided constants.
        for (const auto& constant : d_constants)
        {
            parser.DefineConst(constant.first, constant.second);
        }

        // Variables.
        parser.DefineVar("T", &d_parser_time);
        parser.DefineVar("t", &d_parser_time);
        for (unsigned int d = 0; d < NDIM; ++d)
        {
            const std::string postfix = std::to_string(d);
            parser.DefineVar("X" + postfix, &(d_parser_posn[d]));
            parser.DefineVar("x" + postfix, &(d_parser_posn[d]));
            parser.DefineVar("X_" + postfix, &(d_parser_posn[d]));
            parser.DefineVar("x_" + postfix, &(d_parser_posn[d]));
        }
    }
    return;
} // muParserCartGridFunction

bool
muParserCartGridFunction::isTimeDependent() const
{
    return true;
} // isTimeDependent

void
muParserCartGridFunction::setDataOnPatch(const int data_idx,
                                         Pointer<Variable<NDIM> > /*var*/,
                                         Pointer<Patch<NDIM> > patch,
                                         const double data_time,
                                         const bool /*initial_time*/,
                                         Pointer<PatchLevel<NDIM> > /*level*/)
{
    d_parser_time = data_time;

    const Box<NDIM>& patch_box = patch->getBox();
    const hier::Index<NDIM>& patch_lower = patch_box.lower();
    Pointer<CartesianPatchGeometry<NDIM> > pgeom = patch->getPatchGeometry();

    const double* const XLower = pgeom->getXLower();
    const double* const dx = pgeom->getDx();

    // Set the data in the patch.
    Pointer<PatchData<NDIM> > data = patch->getPatchData(data_idx);
#if !defined(NDEBUG)
    TBOX_ASSERT(data);
#endif
    Pointer<CellData<NDIM, double> > cc_data = data;
    Pointer<FaceData<NDIM, double> > fc_data = data;
    Pointer<NodeData<NDIM, double> > nc_data = data;
    Pointer<SideData<NDIM, double> > sc_data = data;
    Pointer<EdgeData<NDIM, double> > ec_data = data;
    if (cc_data)
    {
#if !defined(NDEBUG)
        TBOX_ASSERT(d_parsers.size() == 1 || d_parsers.size() == static_cast<unsigned int>(cc_data->getDepth()));
#endif
        for (int data_depth = 0; data_depth < cc_data->getDepth(); ++data_depth)
        {
            const int function_depth = (d_parsers.size() == 1 ? 0 : data_depth);
            for (CellIterator<NDIM> ic(patch_box); ic; ic++)
            {
                const CellIndex<NDIM>& i = ic();
                for (unsigned int d = 0; d < NDIM; ++d)
                {
                    d_parser_posn[d] = XLower[d] + dx[d] * (static_cast<double>(i(d) - patch_lower(d)) + 0.5);
                }
                try
                {
                    (*cc_data)(i, data_depth) = d_parsers[function_depth].Eval();
                }
                catch (mu::ParserError& e)
                {
                    TBOX_ERROR("muParserCartGridFunction::setDataOnPatch():\n"
                               << "  error: " << e.GetMsg() << "\n"
                               << "  in:    " << e.GetExpr() << "\n");
                }
                catch (...)
                {
                    TBOX_ERROR("muParserCartGridFunction::setDataOnPatch():\n"
                               << "  unrecognized exception generated by muParser library.\n");
                }
            }
        }
    }
    else if (fc_data)
    {
#if !defined(NDEBUG)
        TBOX_ASSERT(d_parsers.size() == 1 || d_parsers.size() == NDIM ||
                    d_parsers.size() == static_cast<unsigned int>(fc_data->getDepth()) ||
                    d_parsers.size() == NDIM * static_cast<unsigned int>(fc_data->getDepth()));
#endif
        for (int data_depth = 0; data_depth < fc_data->getDepth(); ++data_depth)
        {
            for (unsigned int axis = 0; axis < NDIM; ++axis)
            {
                int function_depth = -1;
                const int parsers_size = static_cast<int>(d_parsers.size());
                const int fc_data_depth = fc_data->getDepth();
                if (parsers_size == 1)
                {
                    function_depth = 0;
                }
                else if (parsers_size == NDIM)
                {
                    function_depth = axis;
                }
                else if (parsers_size == fc_data_depth)
                {
                    function_depth = data_depth;
                }
                else if (parsers_size == NDIM * fc_data_depth)
                {
                    function_depth = NDIM * data_depth + axis;
                }

                for (FaceIterator<NDIM> ic(patch_box, axis); ic; ic++)
                {
                    const FaceIndex<NDIM>& i = ic();
                    const hier::Index<NDIM>& cell_idx = i.toCell(1);
                    for (unsigned int d = 0; d < NDIM; ++d)
                    {
                        if (d == axis)
                        {
                            d_parser_posn[d] = XLower[d] + dx[d] * (static_cast<double>(cell_idx(d) - patch_lower(d)));
                        }
                        else
                        {
                            d_parser_posn[d] =
                                XLower[d] + dx[d] * (static_cast<double>(cell_idx(d) - patch_lower(d)) + 0.5);
                        }
                    }
                    try
                    {
                        (*fc_data)(i, data_depth) = d_parsers[function_depth].Eval();
                    }
                    catch (mu::ParserError& e)
                    {
                        TBOX_ERROR("muParserCartGridFunction::setDataOnPatch():\n"
                                   << "  error: " << e.GetMsg() << "\n"
                                   << "  in:    " << e.GetExpr() << "\n");
                    }
                    catch (...)
                    {
                        TBOX_ERROR("muParserCartGridFunction::setDataOnPatch():\n"
                                   << "  unrecognized exception generated by muParser library.\n");
                    }
                }
            }
        }
    }
    else if (nc_data)
    {
#if !defined(NDEBUG)
        TBOX_ASSERT(d_parsers.size() == 1 || d_parsers.size() == static_cast<unsigned int>(nc_data->getDepth()));
#endif
        for (int data_depth = 0; data_depth < nc_data->getDepth(); ++data_depth)
        {
            const int function_depth = (d_parsers.size() == 1 ? 0 : data_depth);
            for (NodeIterator<NDIM> ic(patch_box); ic; ic++)
            {
                const NodeIndex<NDIM>& i = ic();
                for (unsigned int d = 0; d < NDIM; ++d)
                {
                    d_parser_posn[d] = XLower[d] + dx[d] * (static_cast<double>(i(d) - patch_lower(d)));
                }
                try
                {
                    (*nc_data)(i, data_depth) = d_parsers[function_depth].Eval();
                }
                catch (mu::ParserError& e)
                {
                    TBOX_ERROR("muParserCartGridFunction::setDataOnPatch():\n"
                               << "  error: " << e.GetMsg() << "\n"
                               << "  in:    " << e.GetExpr() << "\n");
                }
                catch (...)
                {
                    TBOX_ERROR("muParserCartGridFunction::setDataOnPatch():\n"
                               << "  unrecognized exception generated by muParser library.\n");
                }
            }
        }
    }
    else if (sc_data)
    {
#if !defined(NDEBUG)
        TBOX_ASSERT(d_parsers.size() == 1 || d_parsers.size() == NDIM ||
                    d_parsers.size() == static_cast<unsigned int>(sc_data->getDepth()) ||
                    d_parsers.size() == NDIM * static_cast<unsigned int>(sc_data->getDepth()));
#endif
        for (int data_depth = 0; data_depth < sc_data->getDepth(); ++data_depth)
        {
            for (unsigned int axis = 0; axis < NDIM; ++axis)
            {
                int function_depth = -1;
                const int parsers_size = static_cast<int>(d_parsers.size());
                const int sc_data_depth = sc_data->getDepth();
                if (parsers_size == 1)
                {
                    function_depth = 0;
                }
                else if (parsers_size == NDIM)
                {
                    function_depth = axis;
                }
                else if (parsers_size == sc_data_depth)
                {
                    function_depth = data_depth;
                }
                else if (parsers_size == NDIM * sc_data_depth)
                {
                    function_depth = NDIM * data_depth + axis;
                }

                for (SideIterator<NDIM> ic(patch_box, axis); ic; ic++)
                {
                    const SideIndex<NDIM>& i = ic();
                    for (unsigned int d = 0; d < NDIM; ++d)
                    {
                        if (d == axis)
                        {
                            d_parser_posn[d] = XLower[d] + dx[d] * (static_cast<double>(i(d) - patch_lower(d)));
                        }
                        else
                        {
                            d_parser_posn[d] = XLower[d] + dx[d] * (static_cast<double>(i(d) - patch_lower(d)) + 0.5);
                        }
                    }
                    try
                    {
                        (*sc_data)(i, data_depth) = d_parsers[function_depth].Eval();
                    }
                    catch (mu::ParserError& e)
                    {
                        TBOX_ERROR("muParserCartGridFunction::setDataOnPatch():\n"
                                   << "  error: " << e.GetMsg() << "\n"
                                   << "  in:    " << e.GetExpr() << "\n");
                    }
                    catch (...)
                    {
                        TBOX_ERROR("muParserCartGridFunction::setDataOnPatch():\n"
                                   << "  unrecognized exception generated by muParser library.\n");
                    }
                }
            }
        }
    }
    else if (ec_data)
    {
#if !defined(NDEBUG)
        TBOX_ASSERT(d_parsers.size() == 1 || d_parsers.size() == NDIM ||
                    d_parsers.size() == static_cast<unsigned int>(ec_data->getDepth()) ||
                    d_parsers.size() == NDIM * static_cast<unsigned int>(ec_data->getDepth()));
#endif
        for (int data_depth = 0; data_depth < ec_data->getDepth(); ++data_depth)
        {
            for (unsigned int axis = 0; axis < NDIM; ++axis)
            {
                int function_depth = -1;
                const int parsers_size = static_cast<int>(d_parsers.size());
                const int ec_data_depth = ec_data->getDepth();
                if (parsers_size == 1)
                {
                    function_depth = 0;
                }
                else if (parsers_size == NDIM)
                {
                    function_depth = axis;
                }
                else if (parsers_size == ec_data_depth)
                {
                    function_depth = data_depth;
                }
                else if (parsers_size == NDIM * ec_data_depth)
                {
                    function_depth = NDIM * data_depth + axis;
                }

                for (EdgeIterator<NDIM> ic(patch_box, axis); ic; ic++)
                {
                    const EdgeIndex<NDIM>& i = ic();
                    for (unsigned int d = 0; d < NDIM; ++d)
                    {
                        if (d == axis)
                        {
                            d_parser_posn[d] = XLower[d] + dx[d] * (static_cast<double>(i(d) - patch_lower(d)) + 0.5);
                        }
                        else
                        {
                            d_parser_posn[d] = XLower[d] + dx[d] * (static_cast<double>(i(d) - patch_lower(d)));
                        }
                    }
                    try
                    {
                        (*ec_data)(i, data_depth) = d_parsers[function_depth].Eval();
                    }
                    catch (mu::ParserError& e)
                    {
                        TBOX_ERROR("muParserCartGridFunction::setDataOnPatch():\n"
                                   << "  error: " << e.GetMsg() << "\n"
                                   << "  in:    " << e.GetExpr() << "\n");
                    }
                    catch (...)
                    {
                        TBOX_ERROR("muParserCartGridFunction::setDataOnPatch():\n"
                                   << "  unrecognized exception generated by muParser library.\n");
                    }
                }
            }
        }
    }
    else
    {
        TBOX_ERROR("muParserCartGridFunction::setDataOnPatch():\n"
                   << "  unsupported patch data type encountered." << std::endl);
    }
    return;
} // setDataOnPatch

//////////////////////////////////////////////////////////////////////////////

} // namespace IBTK

//////////////////////////////////////////////////////////////////////////////
