// ---------------------------------------------------------------------
//
// Copyright (c) 2017 - 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.
//
// ---------------------------------------------------------------------

// Config files
#include <SAMRAI_config.h>

// Headers for basic PETSc functions
#include <petscsys.h>

// Headers for basic SAMRAI objects
#include <BergerRigoutsos.h>
#include <CartesianGridGeometry.h>
#include <LoadBalancer.h>
#include <StandardTagAndInitialize.h>

// Headers for basic libMesh objects
#include <libmesh/boundary_info.h>
#include <libmesh/equation_systems.h>
#include <libmesh/exodusII_io.h>
#include <libmesh/mesh.h>
#include <libmesh/mesh_function.h>
#include <libmesh/mesh_generation.h>
#include <libmesh/mesh_triangle_interface.h>

// Headers for application-specific algorithm/data structure objects
#include <ibamr/IBExplicitHierarchyIntegrator.h>
#include <ibamr/IBFEMethod.h>
#include <ibamr/INSCollocatedHierarchyIntegrator.h>
#include <ibamr/INSStaggeredHierarchyIntegrator.h>

#include <ibtk/AppInitializer.h>
#include <ibtk/IBTKInit.h>
#include <ibtk/IBTK_MPI.h>
#include <ibtk/libmesh_utilities.h>
#include <ibtk/muParserCartGridFunction.h>
#include <ibtk/muParserRobinBcCoefs.h>

#include <boost/multi_array.hpp>

// Set up application namespace declarations
#include <ibamr/app_namespaces.h>

// Elasticity model data.
namespace ModelData
{
// The tether penalty functions each require some data that is set in the
// input file. This data is passed to each object through the void *ctx
// context data pointer. Here we collect all relevant tether data in a struct:
struct ElasticityData
{
    const double c1_s;
    const double kappa_s;
    const double mu_s;
    const double lambda_s;

    ElasticityData(Pointer<Database> input_db)
        : c1_s(input_db->getDouble("C1_S")),
          kappa_s(input_db->getDouble("KAPPA_S")),
          mu_s(input_db->getDouble("MU_S")),
          lambda_s(input_db->getDouble("LAMBDA_S"))
    {
    }
};

// Tether (penalty) force function for the solid block.
void
block_tether_force_function(VectorValue<double>& F,
                            const TensorValue<double>& /*FF*/,
                            const libMesh::Point& X,
                            const libMesh::Point& s,
                            Elem* const /*elem*/,
                            const std::vector<const std::vector<double>*>& /*var_data*/,
                            const std::vector<const std::vector<VectorValue<double> >*>& /*grad_var_data*/,
                            double /*time*/,
                            void* ctx)
{
    const ElasticityData* const elasticity_data = reinterpret_cast<ElasticityData*>(ctx);

    F = elasticity_data->kappa_s * (s - X);
    return;
} // block_tether_force_function

// Tether (penalty) force function for the thin beam.
void
beam_tether_force_function(VectorValue<double>& F,
                           const TensorValue<double>& /*FF*/,
                           const libMesh::Point& X,
                           const libMesh::Point& s,
                           Elem* const /*elem*/,
                           const std::vector<const std::vector<double>*>& /*var_data*/,
                           const std::vector<const std::vector<VectorValue<double> >*>& /*grad_var_data*/,
                           double /*time*/,
                           void* ctx)
{
    const double r = sqrt((s(0) - 0.2) * (s(0) - 0.2) + (s(1) - 0.2) * (s(1) - 0.2));
    if (r <= 0.05)
    {
        const ElasticityData* const elasticity_data = reinterpret_cast<ElasticityData*>(ctx);
        F = elasticity_data->kappa_s * (s - X);
    }
    else
    {
        F.zero();
    }
    return;
} // beam_tether_force_function

// (Penalty) stress tensor function for the solid block.
void
block_PK1_stress_function(TensorValue<double>& PP,
                          const TensorValue<double>& FF,
                          const libMesh::Point& /*X*/,
                          const libMesh::Point& /*s*/,
                          Elem* const /*elem*/,
                          const std::vector<const std::vector<double>*>& /*var_data*/,
                          const std::vector<const std::vector<VectorValue<double> >*>& /*grad_var_data*/,
                          double /*time*/,
                          void* ctx)
{
    const ElasticityData* const elasticity_data = reinterpret_cast<ElasticityData*>(ctx);

    PP = 2.0 * elasticity_data->c1_s * (FF - tensor_inverse_transpose(FF, NDIM));
    return;
} // block_PK1_stress_function

void
beam_PK1_stress_function(TensorValue<double>& PP,
                         const TensorValue<double>& FF,
                         const libMesh::Point& /*X*/,
                         const libMesh::Point& /*s*/,
                         Elem* const /*elem*/,
                         const std::vector<const std::vector<double>*>& /*var_data*/,
                         const std::vector<const std::vector<VectorValue<double> >*>& /*grad_var_data*/,
                         double /*time*/,
                         void* ctx)
{
    const ElasticityData* const elasticity_data = reinterpret_cast<ElasticityData*>(ctx);
    static const TensorValue<double> II(1.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 1.0);
    const TensorValue<double> CC = FF.transpose() * FF;
    const TensorValue<double> EE = 0.5 * (CC - II);
    const TensorValue<double> SS = elasticity_data->lambda_s * EE.tr() * II + 2.0 * elasticity_data->mu_s * EE;
    PP = FF * SS;
    return;
} // beam_PK1_stress_function
} // namespace ModelData
using namespace ModelData;

// Function prototypes
static ofstream drag_stream, lift_stream, A_x_posn_stream, A_y_posn_stream;
void postprocess_data(Pointer<PatchHierarchy<NDIM> > patch_hierarchy,
                      Pointer<INSHierarchyIntegrator> navier_stokes_integrator,
                      Mesh& beam_mesh,
                      EquationSystems* beam_equation_systems,
                      Mesh& block_mesh,
                      EquationSystems* block_equation_systems,
                      const int iteration_num,
                      const double loop_time,
                      const string& data_dump_dirname);

/*******************************************************************************
 * For each run, the input filename and restart information (if needed) must   *
 * be given on the command line.  For non-restarted case, command line is:     *
 *                                                                             *
 *    executable <input file name>                                             *
 *                                                                             *
 * For restarted run, command line is:                                         *
 *                                                                             *
 *    executable <input file name> <restart directory> <restart number>        *
 *                                                                             *
 *******************************************************************************/

int
main(int argc, char* argv[])
{
    // Initialize IBAMR and libraries. Deinitialization is handled by this object as well.
    IBTKInit ibtk_init(argc, argv, MPI_COMM_WORLD);
    const LibMeshInit& init = ibtk_init.getLibMeshInit();

    { // cleanup dynamically allocated objects prior to shutdown

        // Parse command line options, set some standard options from the input
        // file, initialize the restart database (if this is a restarted run),
        // and enable file logging.
        Pointer<AppInitializer> app_initializer = new AppInitializer(argc, argv, "IB.log");
        Pointer<Database> input_db = app_initializer->getInputDatabase();

        // Get various standard options set in the input file.
        const bool dump_viz_data = app_initializer->dumpVizData();
        const int viz_dump_interval = app_initializer->getVizDumpInterval();
        const bool uses_visit = dump_viz_data && app_initializer->getVisItDataWriter();
#ifdef LIBMESH_HAVE_EXODUS_API
        const bool uses_exodus = dump_viz_data && !app_initializer->getExodusIIFilename().empty();
#else
        const bool uses_exodus = false;
        if (!app_initializer->getExodusIIFilename().empty())
        {
            plog << "WARNING: libMesh was compiled without Exodus support, so no "
                 << "Exodus output will be written in this program.\n";
        }
#endif
        const string block_exodus_filename = app_initializer->getExodusIIFilename("block");
        const string beam_exodus_filename = app_initializer->getExodusIIFilename("beam");

        const bool dump_restart_data = app_initializer->dumpRestartData();
        const int restart_dump_interval = app_initializer->getRestartDumpInterval();
        const string restart_dump_dirname = app_initializer->getRestartDumpDirectory();
        const string restart_read_dirname = app_initializer->getRestartReadDirectory();
        const int restart_restore_num = app_initializer->getRestartRestoreNumber();

        const bool dump_postproc_data = app_initializer->dumpPostProcessingData();
        const int postproc_data_dump_interval = app_initializer->getPostProcessingDataDumpInterval();
        const string postproc_data_dump_dirname = app_initializer->getPostProcessingDataDumpDirectory();
        if (dump_postproc_data && (postproc_data_dump_interval > 0) && !postproc_data_dump_dirname.empty())
        {
            Utilities::recursiveMkdir(postproc_data_dump_dirname);
        }

        const bool dump_timer_data = app_initializer->dumpTimerData();
        const int timer_dump_interval = app_initializer->getTimerDumpInterval();

        // Create a simple FE mesh.
        const double dx = input_db->getDouble("DX");
        const double ds = input_db->getDouble("MFAC") * dx;

        Mesh block_mesh(init.comm(), NDIM);
        string block_elem_type = input_db->getString("BLOCK_ELEM_TYPE");
        const double R = 0.05;
        if (block_elem_type == "TRI3" || block_elem_type == "TRI6")
        {
#ifdef LIBMESH_HAVE_TRIANGLE
            const int num_circum_nodes = ceil(2.0 * M_PI * R / ds);
            for (int k = 0; k < num_circum_nodes; ++k)
            {
                const double theta = 2.0 * M_PI * static_cast<double>(k) / static_cast<double>(num_circum_nodes);
                block_mesh.add_point(libMesh::Point(R * cos(theta), R * sin(theta)));
            }
            TriangleInterface triangle(block_mesh);
            triangle.triangulation_type() = TriangleInterface::GENERATE_CONVEX_HULL;
            triangle.elem_type() = Utility::string_to_enum<ElemType>(block_elem_type);
            triangle.desired_area() = sqrt(3.0) / 4.0 * ds * ds;
            triangle.insert_extra_points() = true;
            triangle.smooth_after_generating() = true;
            triangle.triangulate();
            block_mesh.prepare_for_use();
#else
            TBOX_ERROR("ERROR: libMesh appears to have been configured without support for Triangle,\n"
                       << "       but Triangle is required for TRI3 or TRI6 elements.\n");
#endif
        }
        else
        {
            // NOTE: number of segments along boundary is 4*2^r.
            const double num_circum_segments = ceil(2.0 * M_PI * R / ds);
            const int r = log2(0.25 * num_circum_segments);
            MeshTools::Generation::build_sphere(block_mesh, R, r, Utility::string_to_enum<ElemType>(block_elem_type));
        }
        for (MeshBase::node_iterator n_it = block_mesh.nodes_begin(); n_it != block_mesh.nodes_end(); ++n_it)
        {
            Node& n = **n_it;
            n(0) += 0.2;
            n(1) += 0.2;
        }

        Mesh beam_mesh(init.comm(), NDIM);
        string beam_elem_type = input_db->getString("BEAM_ELEM_TYPE");
        MeshTools::Generation::build_square(beam_mesh,
                                            ceil(0.4 / ds),
                                            ceil(0.02 / ds),
                                            0.2,
                                            0.6,
                                            0.19,
                                            0.21,
                                            Utility::string_to_enum<ElemType>(beam_elem_type));
        beam_mesh.prepare_for_use();

        vector<MeshBase*> meshes(2);
        meshes[0] = &block_mesh;
        meshes[1] = &beam_mesh;

        // Create major algorithm and data objects that comprise the
        // application.  These objects are configured from the input database
        // and, if this is a restarted run, from the restart database.
        Pointer<INSHierarchyIntegrator> navier_stokes_integrator;
        const string solver_type = app_initializer->getComponentDatabase("Main")->getString("solver_type");
        if (solver_type == "STAGGERED")
        {
            navier_stokes_integrator = new INSStaggeredHierarchyIntegrator(
                "INSStaggeredHierarchyIntegrator",
                app_initializer->getComponentDatabase("INSStaggeredHierarchyIntegrator"));
        }
        else if (solver_type == "COLLOCATED")
        {
            navier_stokes_integrator = new INSCollocatedHierarchyIntegrator(
                "INSCollocatedHierarchyIntegrator",
                app_initializer->getComponentDatabase("INSCollocatedHierarchyIntegrator"));
        }
        else
        {
            TBOX_ERROR("Unsupported solver type: " << solver_type << "\n"
                                                   << "Valid options are: COLLOCATED, STAGGERED");
        }
        Pointer<IBFEMethod> ib_method_ops =
            new IBFEMethod("IBFEMethod",
                           app_initializer->getComponentDatabase("IBFEMethod"),
                           meshes,
                           app_initializer->getComponentDatabase("GriddingAlgorithm")->getInteger("max_levels"),
                           /*register_for_restart*/ true,
                           restart_read_dirname,
                           restart_restore_num);
        Pointer<IBHierarchyIntegrator> time_integrator =
            new IBExplicitHierarchyIntegrator("IBHierarchyIntegrator",
                                              app_initializer->getComponentDatabase("IBHierarchyIntegrator"),
                                              ib_method_ops,
                                              navier_stokes_integrator);
        Pointer<CartesianGridGeometry<NDIM> > grid_geometry = new CartesianGridGeometry<NDIM>(
            "CartesianGeometry", app_initializer->getComponentDatabase("CartesianGeometry"));
        Pointer<PatchHierarchy<NDIM> > patch_hierarchy = new PatchHierarchy<NDIM>("PatchHierarchy", grid_geometry);
        Pointer<StandardTagAndInitialize<NDIM> > error_detector =
            new StandardTagAndInitialize<NDIM>("StandardTagAndInitialize",
                                               time_integrator,
                                               app_initializer->getComponentDatabase("StandardTagAndInitialize"));
        Pointer<BergerRigoutsos<NDIM> > box_generator = new BergerRigoutsos<NDIM>();
        Pointer<LoadBalancer<NDIM> > load_balancer =
            new LoadBalancer<NDIM>("LoadBalancer", app_initializer->getComponentDatabase("LoadBalancer"));
        Pointer<GriddingAlgorithm<NDIM> > gridding_algorithm =
            new GriddingAlgorithm<NDIM>("GriddingAlgorithm",
                                        app_initializer->getComponentDatabase("GriddingAlgorithm"),
                                        error_detector,
                                        box_generator,
                                        load_balancer);

        // Configure the IBFE solver.
        ElasticityData elasticity_data(input_db);
        void* const elasticity_data_ptr = reinterpret_cast<void*>(&elasticity_data);
        IBFEMethod::LagBodyForceFcnData block_tether_force_data(
            block_tether_force_function, std::vector<IBTK::SystemData>(), elasticity_data_ptr);
        IBFEMethod::PK1StressFcnData block_PK1_stress_data(
            block_PK1_stress_function, std::vector<IBTK::SystemData>(), elasticity_data_ptr);
        ib_method_ops->registerLagBodyForceFunction(block_tether_force_data, 0);
        ib_method_ops->registerPK1StressFunction(block_PK1_stress_data, 0);
        string block_kernel_fcn = input_db->getStringWithDefault("BLOCK_KERNEL_FUNCTION", "PIECEWISE_LINEAR");
        FEDataManager::InterpSpec block_interp_spec = ib_method_ops->getDefaultInterpSpec();
        block_interp_spec.kernel_fcn = block_kernel_fcn;
        ib_method_ops->setInterpSpec(block_interp_spec, 0);
        FEDataManager::SpreadSpec block_spread_spec = ib_method_ops->getDefaultSpreadSpec();
        block_spread_spec.kernel_fcn = block_kernel_fcn;
        ib_method_ops->setSpreadSpec(block_spread_spec, 0);

        IBFEMethod::LagBodyForceFcnData beam_tether_force_data(
            beam_tether_force_function, std::vector<IBTK::SystemData>(), elasticity_data_ptr);
        IBFEMethod::PK1StressFcnData beam_PK1_stress_data(
            beam_PK1_stress_function, std::vector<IBTK::SystemData>(), elasticity_data_ptr);
        ib_method_ops->registerLagBodyForceFunction(beam_tether_force_data, 1);
        ib_method_ops->registerPK1StressFunction(beam_PK1_stress_data, 1);
        string beam_kernel_fcn = input_db->getStringWithDefault("BEAM_KERNEL_FUNCTION", "IB_3");
        FEDataManager::InterpSpec beam_interp_spec = ib_method_ops->getDefaultInterpSpec();
        beam_interp_spec.kernel_fcn = beam_kernel_fcn;
        ib_method_ops->setInterpSpec(beam_interp_spec, 1);
        FEDataManager::SpreadSpec beam_spread_spec = ib_method_ops->getDefaultSpreadSpec();
        beam_spread_spec.kernel_fcn = beam_kernel_fcn;
        ib_method_ops->setSpreadSpec(beam_spread_spec, 1);

        ib_method_ops->initializeFEEquationSystems();
        EquationSystems* block_equation_systems = ib_method_ops->getFEDataManager(0)->getEquationSystems();
        EquationSystems* beam_equation_systems = ib_method_ops->getFEDataManager(1)->getEquationSystems();

        // Create Eulerian initial condition specification objects.
        if (input_db->keyExists("VelocityInitialConditions"))
        {
            Pointer<CartGridFunction> u_init = new muParserCartGridFunction(
                "u_init", app_initializer->getComponentDatabase("VelocityInitialConditions"), grid_geometry);
            navier_stokes_integrator->registerVelocityInitialConditions(u_init);
        }

        if (input_db->keyExists("PressureInitialConditions"))
        {
            Pointer<CartGridFunction> p_init = new muParserCartGridFunction(
                "p_init", app_initializer->getComponentDatabase("PressureInitialConditions"), grid_geometry);
            navier_stokes_integrator->registerPressureInitialConditions(p_init);
        }

        // Create Eulerian boundary condition specification objects (when necessary).
        const IntVector<NDIM>& periodic_shift = grid_geometry->getPeriodicShift();
        vector<RobinBcCoefStrategy<NDIM>*> u_bc_coefs(NDIM);
        if (periodic_shift.min() > 0)
        {
            for (unsigned int d = 0; d < NDIM; ++d)
            {
                u_bc_coefs[d] = NULL;
            }
        }
        else
        {
            for (unsigned int d = 0; d < NDIM; ++d)
            {
                const std::string bc_coefs_name = "u_bc_coefs_" + std::to_string(d);

                const std::string bc_coefs_db_name = "VelocityBcCoefs_" + std::to_string(d);

                u_bc_coefs[d] = new muParserRobinBcCoefs(
                    bc_coefs_name, app_initializer->getComponentDatabase(bc_coefs_db_name), grid_geometry);
            }
            navier_stokes_integrator->registerPhysicalBoundaryConditions(u_bc_coefs);
        }

        // Create Eulerian body force function specification objects.
        if (input_db->keyExists("ForcingFunction"))
        {
            Pointer<CartGridFunction> f_fcn = new muParserCartGridFunction(
                "f_fcn", app_initializer->getComponentDatabase("ForcingFunction"), grid_geometry);
            time_integrator->registerBodyForceFunction(f_fcn);
        }

        // Set up visualization plot file writers.
        Pointer<VisItDataWriter<NDIM> > visit_data_writer = app_initializer->getVisItDataWriter();
        if (uses_visit)
        {
            time_integrator->registerVisItDataWriter(visit_data_writer);
        }
        std::unique_ptr<ExodusII_IO> block_exodus_io(uses_exodus ? new ExodusII_IO(block_mesh) : NULL);
        std::unique_ptr<ExodusII_IO> beam_exodus_io(uses_exodus ? new ExodusII_IO(beam_mesh) : NULL);

        // Check to see if this is a restarted run to append current exodus files
        if (uses_exodus)
        {
            const bool from_restart = RestartManager::getManager()->isFromRestart();
            block_exodus_io->append(from_restart);
            beam_exodus_io->append(from_restart);
        }

        // Initialize hierarchy configuration and data on all patches.
        ib_method_ops->initializeFEData();
        time_integrator->initializePatchHierarchy(patch_hierarchy, gridding_algorithm);

        // Deallocate initialization objects.
        app_initializer.setNull();

        // Print the input database contents to the log file.
        plog << "Input database:\n";
        input_db->printClassData(plog);

        // Write out initial visualization data.
        int iteration_num = time_integrator->getIntegratorStep();
        double loop_time = time_integrator->getIntegratorTime();
        if (dump_viz_data)
        {
            pout << "\n\nWriting visualization files...\n\n";
            if (uses_visit)
            {
                time_integrator->setupPlotData();
                visit_data_writer->writePlotData(patch_hierarchy, iteration_num, loop_time);
            }
            if (uses_exodus)
            {
                block_exodus_io->write_timestep(
                    block_exodus_filename, *block_equation_systems, iteration_num / viz_dump_interval + 1, loop_time);
                beam_exodus_io->write_timestep(
                    beam_exodus_filename, *beam_equation_systems, iteration_num / viz_dump_interval + 1, loop_time);
            }
        }

        // Open streams to save lift and drag coefficients.
        if (IBTK_MPI::getRank() == 0)
        {
            drag_stream.open("C_D.curve", ios_base::out | ios_base::trunc);
            lift_stream.open("C_L.curve", ios_base::out | ios_base::trunc);
            A_x_posn_stream.open("A_x.curve", ios_base::out | ios_base::trunc);
            A_y_posn_stream.open("A_y.curve", ios_base::out | ios_base::trunc);
        }

        // Main time step loop.
        double loop_time_end = time_integrator->getEndTime();
        double dt = 0.0;
        while (!MathUtilities<double>::equalEps(loop_time, loop_time_end) && time_integrator->stepsRemaining())
        {
            iteration_num = time_integrator->getIntegratorStep();
            loop_time = time_integrator->getIntegratorTime();

            pout << "\n";
            pout << "+++++++++++++++++++++++++++++++++++++++++++++++++++\n";
            pout << "At beginning of timestep # " << iteration_num << "\n";
            pout << "Simulation time is " << loop_time << "\n";

            dt = time_integrator->getMaximumTimeStepSize();
            time_integrator->advanceHierarchy(dt);
            loop_time += dt;

            pout << "\n";
            pout << "At end       of timestep # " << iteration_num << "\n";
            pout << "Simulation time is " << loop_time << "\n";
            pout << "+++++++++++++++++++++++++++++++++++++++++++++++++++\n";
            pout << "\n";

            // At specified intervals, write visualization and restart files,
            // print out timer data, and store hierarchy data for post
            // processing.
            iteration_num += 1;
            const bool last_step = !time_integrator->stepsRemaining();
            if (dump_viz_data && (iteration_num % viz_dump_interval == 0 || last_step))
            {
                pout << "\nWriting visualization files...\n\n";
                if (uses_visit)
                {
                    time_integrator->setupPlotData();
                    visit_data_writer->writePlotData(patch_hierarchy, iteration_num, loop_time);
                }
                if (uses_exodus)
                {
                    block_exodus_io->write_timestep(block_exodus_filename,
                                                    *block_equation_systems,
                                                    iteration_num / viz_dump_interval + 1,
                                                    loop_time);
                    beam_exodus_io->write_timestep(
                        beam_exodus_filename, *beam_equation_systems, iteration_num / viz_dump_interval + 1, loop_time);
                }
            }
            if (dump_restart_data && (iteration_num % restart_dump_interval == 0 || last_step))
            {
                pout << "\nWriting restart files...\n\n";
                RestartManager::getManager()->writeRestartFile(restart_dump_dirname, iteration_num);
                ib_method_ops->writeFEDataToRestartFile(restart_dump_dirname, iteration_num);
            }
            if (dump_timer_data && (iteration_num % timer_dump_interval == 0 || last_step))
            {
                pout << "\nWriting timer data...\n\n";
                TimerManager::getManager()->print(plog);
            }
            if (dump_postproc_data && (iteration_num % postproc_data_dump_interval == 0 || last_step))
            {
                pout << "\nWriting state data...\n\n";
                postprocess_data(patch_hierarchy,
                                 navier_stokes_integrator,
                                 beam_mesh,
                                 beam_equation_systems,
                                 block_mesh,
                                 block_equation_systems,
                                 iteration_num,
                                 loop_time,
                                 postproc_data_dump_dirname);
            }
        }

        // Close the logging streams.
        if (IBTK_MPI::getRank() == 0)
        {
            drag_stream.close();
            lift_stream.close();
            A_x_posn_stream.close();
            A_y_posn_stream.close();
        }

        // Cleanup Eulerian boundary condition specification objects (when
        // necessary).
        for (unsigned int d = 0; d < NDIM; ++d) delete u_bc_coefs[d];

    } // cleanup dynamically allocated objects prior to shutdown
} // main

void
postprocess_data(Pointer<PatchHierarchy<NDIM> > /*patch_hierarchy*/,
                 Pointer<INSHierarchyIntegrator> /*navier_stokes_integrator*/,
                 Mesh& beam_mesh,
                 EquationSystems* beam_equation_systems,
                 Mesh& block_mesh,
                 EquationSystems* block_equation_systems,
                 const int /*iteration_num*/,
                 const double loop_time,
                 const string& /*data_dump_dirname*/)
{
    double F_integral[NDIM];
    for (unsigned int d = 0; d < NDIM; ++d) F_integral[d] = 0.0;
    Mesh* mesh[2] = { &beam_mesh, &block_mesh };
    EquationSystems* equation_systems[2] = { beam_equation_systems, block_equation_systems };
    for (unsigned int k = 0; k < 2; ++k)
    {
        System& F_system = equation_systems[k]->get_system<System>(IBFEMethod::FORCE_SYSTEM_NAME);
        NumericVector<double>* F_vec = F_system.solution.get();
        NumericVector<double>* F_ghost_vec = F_system.current_local_solution.get();
        copy_and_synch(*F_vec, *F_ghost_vec);
        DofMap& F_dof_map = F_system.get_dof_map();
        std::vector<std::vector<unsigned int> > F_dof_indices(NDIM);
        std::unique_ptr<FEBase> fe(FEBase::build(NDIM, F_dof_map.variable_type(0)));
        std::unique_ptr<QBase> qrule = QBase::build(QGAUSS, NDIM, FIFTH);
        fe->attach_quadrature_rule(qrule.get());
        const std::vector<std::vector<double> >& phi = fe->get_phi();
        const std::vector<double>& JxW = fe->get_JxW();
        boost::multi_array<double, 2> F_node;
        const MeshBase::const_element_iterator el_begin = mesh[k]->active_local_elements_begin();
        const MeshBase::const_element_iterator el_end = mesh[k]->active_local_elements_end();
        for (MeshBase::const_element_iterator el_it = el_begin; el_it != el_end; ++el_it)
        {
            Elem* const elem = *el_it;
            fe->reinit(elem);
            for (unsigned int d = 0; d < NDIM; ++d)
            {
                F_dof_map.dof_indices(elem, F_dof_indices[d], d);
            }
            const int n_qp = qrule->n_points();
            const int n_basis = static_cast<int>(F_dof_indices[0].size());
            get_values_for_interpolation(F_node, *F_ghost_vec, F_dof_indices);
            for (int qp = 0; qp < n_qp; ++qp)
            {
                for (int k = 0; k < n_basis; ++k)
                {
                    for (int d = 0; d < NDIM; ++d)
                    {
                        F_integral[d] += F_node[k][d] * phi[k][qp] * JxW[qp];
                    }
                }
            }
        }
    }
    IBTK_MPI::sumReduction(F_integral, NDIM);
    if (IBTK_MPI::getRank() == 0)
    {
        drag_stream.precision(12);
        drag_stream.setf(ios::fixed, ios::floatfield);
        drag_stream << loop_time << " " << -F_integral[0] << endl;
        lift_stream.precision(12);
        lift_stream.setf(ios::fixed, ios::floatfield);
        lift_stream << loop_time << " " << -F_integral[1] << endl;
    }

    System& X_system = beam_equation_systems->get_system<System>(IBFEMethod::COORDS_SYSTEM_NAME);
    NumericVector<double>* X_vec = X_system.solution.get();
    std::unique_ptr<NumericVector<Number> > X_serial_vec = NumericVector<Number>::build(X_vec->comm());
    X_serial_vec->init(X_vec->size(), true, SERIAL);
    X_vec->localize(*X_serial_vec);
    DofMap& X_dof_map = X_system.get_dof_map();
    vector<unsigned int> vars(2);
    vars[0] = 0;
    vars[1] = 1;
    MeshFunction X_fcn(*beam_equation_systems, *X_serial_vec, X_dof_map, vars);
    X_fcn.init();
    DenseVector<double> X_A(2);
    X_fcn(libMesh::Point(0.6, 0.2, 0), 0.0, X_A);
    if (IBTK_MPI::getRank() == 0)
    {
        A_x_posn_stream.precision(12);
        A_x_posn_stream.setf(ios::fixed, ios::floatfield);
        A_x_posn_stream << loop_time << " " << X_A(0) << endl;
        A_y_posn_stream.precision(12);
        A_y_posn_stream.setf(ios::fixed, ios::floatfield);
        A_y_posn_stream << loop_time << " " << X_A(1) << endl;
    }
    return;
} // postprocess_data


#include "TimeIntegrator.h"

#include <exception>
#include <Eigen/Dense>
#include <Eigen/IterativeLinearSolvers>
#include <Eigen/SparseLU>
#include <cinder/Log.h>

namespace ar
{
    VectorX TimeIntegrator::solveDense(const MatrixX& A, const VectorX& b) const
    {
        switch (denseSolver)
        {
        case DenseLinearSolver::PartialPivLU: {return A.partialPivLu().solve(b); }
        case DenseLinearSolver::FullPivLU:
            {
                Eigen::FullPivLU<MatrixX> decomp = A.fullPivLu();
                decomp.setThreshold(1e-5);
                int rank = decomp.rank();
                if (rank == std::min(A.rows(), A.cols()))
                    CI_LOG_D("rank is " << decomp.rank() << " (of " << std::min(A.rows(), A.cols()) << ")");
                else
                    CI_LOG_W("Matrix is not of full rank, only " << rank << " of " << std::min(A.rows(), A.cols()));
                return decomp.solve(b);
            }
        case DenseLinearSolver::HouseholderQR: {return A.householderQr().solve(b); }
        case DenseLinearSolver::ColPivHousholderQR:
            {
                Eigen::ColPivHouseholderQR<MatrixX> decomp = A.colPivHouseholderQr();
                int rank = decomp.rank();
                if (rank == std::min(A.rows(), A.cols()))
                    CI_LOG_D("rank is " << decomp.rank() << " (of " << std::min(A.rows(), A.cols()) << ")");
                else
                    CI_LOG_W("Matrix is not of full rank, only " << rank << " of " << std::min(A.rows(), A.cols()));
                return decomp.solve(b);
            }
        case DenseLinearSolver::FullPivHouseholderQR:
            {
                Eigen::FullPivHouseholderQR<MatrixX> decomp = A.fullPivHouseholderQr();
                int rank = decomp.rank();
                if (rank == std::min(A.rows(), A.cols()))
                    CI_LOG_D("rank is " << decomp.rank() << " (of " << std::min(A.rows(), A.cols()) << ")");
                else
                    CI_LOG_W("Matrix is not of full rank, only " << rank << " of " << std::min(A.rows(), A.cols()));
                return decomp.solve(b);
            }
        case DenseLinearSolver::CompleteOrthogonalDecomposition: {return A.completeOrthogonalDecomposition().solve(b); }
        case DenseLinearSolver::LLT: {return A.llt().solve(b); }
        case DenseLinearSolver::LDLT: {return A.ldlt().solve(b); }
        default: throw std::exception("unsupported solver");
        }
    }

    VectorX TimeIntegrator::solveSparse(const SparseMatrixRowMajor & A, const VectorX & b, const VectorX& guess) const
    {
        //TODO: do I want to support custom number of iterations and thesholds?
        switch (sparseSolver)
        {
        case SparseLinearSolver::ConjugateGradient:
            {
                Eigen::ConjugateGradient<SparseMatrixRowMajor, Eigen::Lower | Eigen::Upper> cg;
                if (sparseSolverIterations > 0) cg.setMaxIterations(sparseSolverIterations);
                if (sparseSolverTolerance > 0) cg.setTolerance(sparseSolverTolerance);
                cg.compute(A);
                if (guess.size() > 0)
                    return cg.solveWithGuess(b, guess);
                else
                    return cg.solve(b);
            }
        case SparseLinearSolver::BiCGSTAB:
        {
            Eigen::BiCGSTAB<SparseMatrixRowMajor> cg;
            if (sparseSolverIterations > 0) cg.setMaxIterations(sparseSolverIterations);
            if (sparseSolverTolerance > 0) cg.setTolerance(sparseSolverTolerance);
            cg.compute(A);
            if (guess.size() > 0)
                return cg.solveWithGuess(b, guess);
            else
                return cg.solve(b);
        }
        case SparseLinearSolver::LU:
        {
            CI_LOG_E("Sparse LU does only work with column-major matrices");
            return VectorX::Zero(b.size());
        }
        default: throw std::exception("unsupported solver");
        }
    }

    VectorX TimeIntegrator::solveSparse(const SparseMatrixColumnMajor & A, const VectorX & b, const VectorX& guess) const
    {
        //TODO: do I want to support custom number of iterations and thesholds?
        switch (sparseSolver)
        {
        case SparseLinearSolver::ConjugateGradient:
        {
            Eigen::ConjugateGradient<SparseMatrixColumnMajor, Eigen::Lower | Eigen::Upper> cg;
            if (sparseSolverIterations > 0) cg.setMaxIterations(sparseSolverIterations);
            if (sparseSolverTolerance > 0) cg.setTolerance(sparseSolverTolerance);
            cg.compute(A);
            if (guess.size() > 0)
                return cg.solveWithGuess(b, guess);
            else
                return cg.solve(b);
        }
        case SparseLinearSolver::BiCGSTAB:
        {
            Eigen::BiCGSTAB<SparseMatrixColumnMajor> cg;
            if (sparseSolverIterations > 0) cg.setMaxIterations(sparseSolverIterations);
            if (sparseSolverTolerance > 0) cg.setTolerance(sparseSolverTolerance);
            cg.compute(A);
            if (guess.size() > 0)
                return cg.solveWithGuess(b, guess);
            else
                return cg.solve(b);
        }
        case SparseLinearSolver::LU:
        {
            Eigen::SparseLU<SparseMatrixColumnMajor, Eigen::COLAMDOrdering<int>> lu;
            lu.analyzePattern(A);
            lu.factorize(A);
            return lu.solve(b);
        }
        default: throw std::exception("unsupported solver");
        }
    }

    std::shared_ptr<TimeIntegrator> TimeIntegrator::createIntegrator(Integrator type, Eigen::Index dof)
    {
        switch (type)
        {
        case Integrator::Newmark1: return std::make_shared<TimeIntegrator_Newmark1>(dof);
        case Integrator::Newmark2: return std::make_shared<TimeIntegrator_Newmark2>(dof);
        case Integrator::ExplicitCentralDifferences: return std::make_shared<TimeIntegrator_ExplicitCentralDifferences>(dof);
        case Integrator::ImplicitLinearAcceleartion: return std::make_shared<TimeIntegrator_ImplicitLinearAcceleration>(dof);
        case Integrator::Newmark3: return std::make_shared<TimeIntegrator_Newmark3>(dof);
        case Integrator::HHTAlpha: return std::make_shared<TimeIntegrator_HHTalpha>(dof);
        default: throw std::exception("unknown integrator");
        }
    }
}


/*

copyright (c) 2021, adrien blassiau and corentin juvigny

permission to use, copy, modify, and/or distribute this software
for any purpose with or without fee is hereby granted, provided
that the above copyright notice and this permission notice appear
in all copies.

the software is provided "as is" and the author disclaims all
warranties with regard to this software including all implied
warranties of merchantability and fitness. in no event shall the
author be liable for any special, direct, indirect, or
consequential damages or any damages whatsoever resulting from
loss of use, data or profits, whether in an action of contract,
negligence or other tortious action, arising out of or in
connection with the use or performance of this software.

*/

/** @file mip.cpp
 *
 * @brief Mixed integer algorithm to construct the optimal solution of the problem
 *
 */

#define TBB_SUPPRESS_DEPRECATED_MESSAGES 1

#include <boost/config.hpp>
#include <boost/graph/adjacency_list.hpp>
#include <boost/graph/connected_components.hpp>
#include <algorithm>
#include <mip.hpp>
#include <utility>
#include <ilcplex/ilocplex.h>

#define EPSILON 10e-10

typedef std::map<Node<2>*,std::pair<IloNum,long>,Node<2>::NodeCmp> IloNumIndexMap;
typedef std::map<Node<2>*,IloNumVar,Node<2>::NodeCmp> IloNumVarMap;
typedef std::map<Node<2>*,IloNum,Node<2>::NodeCmp> IloNumMap;
typedef boost::adjacency_list < boost::vecS, boost::vecS, boost::undirectedS > graph_t;
typedef boost::graph_traits < graph_t >::vertex_descriptor vertex_descriptor;

ILOLAZYCONSTRAINTCALLBACK1( ValidityCallback
                          , const IloNumVarMap&
                          , X )
{
   IloNumIndexMap Xr;
   size_t index = 0;
#if 0
   for (const auto &[name,value] : X)
      Xr[name] = std::make_pair(getValue(value),index++);
#else
   for (const auto &[node,value] : X) {
      IloInt x = getValue(value);
      if ( abs(1 - x) <= EPSILON )
         Xr[node] = std::make_pair(x,index++);
      else
         Xr[node] = std::make_pair(x,-1);
   }
#endif

   graph_t graph;
#if 0
   for (struct {size_t i; IloNumIndexMap::const_iterator it;} ist = {0, Xr.cbegin()}; ist.it != Xr.cend(); ist.it++, ist.i++) {
      for (struct {size_t j; IloNumIndexMap::const_iterator jt;} jst = {ist.i, ist.it}; jst.jt != Xr.cend(); jst.jt++, jst.j++) {
         if ( abs(ist.it->second.first - 1) <= 10e-10
            && (std::find( ist.it->first->communication_queue().cbegin()
                         , ist.it->first->communication_queue().cend()
                         , jst.jt->first ) != ist.it->first->communication_queue().end()) ) {
            boost::add_edge(ist.i,jst.j,graph);
         }
      }
   }
#endif
#if 1
   for (const auto &[node,pair] : Xr) {
      const auto [value1,index1] = pair;
      if ( abs(value1 - 1) <= EPSILON ) {
         for (const auto &neighbour : node->communication_queue()) {
            const auto [value2,index2] = Xr[neighbour];
            if ( abs(value2 - 1) <= EPSILON )
               boost::add_edge(index1,index2,graph);
         }
      }
   }
#endif
#if 0
   std::vector<Node<2>*> nodes(X.size());
   std::transform( Xr.cbegin()
                 , Xr.cend()
                 , nodes.begin()
                 , [](const auto &node) { return node.first; });
   for (size_t i = 0; i < nodes.size(); i++) {
      for (size_t j = i+1; j < nodes.size(); j++) {
         if ( Xr[nodes[i]].value == 1
            && (std::find( nodes[i]->communication_queue().cbegin()
                         , nodes[i]->communication_queue().cend()
                         , nodes[j] ) != nodes[i]->communication_queue().end()) ) {
            boost::add_edge(i,j,graph);
         }
      }
   }

#endif

   std::vector<size_t> component(boost::num_vertices(graph));
   size_t num = boost::connected_components(graph,&component[0]);

   if ( num > 1 ) { // Graph is disconnected
      IloExpr e(getEnv());
      IloInt nbr = 0;
      for (const auto &[node,tuple] : Xr) {
         if ( tuple.first == 1 ) {
            e += X.find(node)->second;
            nbr++;
         } else {
            //e -= X.find(node)->second;
         }
      }
      addLocal(e <= nbr-1);
      e.end();
   }
}

void mip_resolution_2D(Grid<2> &g)
{
   IloEnv env;

   IloModel model(env);

   // Decision var
   IloNumVarMap X;
   for (const auto &s : g.nodes()) {
      std::stringstream ss;
      ss << "x_" << s->name();
      X[s.get()] = IloNumVar(env,0,1,ILOBOOL,ss.str().c_str());
   }

   // Objective
   IloExpr obj(env);
   for (auto &x_s : X)
      obj += x_s.second;
   model.add(IloMinimize(env,obj,"obj"));
   obj.end();

   // Constraints
   for (const auto &s : g.nodes()) {
      IloExpr e_com(env);
      IloExpr e_capt(env);
      for (const auto &v : s->communication_queue())
         e_com += X[v];
      for (const auto &v : s->capture_queue())
         e_capt += X[v];
      model.add(e_com >= 1);
      model.add(e_capt >= 1);
      e_capt.end();
      e_com.end();
   }
   model.add(X[g.well().get()] == 1);

   IloCplex cplex(model);

   cplex.setParam(IloCplex::Param::Threads,8);
   cplex.setParam(IloCplex::Param::MIP::Strategy::File,1);
   cplex.setParam(IloCplex::Param::MIP::Limits::Solutions,1000);
   cplex.use(ValidityCallback(env,X));
   cplex.solve();

   IloAlgorithm::Status status = cplex.getStatus();
   std::stringstream ss;
   ss << "Solution status:                   " << cplex.getStatus();
   std::string separator_line(ss.str().length(),'-');
   std::cout << separator_line << std::endl;
   std::cout << ss.str() << std::endl;
   IloNumIndexMap Xr;
   graph_t graph;
   size_t index = 0;
   switch ( status ) {
      case IloAlgorithm::Optimal:
         std::cout << "Nodes processed:                   " << cplex.getNnodes() << std::endl;
         std::cout << "Active user cuts/lazy constraints: " << cplex.getNcuts(IloCplex::CutUser) << std::endl;
         std::cout << "Time elapsed:                      " << cplex.getTime() << std::endl;
         std::cout << "Optimal value:                     " << cplex.getObjValue()-1 << std::endl;
         for (const auto &[node,value] : X) {
            IloInt x = cplex.getValue(value);
            if ( abs(1 - x) <= EPSILON )
               Xr[node] = std::make_pair(x,index++);
            else
               Xr[node] = std::make_pair(x,-1);
         }
         for (const auto &v : g.nodes()) {
            if ( Xr[v.get()].first == 1 ) {
               typename Node<2>::Queue queue;
               for (const auto &node : v->communication_queue())
                  if ( Xr[node].first == 1 ) {
                     queue.push_front(node);
                  }
               v->set_new_sensor(queue);
            }
         }
         {
            for (const auto &[node,pair] : Xr) {
               const auto [value1,index1] = pair;
               if ( abs(value1 - 1) <= EPSILON ) {
                  for (const auto &neighbour : node->communication_queue()) {
                     const auto [value2,index2] = Xr[neighbour];
                     if ( abs(value2 - 1) <= EPSILON ) {
                        boost::add_edge(index1,index2,graph);
                        std::cout << "Node "<<*node<<" connected to node "<<*neighbour<< " with indices : " << index1 << " " << index2 << std::endl;
                     }
                  }
               }
            }
            std::vector<size_t> component(boost::num_vertices(graph));
            long num = boost::connected_components(graph,&component[0]);
            std::cout << "Number of connected components " << num << std::endl;
            std::cout << "Number of vertices in the graph : " << boost::num_vertices(graph) << std::endl;

            for (const auto &[node,pair] : Xr) {
               const auto [value1,index1] = pair;
               if ( abs(value1 - 1) <= EPSILON ) {
                  std::cout << "Node " << *node << " in component " << component[index1] << std::endl;
               }
            }

         }
         break;
      default:
         break;
   }
   std::cout << separator_line << std::endl;

   model.end();
   cplex.end();
   env.end();
}


/*
 * poh_color.hpp
 * Author: Aven Bross
 *
 * Implementation of Poh path 3-coloring algorithm for triangulated plane graphs
 * that traces a chordless path along the inside of the outer face.
 */

#ifndef __POH_COLOR_HPP
#define __POH_COLOR_HPP

// STL headers
#include <vector>
#include <queue>
#include <stdexcept>
#include <utility>

// Basic graph headers
#include <boost/graph/properties.hpp>
#include <boost/graph/graph_traits.hpp>
#include <boost/property_map/property_map.hpp>

// Local project headers
#include "incidence_list_helpers.hpp"

/*
 * poh_color_recursive
 *
 * assumptions: There are two disjoint colored paths P=p_0...p_n and Q=q_0...q_m
 *     such that p_0...p_nq_m...q_0 is a cycle, and all vertices on the interior
 *     of the subgrapph bounded by this cycle with at least one neighbor in P
 *     have been marked with an identifying integer.
 *
 * inputs: A weakly triangulated planar graph with vertex indices (predfined
 *     boost property), a valid planar embedding of the graph modeling the boost
 *     PlanarEmbedding concept, a read-write-able vertex property map to store
 *     vertex colors, a read-write-able vertex property map to store integer
 *     marks for BFS, a read-write-able vertex property map to store start and
 *     stop iterators in each incidence list, the vertex u such that p_0uq_0
 *     is a triangle, the integer mark for vertices with neighbors in P, an
 *     unsigned integer count such that all vertex marks assigned so far are
 *     less than count, and finally the color for the path T to be constructed,
 *     the color of the path P, and the color of the path Q.
 *
 * output: The coloring vertex property will contain a valid path 3-coloring of
 *     the subgraph bounded by p_0...p_nq_m...q_0 such that no interior vertex
 *     shares a color with a neighbor in P or Q.
 */

namespace {
    template<
            typename graph_t, typename planar_embedding_t, typename color_map_t,
            typename mark_map_t, typename neighbor_range_map_t,
            typename color_t, typename vertex_t
                = typename boost::graph_traits<graph_t>::vertex_descriptor
        >
    void poh_color_recursive(
            const graph_t & graph, const planar_embedding_t & planar_embedding,
            color_map_t & color_map, mark_map_t & mark_map,
            neighbor_range_map_t & neighbor_range_map, vertex_t u,
            int face_mark, int count,
            color_t t_color, color_t p_color, color_t q_color
        )
    {
        // t_i will track the last vertex of T, and w the first interior vertex
        vertex_t t_i = u, w = u;

        // Initialize remaining variables for the loop constructing T
        auto edge_iter = neighbor_range_map[u].first;
        int below_t_mark = count++;
        color_map[u] = t_color;

        // If u is the only remaining interior vertex we are are done
        if(neighbor_range_map[u].first == neighbor_range_map[u].second){
            return;
        }

        // Construct the path T one vertex at a time, with current vertex t_i
        while(edge_iter++ != neighbor_range_map[t_i].second)
        {
            // If we hit the end of the incidence list, wrap to the start
            if(edge_iter == planar_embedding[t_i].end())
                edge_iter = planar_embedding[t_i].begin();

            vertex_t n = get_incident_vertex(t_i, *edge_iter, graph);

            // Note if the vertex n may be added to T
            bool continue_path = (mark_map[n] == face_mark
                && color_map[n] != p_color && color_map[n] != q_color);

            // If we will add n to T or we hit the end, color vertices above T
            if(continue_path || color_map[n] == p_color) {
                if(continue_path) {
                    color_map[n] = t_color;
                }

                vertex_t l = n;

                // Loop through neighbors of t_i that lie above T
                while(edge_iter++ != neighbor_range_map[t_i].second) {
                    // If we hit the end of the list, wrap to the start
                    if(edge_iter == planar_embedding[t_i].end())
                        edge_iter = planar_embedding[t_i].begin();

                    vertex_t v = get_incident_vertex(t_i, *edge_iter, graph);

                    // If v is in P and l is uncolored, make recursive call
                    if(color_map[v] == p_color)
                    {
                        if(color_map[l] != p_color && color_map[l] != q_color
                            && color_map[l] != t_color)
                        {
                            poh_color_recursive(
                                    graph, planar_embedding, color_map,
                                    mark_map, neighbor_range_map, l, face_mark,
                                    count, q_color, p_color, t_color
                                );
                        }
                    }

                    l = v;
                }

                // If we are adding n to T, setup next loop with t_i = n
                if(continue_path) {
                    auto n_back_iter = find_edge_iterator_restricted(
                            n, t_i, neighbor_range_map[n].first,
                            neighbor_range_map[n].second, planar_embedding,
                            graph
                        );

                    neighbor_range_map[n].first = n_back_iter;
                    edge_iter = n_back_iter;
                    t_i = n;
                }
                // Otherwise we have completed T and we are done
                else {
                    return;
                }
            }
            // If n is in Q we have an edge T to Q and may split along it
            else if(color_map[n] == q_color) {
                // Color the left cycle if it has uncolored vertices
                if(w != u) {
                    poh_color_recursive(
                            graph, planar_embedding, color_map, mark_map,
                            neighbor_range_map, w, below_t_mark, count,
                            p_color, t_color, q_color
                        );

                    u = t_i;
                    w = t_i;
                }
            }
            // If n is interior and has not been visited, mark and initialize it
            else if(mark_map[n] != below_t_mark) {
                auto back_iter = find_edge_iterator(
                        n, t_i, planar_embedding, graph
                    );

                // If this is the first interior vertex hit, save it as w
                if(w == u) {
                    w = n;
                }

                // Increment back_iter counterclockwise, wrapping in list
                if(++back_iter == planar_embedding[n].end()) {
                    back_iter = planar_embedding[n].begin();
                }

                initialize_neighbor_range(
                        n, back_iter, neighbor_range_map, planar_embedding
                    );
                mark_map[n] = below_t_mark;
            }
        }
    }
}

/*
 * poh_color
 *
 * inputs: A weakly triangulated planar graph with vertex indices (predfined
 *     boost property), a valid planar embedding of the graph modeling the boost
 *     PlanarEmbedding concept, a read-write-able vertex property map to store
 *     vertex colors, pairs of bidirectional iterators for lists of vertices
 *     of two disjoint colored paths P=p_0...p_n and Q=q_0...q_m such that
 *     p_0...p_nq_m...q_0 is a cycle, and three colors for the path 3-coloring.
 *
 * output: The coloring vertex property will contain a valid path 3-coloring of
 *     the subgraph bounded by p_0...p_nq_m...q_0 such that vertices in P are
 *     recieve the first color, vertices in Q the second color, and no interior
 *     vertex shares a color with a neighbor in P or Q.
 */

template<
        typename graph_t,
        typename planar_embedding_t,
        typename color_map_t,
        typename mark_map_t,
        typename neighbor_range_map_t,
        typename vertex_iterator_t,
        typename color_t
    >
void poh_color(
        const graph_t & graph,
        const planar_embedding_t & planar_embedding,
        vertex_iterator_t p_begin, vertex_iterator_t p_end,
        vertex_iterator_t q_begin, vertex_iterator_t q_end,
        color_t c_0, color_t c_1, color_t c_2,
        neighbor_range_map_t & neighbor_range_map,
        mark_map_t & mark_map,
        color_map_t & color_map
    )
{
    // Type definitions
    typedef typename boost::graph_traits<graph_t>::vertex_descriptor vertex_t;

    // Intitialize neighbor ranges for vertices in the path P
    vertex_t l = *q_begin;
    for(vertex_iterator_t p_iter = p_begin; p_iter != p_end; ++p_iter) {
        vertex_t v = *p_iter;

        auto back_iter = find_edge_iterator(v, l, planar_embedding, graph);
        initialize_neighbor_range(
                v, back_iter, neighbor_range_map, planar_embedding
            );

        l = v;
    }

    // Color the path Q
    for(vertex_iterator_t q_iter = q_begin; q_iter != q_end; ++q_iter) {
        vertex_t v = *q_iter;

        color_map[v] = c_1;
    }

    // Construct the path 3-coloring
    poh_color_recursive(
            graph, planar_embedding, color_map, mark_map,
            neighbor_range_map, *p_begin, 1, 2,
            c_0, c_2, c_1
        );
}

/*
 * A wrapper that automatically constructs fast property maps for the mark_map
 * and neighbor_range_map, but requires that graph_t is some definition of
 * boost::adjacency_list.
 */

template<
        typename graph_t,
        typename planar_embedding_t,
        typename color_map_t,
        typename vertex_iterator_t,
        typename color_t
    >
void poh_color(
        const graph_t & graph,
        const planar_embedding_t & planar_embedding,
        vertex_iterator_t p_begin, vertex_iterator_t p_end,
        vertex_iterator_t q_begin, vertex_iterator_t q_end,
        color_t c_0, color_t c_1, color_t c_2,
        color_map_t & color_map
    )
{
    // Vertex property map type to store vertex marks
    typedef boost::iterator_property_map<
            std::vector<int>::iterator,
            typename boost::property_map<
                    graph_t, boost::vertex_index_t
                >::const_type
        > integer_property_map_t;

    // Vertex property map type for the neighbor ranges of planar_embedding_t
    typedef typename boost::property_traits<planar_embedding_t>::value_type
            ::const_iterator embedding_iterator_t;
    typedef typename std::vector<
            std::pair<embedding_iterator_t, embedding_iterator_t>
        > neighbor_range_storage_t;
    typedef boost::iterator_property_map<
            typename neighbor_range_storage_t::iterator,
            typename boost::property_map<
                    graph_t, boost::vertex_index_t
                >::const_type
        > neighbor_range_map_t;

    // Construct a vertex property map to store vertex marks
    std::vector<int> mark_storage(boost::num_vertices(graph));
    integer_property_map_t mark_map(
            mark_storage.begin(), boost::get(boost::vertex_index, graph)
        );

    // Construct a vertex property map to store neighbor ranges
    neighbor_range_storage_t neighbor_range_storage(num_vertices(graph));
    neighbor_range_map_t neighbor_range_map(
            neighbor_range_storage.begin(),
            boost::get(boost::vertex_index, graph)
        );

    // Construct the path 3-coloring
    poh_color(
            graph,
            planar_embedding,
            p_begin, p_end,
            q_begin, q_end,
            1, 2, 3,
            neighbor_range_map,
            mark_map,
            color_map
        );
}

#endif


#pragma once

#include <boost/functional/hash.hpp>
#include <boost/math/special_functions/binomial.hpp>
#include <boost/math/special_functions/gamma.hpp>
#include <boost/math/special_functions/laguerre.hpp>
#include <stdexcept>

#include "spectral/basis/spectral_function/spectral_function_base.hpp"
#include "spectral/basis/spectral_function/spectral_weight_function.hpp"

namespace boltzmann {

namespace local_ {
struct laguerreKS_id_t
{
 private:
  constexpr const static double FUZZY = 1e6;

 public:
  typedef laguerreKS_id_t id_t;

  /// Default constructor
  laguerreKS_id_t()
      : j(std::nan("nan"))
      , k(std::nan("nan"))
      , fw(std::nan("nan"))
      , idw(std::nan("nan"))
  {
  }

  laguerreKS_id_t(int k_, int j_, double fw_)
      : j(j_)
      , k(k_)
      , fw(fw_)
      , idw(FUZZY * fw_)
  {
  }

  laguerreKS_id_t(const id_t& id)
      : j(id.j)
      , k(id.k)
      , fw(id.fw)
      , idw(id.idw)
  {
  }

  /// order parameter
  int j;
  /// Polynomial degree
  int k;
  /// weight exponent
  double fw;
  /// weight id
  long int idw;

  bool operator<(const id_t& other) const
  {
    return std::tie(j, k, idw) < std::tie(other.j, other.k, other.idw);
    //#warning "fix this!"
    //    return k < other.k;
  }

  // ----------------------------------------------------------------------
  inline bool operator==(const id_t& other) const
  {
    return std::tie(j, k, idw) == std::tie(other.j, other.k, other.idw);
    //#warning "fix this!"
    //    return k == other.k;
  }

  // ----------------------------------------------------------------------
  friend std::ostream& operator<<(std::ostream& stream, const id_t& x)
  {
    stream << x.to_string();
    return stream;
  }

  // ----------------------------------------------------------------------
  std::string to_string() const
  {
    return "(fw_" + boost::lexical_cast<std::string>(fw) + ", k_" +
           boost::lexical_cast<std::string>(k) + ", j_" + boost::lexical_cast<std::string>(j) +
           ") ";
  }

  // ----------------------------------------------------------------------
  inline std::tuple<int, int, long int> key() const { return std::make_tuple(j, k, idw); }
};
}  // end namespace local_
}  // end namespace boltzmann

namespace std {
// hash functions for id's
template <>
class hash<boltzmann::local_::laguerreKS_id_t>
{
 public:
  size_t operator()(const boltzmann::local_::laguerreKS_id_t& id) const
  {
    std::size_t current = std::hash<double>()(id.fw);
    boost::hash_combine(current, std::hash<int>()(id.k));
    boost::hash_combine(current, std::hash<int>()(id.j));
    return current;
  }
};
}  // end namespace std

namespace boltzmann {

/**
 * @brief Laguerre Radial Polynomial used by Kitzler & Schoeberl
 *
 * Normalization & Orthogonality
 * =============================
 *
 * For \f$ k_1 \equiv k_2 \operatorname{mod} 2 \f$, \f$k_1\f$ even:
 * \f[
 *   \int_0^{\infty} r^{2j}  r^{2j} L^{(2j)}_{\frac{k_1}{2}-j}(r^2)
 * L^{(2j)}_{\frac{k_2}{2}-j} (r^2) e^{-r^2}\; r \operatorname{d} r= \frac{1}{2} \delta_{k_1, k_2}
 * \f]
 *
 *  For \f$ k_1 \equiv k_2 \operatorname{mod} 2 \f$, \f$k_1\f$ odd:
 * \f[
 *   \int_0^{\infty} r^{2j+1}  r^{2j+1} L^{(2j+1)}_{\frac{k_1-1}{2}-j}(r^2)
 * L^{(2j+1)}_{\frac{k_2-1}{2}-j} (r^2) e^{-r^2}\; r \operatorname{d} r= \frac{1}{2} \delta_{k_1,
 * k_2}
 * \f]
 *
 */
class LaguerreKS : public weighted<LaguerreKS, true>,
                   public local_::index_policy<local_::laguerreKS_id_t>
{
 public:
  typedef double numeric_t;

 public:
  /**
   * @brief Basis functions of the form \f$ r^{2*j + k\mod2} L_{(k-k\mod2)/2-j}^{2*j + k\mod2}(r^2)
   * \f$
   *
   *
   * @param k  polynomial degree
   * @param j  angular parameter
   * @param w  weight
   *
   * @return
   */
  explicit LaguerreKS(int k, int j, double w = 0.5);
  explicit LaguerreKS(const id_t& id)
      : id_(id)
  {
  }
  LaguerreKS(){};

  /// evaluate polynomial part
  numeric_t evaluate(double r) const;

  /// evaluate weight
  numeric_t weight(double r) const;

  /// return weight
  numeric_t w() const { return id_.fw; }

  const id_t& get_id() const { return id_; }

  unsigned int get_order() const { return 2 * id_.j + id_.k % 2; }
  unsigned int get_degree() const { return id_.k / 2 - id_.j; }

 private:
  id_t id_;
};

// ----------------------------------------------------------------------
inline LaguerreKS::LaguerreKS(int k, int j, double w)
    : id_(k, j, w)
{
  if (j > k / 2) {
    std::runtime_error("invalid combination of parameters in LaguerreKS");
  }
}

// ----------------------------------------------------------------------
inline typename LaguerreKS::numeric_t
LaguerreKS::evaluate(double r) const
{
  int alpha = 2 * id_.j + (id_.k % 2);
  int n = id_.k / 2 - id_.j;

  // normalization factor
  double a = 1;
  for (int i = n + alpha; i > n; --i) {
    a *= i;
  }

  return std::pow(r, 2 * id_.j + (id_.k % 2)) * boost::math::laguerre(n, alpha, r * r) /
         std::sqrt(a);
}

// ----------------------------------------------------------------------
inline typename LaguerreKS::numeric_t
LaguerreKS::weight(double r) const
{
  return std::exp(-r * r * id_.fw);
}

}  // end namespace boltzmann


#pragma once
#include <random>
#include <bitset>
#include <fstream>
#include <ios>
#include <string>
#include <cassert>
#include <Eigen/Eigen>

#include <tbb/tbb.h>

#include <nlohmann/json.hpp>

#include "Utilities/type_traits.hpp"
#include "Utilities/Utility.hpp"
#include "Serializers/SerializeEigen.hpp"

namespace yannq
{
//! \ingroup Machines
//! RBM machine
template<typename T>
class RBM
{
	static_assert(std::is_floating_point<T>::value || is_complex_type<T>::value, "T must be floating or complex");
public:
	using Scalar = T;
	using RealScalar = typename yannq::remove_complex<T>::type;

	using Matrix = Eigen::Matrix<T, Eigen::Dynamic, Eigen::Dynamic>;
	using Vector = Eigen::Matrix<T, Eigen::Dynamic, 1>;
	using VectorRef = Eigen::Ref<Vector>;
	using VectorConstRef = Eigen::Ref<const Vector>;

	using RealVector = Eigen::Matrix<RealScalar, Eigen::Dynamic, 1>;

	using DataT = std::tuple<Eigen::VectorXi, Vector>;

private:
	uint32_t n_; //# of qubits
	uint32_t m_; //# of hidden units

	bool useBias_;

	Matrix W_; //W should be m by n
	Vector a_; //a is length n
	Vector b_; //b is length m

public:
	RBM(uint32_t n, uint32_t m, bool useBias = true) noexcept
		: n_(n), m_(m), useBias_(useBias), W_(m, n), a_(n), b_(m)
	{
		a_.setZero();
		b_.setZero();
		W_.setZero();
	}

	RBM() noexcept
		: useBias_{true}
	{
	}

	template<typename U, std::enable_if_t<std::is_convertible_v<U, T> && !std::is_same_v<U, T>, int> = 0>
	RBM(const RBM<U>& rhs)
		: n_{rhs.getN()}, m_{rhs.getM()}, useBias_{rhs.useBias()}
	{
		W_ = rhs.getW().template cast<T>();
		a_ = rhs.getA().template cast<T>();
		b_ = rhs.getB().template cast<T>();
	}

	RBM(const RBM& rhs) /* noexcept */ = default;
	RBM(RBM&& rhs) /* noexcept */ = default;

	template<typename U, std::enable_if_t<std::is_convertible_v<U, T> && !std::is_same_v<U, T>, int> = 0>
	RBM& operator=(const RBM<U>& rhs)
	{
		n_ = rhs.n_;
		m_ = rhs.m_;
		useBias_ = rhs.useBias_;

		W_ = rhs.W_.template cast<T>();
		a_ = rhs.a_.template cast<T>();
		b_ = rhs.b_.template cast<T>();

		return *this;
	}

	RBM& operator=(const RBM& rhs) /* noexcept */ = default;
	RBM& operator=(RBM&& rhs) /* noexcept */ = default;

	template<typename U, std::enable_if_t<std::is_convertible_v<T, U>, int> = 0>
	RBM<U> cast() const
	{
		RBM<U> res(n_, m_, useBias_);
		res.setA(a_.template cast<U>());
		res.setB(b_.template cast<U>());
		res.setW(W_.template cast<U>());
		return res;
	}

	nlohmann::json desc() const
	{
		return nlohmann::json
		{
			{"name", "RBM"},
			{"useBias", useBias_},
			{"n", n_},
			{"m", m_}
		};
	}

	inline uint32_t getN() const
	{
		return n_;
	}
	inline uint32_t getM() const
	{
		return m_;
	}

	inline uint32_t getDim() const
	{
		if(useBias_)
			return n_*m_ + n_ + m_;
		else
			return n_*m_;
	}

	inline bool useBias() const
	{
		return useBias_;
	}

	inline Vector calcTheta(const Eigen::VectorXi& sigma) const
	{
		Vector s = sigma.cast<T>();
		return W_*s + b_;
	}
	inline Vector calcGamma(const Eigen::VectorXi& hidden) const
	{
		Vector h = hidden.cast<T>();
		return W_.transpose()*h + a_;
	}

	void setUseBias(bool newBias)
	{
		useBias_ = newBias;
	}

	void resize(uint32_t n, uint32_t m)
	{
		n_ = n;
		m_ = m;

		a_.resize(n);
		b_.resize(m);
		W_.resize(m,n);

		if(!useBias_)
		{
			a_.setZero();
			b_.setZero();
		}
	}

	void conservativeResize(uint32_t newM)
	{
		Vector newB = Vector::Zero(newM);
		newB.head(m_) = b_;

		Matrix newW = Matrix::Zero(newM, n_);
		newW.topRows(m_) = W_;

		m_ = newM;
		b_ = std::move(newB);
		W_ = std::move(newW);
	}

	void setW(const Eigen::Ref<const Matrix>& m)
	{
		assert(m.rows() == W_.rows() && m.cols() == W_.cols());
		W_ = m;
	}

	void setA(const VectorConstRef& A)
	{
		assert(A.size() == a_.size());
		if(!useBias_)
			return ;
		a_ = A;
	}

	void setB(const VectorConstRef& B)
	{
		assert(B.size() == b_.size());
		if(!useBias_)
			return ;
		b_ = B;
	}

	inline const T& W(uint32_t j, uint32_t i) const
	{
		return W_.coeff(j,i);
	}
	inline const T& A(uint32_t i) const
	{
		return a_.coeff(i);
	}
	inline const T& B(uint32_t j) const
	{
		return b_.coeff(j);
	}

	inline T& W(uint32_t j, uint32_t i)
	{
		return W_.coeffRef(j,i);
	}
	inline T& A(uint32_t i)
	{
		return a_.coeffRef(i);
	}
	inline T& B(uint32_t j)
	{
		return b_.coeffRef(j);
	}

	const Matrix& getW() const & { return W_; }
	Matrix getW() && { return std::move(W_); }

	const Vector& getA() const & { return a_; }
	Vector getA() && { return std::move(a_); }

	const Vector& getB() const & { return b_; }
	Vector getB() && { return std::move(b_); }

	//! update Bias A by adding v
	void updateA(const VectorConstRef& v)
	{
		assert(useBias_);
		a_ += v;
	}
	//! update Bias B by adding v
	void updateB(const VectorConstRef& v)
	{
		assert(useBias_);
		b_ += v;
	}
	//! update the weight W by adding m
	void updateW(const Eigen::Ref<const Matrix>& m)
	{
		W_ += m;
	}

	//! update all parameters.
	void updateParams(const VectorConstRef& m)
	{
		assert(m.size() == getDim());
		W_ += Eigen::Map<const Matrix>(m.data(), m_, n_);
		if(!useBias_)
			return ;
		a_ += Eigen::Map<const Vector>(m.data() + m_*n_, n_);
		b_ += Eigen::Map<const Vector>(m.data() + m_*n_ + n_, m_);
	}

	Vector getParams() const
	{
		Vector res(getDim());
		res.head(n_*m_) = Eigen::Map<const Vector>(W_.data(), W_.size());
		if(!useBias_)
			return res;

		res.segment(n_*m_, n_) = a_;
		res.segment(n_*m_ + n_, m_) = b_;
		return res;
	}

	void setParams(const VectorConstRef& r)
	{
		assert(r.size() == getDim());
		Eigen::Map<Vector>(W_.data(), W_.size()) = r.head(n_*m_);
		if(!useBias_)
			return ;
		a_ = r.segment(n_*m_, n_);
		b_ = r.segment(n_*m_ + n_, m_);
	}

	bool hasNaN() const
	{
		return a_.hasNaN() || b_.hasNaN() || W_.hasNaN();
	}

	/* When T is real type */
	template <typename RandomEngine, class U=T,
            	std::enable_if_t < !is_complex_type<U>::value, int > = 0 >
	void initializeRandom(RandomEngine& re, T sigma = 1e-3)
	{
		std::normal_distribution<T> nd{0, sigma};
		if(useBias_)
		{
			for(uint32_t i = 0u; i < n_; i++)
			{
				a_.coeffRef(i) = nd(re);
			}
			for(uint32_t i = 0u; i < m_; i++)
			{
				b_.coeffRef(i) = nd(re);
			}
		}
		for(uint32_t j = 0u; j < n_; j++)
		{
			for(uint32_t i = 0u; i < m_; i++)
			{
				W_.coeffRef(i, j) = nd(re);
			}
		}
	}

	/* When T is complex type */
	template <typename RandomEngine, class U=T,
               std::enable_if_t < is_complex_type<U>::value, int > = 0 >
	void initializeRandom(RandomEngine& re, typename remove_complex<T>::type sigma = 1e-3)
	{
		std::normal_distribution<typename remove_complex<T>::type> nd{0, sigma};

		if(useBias_)
		{
			for(uint32_t i = 0u; i < n_; i++)
			{
				a_.coeffRef(i) = T{nd(re), nd(re)};
			}
			for(uint32_t i = 0u; i < m_; i++)
			{
				b_.coeffRef(i) = T{nd(re), nd(re)};
			}
		}
		for(uint32_t j = 0; j < n_; j++)
		{
			for(uint32_t i = 0u; i < m_; i++)
			{
				W_.coeffRef(i, j) = T{nd(re), nd(re)};
			}
		}
	}

	bool operator==(const RBM<T>& rhs) const
	{
		if(n_ != rhs.n_ || m_ != rhs.m_)
			return false;
		bool equalW = (W_ == rhs.W_);
		if(useBias_)
			return equalW && (a_ == rhs.a_) && (b_ == rhs.b_);
		else
			return equalW;
	}

	std::tuple<Eigen::VectorXi, Vector> makeData(const Eigen::VectorXi& sigma) const
	{
		return std::make_tuple(sigma, calcTheta(sigma));
	}

	T logCoeff(const std::tuple<Eigen::VectorXi, Vector>& t) const
	{
		using std::cosh;

		Vector ss = std::get<0>(t).template cast<T>();
		T s = a_.transpose()*ss;
		for(uint32_t j = 0u; j < m_; j++)
		{
			s += logCosh(std::get<1>(t).coeff(j));
		}
		return s;
	}

	T coeff(const std::tuple<Eigen::VectorXi, Vector>& t) const
	{
		using std::cosh;

		Vector ss = std::get<0>(t).template cast<T>();
		T s = a_.transpose()*ss;
		T p = exp(s) * std::get<1>(t).array().cosh().prod();
		return p;
	}

	Vector logDeriv(const std::tuple<Eigen::VectorXi, Vector>& t) const
	{
		Vector res(getDim());

		Vector tanhs = std::get<1>(t).array().tanh();
		Vector sigma = std::get<0>(t).template cast<Scalar>();

		for(uint32_t i = 0u; i < n_; i++)
		{
			res.segment(i*m_, m_) = sigma(i)*tanhs;
		}
		if(!useBias_)
			return res;
		res.segment(n_*m_, n_) = sigma;
		res.segment(n_*m_ + n_, m_) = tanhs;
		return res;
	}
};

template<typename T>
typename RBM<T>::Vector getPsi(const RBM<T>& qs, bool normalize)
{
	const uint32_t n = qs.getN();
	typename RBM<T>::Vector psi(1u<<n);
	tbb::parallel_for(0u, (1u << n),
		[n, &qs, &psi](uint32_t idx)
	{
		auto s = toSigma(n, idx);
		psi(idx) = qs.coeff(qs.makeData(s));
	});
	if(normalize)
		psi.normalize();
	return psi;
}

template<typename T, typename Iterable> //Iterable must be random access iterable
typename RBM<T>::Vector getPsi(const RBM<T>& qs, Iterable&& basis, bool normalize)
{
	const uint32_t n = qs.getN();
	typename RBM<T>::Vector psi(basis.size());

	tbb::parallel_for(std::size_t(0u), basis.size(),
		[n, &qs, &psi, &basis](std::size_t idx)
	{
		auto s = toSigma(n, basis[idx]);
		psi(idx) = qs.coeff(qs.makeData(s));
	});
	if(normalize)
		psi.normalize();
	return psi;
}
template<typename T>
typename RBM<T>::RealVector getProbs(const RBM<T>& qs, bool normalize)
{
	return getPsi(qs, normalize).cwiseAbs2();
}

template<typename T, typename Iterable> //Iterable must be random access iterable
typename RBM<T>::RealVector getProbs(const RBM<T>& qs, Iterable&& basis, bool normalize)
{
	return getPsi(qs, std::forward<Iterable>(basis), normalize).cwiseAbs2();
}
}//namespace yannq


// Copyright Aleksey Gurtovoy 2001-2004
//
// Distributed under the Boost Software License, Version 1.0.
// (See accompanying file LICENSE_1_0.txt or copy at
// http://www.boost.org/LICENSE_1_0.txt)
//
// See http://www.boost.org/libs/mpl for documentation.

// $Id: integer.cpp 49268 2008-10-11 06:26:17Z agurtovoy $
// $Date: 2008-10-11 02:26:17 -0400 (Sat, 11 Oct 2008) $
// $Revision: 49268 $

#include <boost/mpl/multiplies.hpp>
#include <boost/mpl/list.hpp>
#include <boost/mpl/lower_bound.hpp>
#include <boost/mpl/transform_view.hpp>
#include <boost/mpl/sizeof.hpp>
#include <boost/mpl/int.hpp>
#include <boost/mpl/identity.hpp>
#include <boost/mpl/base.hpp>
#include <boost/mpl/eval_if.hpp>
#include <boost/mpl/deref.hpp>
#include <boost/mpl/begin_end.hpp>
#include <boost/mpl/assert.hpp>

#include <boost/type_traits/is_same.hpp>

namespace mpl = boost::mpl;
using namespace mpl::placeholders;

template< int bit_size >
class big_int
{
    // ...
};

template< int bit_size >
struct integer
{
    typedef mpl::list<char,short,int,long> builtins_;
    typedef typename mpl::base< typename mpl::lower_bound<
          mpl::transform_view< builtins_
            , mpl::multiplies< mpl::sizeof_<_1>, mpl::int_<8> >
            >
        , mpl::int_<bit_size>
        >::type >::type iter_;

    typedef typename mpl::end<builtins_>::type last_;
    typedef typename mpl::eval_if<
          boost::is_same<iter_,last_>
        , mpl::identity< big_int<bit_size> >
        , mpl::deref<iter_>
        >::type type;
};

typedef integer<1>::type int1;
typedef integer<5>::type int5;
typedef integer<15>::type int15;
typedef integer<32>::type int32;
typedef integer<100>::type int100;

BOOST_MPL_ASSERT(( boost::is_same< int1, char > ));
BOOST_MPL_ASSERT(( boost::is_same< int5, char > ));
BOOST_MPL_ASSERT(( boost::is_same< int15, short > ));
BOOST_MPL_ASSERT(( boost::is_same< int32, int > ));
BOOST_MPL_ASSERT(( boost::is_same< int100, big_int<100> > ));


/**
 * @file print_helper.hpp
 * @author francois.hamonic@gmail.com
 * @brief
 * @version 0.1
 * @date 2021-07-19
 */
#ifndef PRINT_HELPER_HPP
#define PRINT_HELPER_HPP

#include <math.h>
#include <filesystem>

#include <boost/range/adaptors.hpp>
#include <boost/range/algorithm.hpp>

#include <fmt/os.h>

#include "landscape/decored_landscape.hpp"
#include "landscape/mutable_landscape.hpp"

#include "Eigen/Dense"
#include "algorithms/multiplicative_dijkstra.hpp"
#include "lemon/dijkstra.h"

#include "fast-cpp-csv-parser/csv.h"
#include "solvers/concept/solver.hpp"

#include "lemon/dim2.h"
#include "lemon/graph_to_eps.h"

#include "indices/eca.hpp"

#include "helper.hpp"

namespace Helper {
/**
 * @brief Finds the maximum node scaling for which no nodes overlaps.
 *
 * Finds the maximum scaling coeficient of nodes radius for which their is no
 * overlaps in their graphical representation. Runs in $O(n^2)$, where $n$ is
 * the number of nodes.
 *
 * @tparam GR
 * @tparam QM
 * @tparam DM
 * @tparam CM
 * @param landscape
 * @return double
 */
template <typename LS>
std::pair<double, double> findNodeScale(const LS & landscape) {
    using Graph = typename LS::Graph;
    const Graph & graph = landscape.getNetwork();
    const typename LS::CoordsMap & coordsMap = landscape.getCoordsMap();
    const typename LS::QualityMap & qualityMap = landscape.getQualityMap();

    if(lemon::countNodes(graph) < 2) return std::make_pair(1, 0);

    auto dist = [&](typename Graph::Node u, typename Graph::Node v) {
        const double dx = coordsMap[u].x - coordsMap[v].x;
        const double dy = coordsMap[u].y - coordsMap[v].y;
        return std::sqrt(dx * dx + dy * dy);
    };
    auto radius = [&](typename Graph::Node u) {
        return std::sqrt(qualityMap[u] / (2 * M_PI));
    };
    double r_max = std::numeric_limits<double>::max();
    double min_dist = 0.0;
    for(typename Graph::NodeIt u(graph); u != lemon::INVALID; ++u) {
        for(typename Graph::NodeIt v(graph); v != lemon::INVALID; ++v) {
            if(u == v) continue;
            const double r = dist(u, v) / (radius(u) + radius(v));
            if(r < r_max) {
                r_max = r;
                min_dist = dist(u, v);
            }
        }
    }
    return std::make_pair(r_max, min_dist);
}

template <typename LS>
void printLandscape(const LS & landscape, std::filesystem::path path) {
    const MutableLandscape::Graph & graph = landscape.getNetwork();
    std::string name = path.stem();

    auto radius = [&](double area) { return std::sqrt(area / (2 * M_PI)); };
    const auto & [r_max, min_dist] = findNodeScale(landscape);
    const double a_max = r_max * 2 * radius(minNonZeroQuality(landscape));

    const bool directed = false;

    double node_scale = 0.7;
    double text_scale =
        0.75 / (1 + static_cast<int>(std::log10(lemon::countNodes(graph))));
    double arrow_scale = (1 - node_scale) * 2 / 3;
    double arc_width = node_scale * a_max / 16;

    MutableLandscape::Graph::NodeMap<std::string> node_idsMap(graph, "\7");
    MutableLandscape::Graph::NodeMap<lemon::Color> node_colorsMap(graph,
                                                                  lemon::WHITE);
    MutableLandscape::Graph::NodeMap<double> node_sizesMap(graph,
                                                           arc_width * 0.9);
    for(MutableLandscape::NodeIt v(graph); v != lemon::INVALID; ++v) {
        if(landscape.getQuality(v) == 0) continue;
        node_idsMap[v] = std::to_string(graph.id(v));
        node_sizesMap[v] = radius(landscape.getQuality(v));
    }

    MutableLandscape::Graph::ArcMap<lemon::Color> arcs_colorsMap(graph,
                                                                 lemon::BLACK);
    MutableLandscape::Graph::ArcMap<double> arc_widths(graph, arc_width);

    return lemon::graphToEps(graph, "output/" + name)
        .title(name)
        .coords(landscape.getCoordsMap())
        .autoNodeScale(false)
        .absoluteNodeSizes(true)
        .nodeSizes(node_sizesMap)
        .nodeScale(node_scale * r_max)
        .autoArcWidthScale(false)
        .absoluteArcWidths(true)
        .arcWidths(arc_widths)
        .arcWidthScale(1)
        .drawArrows(directed)
        .arrowLength(arrow_scale * min_dist)
        .arrowWidth(arc_width * 2)
        .nodeTexts(node_idsMap)
        .nodeTextSize(text_scale * a_max)
        .nodeColors(node_colorsMap)
        .arcColors(arcs_colorsMap)
        .enableParallel(directed)
        .parArcDist(2 * arrow_scale * min_dist)
        .border(20)
        .run();
}

template <typename LS>
void printInstance(const LS & landscape, const RestorationPlan<LS> & plan,
                   std::filesystem::path path) {
    using Graph = typename LS::Graph;

    const Graph & graph = landscape.getNetwork();
    std::string name = path.stem();

    auto radius = [&](double area) { return std::sqrt(area / (2 * M_PI)); };
    const auto & [r_max, min_dist] = findNodeScale(landscape);
    const double a_max = r_max * 2 * radius(minNonZeroQuality(landscape));

    const bool directed = false;

    double node_scale = 0.7;
    double text_scale =
        0.75 / (1 + static_cast<int>(std::log10(lemon::countNodes(graph))));
    double arrow_scale = (1 - node_scale) * 2 / 3;
    double arc_width = node_scale * a_max / 16;

    typename Graph::template NodeMap<std::string> node_idsMap(graph, "");
    typename Graph::template NodeMap<lemon::Color> node_colorsMap(graph,
                                                                  lemon::BLACK);
    typename Graph::template NodeMap<double> node_sizesMap(graph,
                                                           arc_width * 0.9);
    for(typename Graph::NodeIt v(graph); v != lemon::INVALID; ++v) {
        if(landscape.getQuality(v) == 0) continue;
        node_idsMap[v] = std::to_string(graph.id(v));
        node_colorsMap[v] = lemon::WHITE;
        node_sizesMap[v] = radius(landscape.getQuality(v));
    }

    typename Graph::template ArcMap<lemon::Color> arcs_colorsMap(graph,
                                                                 lemon::BLACK);
    typename Graph::template ArcMap<double> arc_widths(graph, arc_width);

    for(typename Graph::NodeIt u(graph); u != lemon::INVALID; ++u)
        node_colorsMap[u] = plan[u].empty() ? lemon::BLACK : lemon::RED;

    for(typename Graph::ArcIt a(graph); a != lemon::INVALID; ++a)
        arcs_colorsMap[a] = plan[a].empty() ? lemon::BLACK : lemon::RED;

    return lemon::graphToEps(graph, path)
        .title(name)
        .coords(landscape.getCoordsMap())
        .autoNodeScale(false)
        .absoluteNodeSizes(true)
        .nodeSizes(node_sizesMap)
        .nodeScale(node_scale * r_max)
        .autoArcWidthScale(false)
        .absoluteArcWidths(true)
        .arcWidths(arc_widths)
        .arcWidthScale(1)
        .drawArrows(directed)
        .arrowLength(arrow_scale * min_dist)
        .arrowWidth(arc_width * 2)
        .nodeTexts(node_idsMap)
        .nodeTextSize(text_scale * a_max)
        .nodeColors(node_colorsMap)
        .arcColors(arcs_colorsMap)
        .enableParallel(directed)
        .parArcDist(2 * arrow_scale * min_dist)
        .border(2)
        .run();
}

void printSolution(const MutableLandscape & landscape,
                   const RestorationPlan<MutableLandscape> & plan,
                   std::string name, concepts::Solver & solver, double B,
                   const Solution & solution);

// need to include the binary search tree for y-h , y+h search
std::pair<MutableLandscape::Node, MutableLandscape::Node> neerestNodes(
    const MutableLandscape & landscape);

template <typename Graph>
class GraphToGraphviz {
public:
    using NodePosMap =
        typename Graph::template NodeMap<lemon::dim2::Point<double>>;
    using NodeSizeMap = typename Graph::template NodeMap<double>;
    using NodeColorMap = typename Graph::template NodeMap<int>;
    using ArcSizeMap = typename Graph::template ArcMap<double>;
    using ArcColorMap = typename Graph::template ArcMap<int>;

    using Node = typename Graph::Node;
    using NodeIt = typename Graph::NodeIt;
    using Arc = typename Graph::Arc;
    using ArcIt = typename Graph::ArcIt;

private:
    const Graph & _graph;
    std::filesystem::path _path;
    NodePosMap _nodePos;
    NodeSizeMap _nodeSizes;
    NodeColorMap _nodeColors;
    ArcSizeMap _arcSizes;
    ArcColorMap _arcColors;

    double _pageWidth;
    double _pageHeight;

    double _node_size_scale;
    double _arc_size_scale;

public:
    GraphToGraphviz(const Graph & g, const std::filesystem::path & p)
        : _graph(g)
        , _path(p)
        , _nodePos(g)
        , _nodeSizes(g, 1)
        , _nodeColors(g, 0xffffff)
        , _arcSizes(g, 1)
        , _arcColors(g, 0x000000)
        , _pageWidth(8)
        , _pageHeight(11)
        , _node_size_scale(1)
        , _arc_size_scale(1) {}

    template <typename PM>
    GraphToGraphviz<Graph> & nodePos(const PM & posMap) {
        for(NodeIt u(_graph); u != lemon::INVALID; ++u) _nodePos[u] = posMap[u];
        return *this;
    }
    template <typename SM>
    GraphToGraphviz<Graph> & nodeSizes(const SM & sizeMap) {
        for(NodeIt u(_graph); u != lemon::INVALID; ++u)
            _nodeSizes[u] = sizeMap[u];
        return *this;
    }
    GraphToGraphviz<Graph> & nodeScale(const double scale) {
        _node_size_scale = scale;
        return *this;
    }
    template <typename CM>
    GraphToGraphviz<Graph> & nodeColors(const CM & colorMap) {
        for(NodeIt u(_graph); u != lemon::INVALID; ++u)
            _nodeColors[u] = colorMap[u];
        return *this;
    }
    template <typename AM>
    GraphToGraphviz<Graph> & arcSizes(const AM & sizeMap) {
        for(ArcIt a(_graph); a != lemon::INVALID; ++a)
            _arcSizes[a] = sizeMap[a];
        return *this;
    }
    GraphToGraphviz<Graph> & arcScale(const double scale) {
        _arc_size_scale = scale;
        return *this;
    }
    template <typename CM>
    GraphToGraphviz<Graph> & arcColors(const CM & colorMap) {
        for(ArcIt a(_graph); a != lemon::INVALID; ++a)
            _arcColors[a] = colorMap[a];
        return *this;
    }
    GraphToGraphviz<Graph> & pageSize(double width, double height) {
        _pageWidth = width;
        _pageHeight = height;
        return *this;
    }

    void run() const {
        double min_x, max_x, min_y, max_y;
        min_x = max_x = _nodePos[_graph.nodeFromId(0)].x;
        min_y = max_y = _nodePos[_graph.nodeFromId(0)].y;
        for(NodeIt u(_graph); u != lemon::INVALID; ++u) {
            min_x = std::min(min_x, _nodePos[u].x);
            max_x = std::max(max_x, _nodePos[u].x);
            min_y = std::min(min_y, _nodePos[u].y);
            max_y = std::max(max_y, _nodePos[u].y);
        }
        const double scale = std::min(_pageWidth / (max_x - min_x),
                                      _pageHeight / (max_y - min_y));
        auto scale_x = [&](double x) { return scale * (x - min_x); };
        auto scale_y = [&](double y) { return scale * (y - min_y); };
        auto scale_size = [&](double s) {
            return _node_size_scale * scale * s;
        };

        std::vector<Node> colorSortedNodes;
        for(NodeIt u(_graph); u != lemon::INVALID; ++u)
            colorSortedNodes.push_back(u);
        std::sort(colorSortedNodes.begin(), colorSortedNodes.end(),
                  [&](Node & a, Node & b) {
                      return _nodeColors[a] < _nodeColors[b];
                  });

        std::vector<Arc> colorSortedArcs;
        for(ArcIt a(_graph); a != lemon::INVALID; ++a)
            colorSortedArcs.push_back(a);
        std::sort(
            colorSortedArcs.begin(), colorSortedArcs.end(),
            [&](Arc & a, Arc & b) { return _arcColors[a] < _arcColors[b]; });

        auto dot_file = fmt::output_file(_path.generic_string());
        dot_file.print("digraph {{size=\"{},{}\";\n", _pageWidth, _pageHeight);
        dot_file.print(
            "graph [pad=\"0.2,0.1\" bgcolor=transparent overlap=scale]\n");
        dot_file.print("node [style=filled shape=\"circle\"]\n");
        dot_file.print("edge [style=filled]\n");

        int prev_color = -1;
        for(Node u : colorSortedNodes) {
            if(_nodeColors[u] != prev_color) {
                dot_file.print("node [fillcolor=\"#{:06x}\"]\n",
                               _nodeColors[u]);
                prev_color = _nodeColors[u];
            }
            dot_file.print("{} [width=\"{}\" pos=\"{},{}!\"]\n", _graph.id(u),
                           scale_size(std::sqrt(_nodeSizes[u])), scale_x(_nodePos[u].x),
                           scale_y(_nodePos[u].y));
        }

        prev_color = -1;
        for(Arc a : colorSortedArcs) {
            if(_arcColors[a] != prev_color) {
                dot_file.print("edge [color=\"#{:06x}\"]\n", _arcColors[a]);
                prev_color = _arcColors[a];
            }
            dot_file.print(
                "{} -> {} [penwidth=\"{}\"]\n", _graph.id(_graph.source(a)),
                _graph.id(_graph.target(a)), _arc_size_scale * _arcSizes[a]);
        }

        dot_file.print("}}");
    }
};

template <typename LS>
void printLandscapeGraphviz(const LS & landscape, std::filesystem::path path) {
    using Graph = typename LS::Graph;
    using Node = typename LS::Graph::Node;
    using NodeIt = typename LS::Graph::NodeIt;
    using Arc = typename LS::Graph::Arc;
    using ArcIt = typename LS::Graph::ArcIt;
    using NodeColorMap = typename Graph::template NodeMap<int>;
    using ArcColorMap = typename Graph::template ArcMap<int>;

    const Graph & graph = landscape.getNetwork();

    NodeColorMap nodeColors(graph);
    for(NodeIt u(graph); u != lemon::INVALID; ++u)
        nodeColors[u] = landscape.getQuality(u) > 0 ? 0x50e050 : 0x101010;

    GraphToGraphviz<Graph>(graph, path)
        .nodePos(landscape.getCoordsMap())
        .nodeSizes(landscape.getQualityMap())
        .nodeColors(nodeColors)
        .run();
};

template <typename LS>
void printInstanceGraphviz(const LS & landscape,
                           const RestorationPlan<LS> & plan,
                           std::filesystem::path path) {
    using Graph = typename LS::Graph;
    using Node = typename LS::Graph::Node;
    using NodeIt = typename LS::Graph::NodeIt;
    using Arc = typename LS::Graph::Arc;
    using ArcIt = typename LS::Graph::ArcIt;
    using NodeColorMap = typename Graph::template NodeMap<int>;
    using ArcColorMap = typename Graph::template ArcMap<int>;

    const Graph & graph = landscape.getNetwork();

    NodeColorMap nodeColors(graph);
    for(NodeIt u(graph); u != lemon::INVALID; ++u)
        nodeColors[u] = plan.contains(u) ? 0xe05050 : 0x50e050;

    ArcColorMap arcColors(graph);
    for(ArcIt a(graph); a != lemon::INVALID; ++a) {
        if(!plan.contains(a)) {
            arcColors[a] = 0x101010;
            continue;
        }
        arcColors[a] = 0xe05050;
        Node u = graph.source(a);
        Node v = graph.target(a);
        nodeColors[u] = (nodeColors[u] & 0xffff00) | 0xe00000;
        nodeColors[v] = (nodeColors[v] & 0xffff00) | 0xe00000;
    }
    for(NodeIt u(graph); u != lemon::INVALID; ++u)
        nodeColors[u] = landscape.getQuality(u) > 0 ? nodeColors[u] : 0x101010;

    GraphToGraphviz<Graph>(graph, path)
        .nodePos(landscape.getCoordsMap())
        .nodeSizes(landscape.getQualityMap())
        .nodeColors(nodeColors)
        .arcScale(3)
        .arcColors(arcColors)
        .run();
};

template <typename LS>
void printSolutionGraphviz(const LS & landscape,
                           const RestorationPlan<LS> & plan,
                           const Solution & solution,
                           std::filesystem::path path) {
    using Graph = typename LS::Graph;
    using Node = typename LS::Graph::Node;
    using NodeIt = typename LS::Graph::NodeIt;
    using Arc = typename LS::Graph::Arc;
    using ArcIt = typename LS::Graph::ArcIt;
    using NodeColorMap = typename Graph::template NodeMap<int>;
    using ArcColorMap = typename Graph::template ArcMap<int>;
    using ArcSizesMap = typename Graph::template ArcMap<double>;

    const Graph & graph = landscape.getNetwork();

    NodeColorMap nodeColors(graph);
    for(NodeIt u(graph); u != lemon::INVALID; ++u)
        nodeColors[u] =
            plan.contains(u)
                ? 0xe05050
                : (landscape.getQuality(u) > 0 ? 0x50e050 : 0x101010);

    ArcColorMap arcColors(graph);
    for(ArcIt a(graph); a != lemon::INVALID; ++a) {
        arcColors[a] = plan.contains(a) ? 0xe05050 : 0x101010;
        if(plan.contains(a)) {
            Node u = graph.source(a);
            Node v = graph.target(a);
            nodeColors[u] = (nodeColors[u] & 0x00ffff) | 0xe00000;
            nodeColors[v] = (nodeColors[v] & 0x00ffff) | 0xe00000;
        }
    }

    ArcSizesMap arcSizes(graph, 1);
    const auto & arcCentrality = *Helper::corridorCentralityMap(landscape);
    double max_centrality = 0;
    for(ArcIt a(graph); a != lemon::INVALID; ++a)
        max_centrality = std::max(max_centrality, arcCentrality[a]);

    const double pow = std::log(10) / std::log(max_centrality);

    for(ArcIt a(graph); a != lemon::INVALID; ++a) {
        if(arcCentrality[a] == 0) continue;
        arcSizes[a] = std::pow(arcCentrality[a], pow);
        if(arcSizes[a] > 0) arcColors[a] = 0x5050e0;
        if(Node u = graph.source(a); landscape.getQuality(u) == 0)
            nodeColors[u] = (nodeColors[u] & 0xffff00) | 0x0000e0;
        if(Node v = graph.target(a); landscape.getQuality(v) == 0)
            nodeColors[v] = (nodeColors[v] & 0xffff00) | 0x0000e0;
    }

    GraphToGraphviz<Graph>(graph, path)
        .nodePos(landscape.getCoordsMap())
        .nodeSizes(landscape.getQualityMap())
        .nodeColors(nodeColors)
        .arcSizes(arcSizes)
        .arcColors(arcColors)
        .run();
};

}  // namespace Helper

#endif  // PRINT_HELPER_HPP


#pragma once

// system includes --------------------------------------------------
#include <Eigen/Core>
#include <Eigen/Dense>
#include <cmath>
#include <stdexcept>
#include <vector>
// own includes -----------------------------------------------------
#include "aux/message.hpp"
#include "aux/timer.hpp"
#include "shift_hermite.hpp"

namespace boltzmann {

template <typename BASIS, typename NUMERIC = double>
class ShiftHermite2D
{
 public:
  typedef NUMERIC numeric_t;
  typedef BASIS basis_t;
  typedef std::vector<numeric_t> std_vec_t;

 public:
  /**
   *
   *
   * @param basis
   * @param N      max. polynomial degree+1 in one variable
   */
  ShiftHermite2D(const basis_t& basis);

  void init();

  void shift(numeric_t* c, numeric_t x, numeric_t y);

  // debug
  const HShiftMatrix<numeric_t>& get_sx() { return sx_op_; }
  const HShiftMatrix<numeric_t>& get_sy() { return sy_op_; }

 private:
  typedef typename BASIS::elem_t elem_t;
  typedef typename boost::mpl::at_c<typename elem_t::types_t, 0>::type hx_t;
  typedef typename boost::mpl::at_c<typename elem_t::types_t, 1>::type hy_t;
  typedef Eigen::Matrix<numeric_t, Eigen::Dynamic, Eigen::Dynamic, Eigen::RowMajor> matrix_t;

 private:
  const basis_t& basis_;
  /// buffers
  matrix_t buf_in;
  matrix_t buf_out;
  /// permutation vector
  std::vector<unsigned int> perm_;
  bool is_initialized_;

  typename elem_t::Acc::template get<hx_t> get_hx;
  typename elem_t::Acc::template get<hy_t> get_hy;

  HShiftMatrix<numeric_t> sx_op_;
  HShiftMatrix<numeric_t> sy_op_;
};

// ----------------------------------------------------------------------
template <typename BASIS, typename NUMERIC>
ShiftHermite2D<BASIS, NUMERIC>::ShiftHermite2D(const basis_t& basis)
    : basis_(basis)
    ,

    perm_(basis.n_dofs())
    , is_initialized_(false)
{ /* empty */
}

// ----------------------------------------------------------------------
template <typename BASIS, typename NUMERIC>
void
ShiftHermite2D<BASIS, NUMERIC>::init()
{
  // max degree
  unsigned int max_kx =
      get_hx(*std::max_element(basis_.begin(),
                               basis_.end(),
                               [&](const elem_t& e1, const elem_t& e2) {
                                 return get_hx(e1).get_id().k < get_hx(e2).get_id().k;
                               }))
          .get_id()
          .k;
  unsigned int max_ky =
      get_hy(*std::max_element(basis_.begin(),
                               basis_.end(),
                               [&](const elem_t& e1, const elem_t& e2) {
                                 return get_hy(e1).get_id().k < get_hy(e2).get_id().k;
                               }))
          .get_id()
          .k;
  buf_in.resize(max_kx + 1, max_ky + 1);
  buf_out.resize(max_kx + 1, max_ky + 1);

  buf_in.fill(0);
  buf_out.fill(0);

  unsigned int stride = buf_in.cols();

  // build permutation vector
  unsigned int i = 0;
  for (auto elem = basis_.begin(); elem < basis_.end(); ++elem, ++i) {
    unsigned int kx = get_hx(*elem).get_id().k;
    unsigned int ky = get_hy(*elem).get_id().k;
    perm_[i] = kx * stride + ky;
  }

  Timer timer;
  // initialize shift matrices
  timer.restart();
  sx_op_.init(max_kx);
  sy_op_.init(max_ky);
  print_timer(timer.stop(), "initialize shift matrices");
  is_initialized_ = true;
}

// ----------------------------------------------------------------------
template <typename BASIS, typename NUMERIC>
void
ShiftHermite2D<BASIS, NUMERIC>::shift(numeric_t* c, numeric_t x, numeric_t y)
{
  assert(is_initialized_);

  // permute
  numeric_t* in = buf_in.data();

  for (unsigned int i = 0; i < basis_.n_dofs(); ++i) {
    in[perm_[i]] = c[i];
  }

  // apply shift matrices
  sx_op_.setx(x);
  sy_op_.setx(y);

  auto& Sx = sx_op_.get();
  auto& Sy = sy_op_.get();

  buf_out = Sx * buf_in * Sy.transpose();

  numeric_t* out = buf_out.data();

  // invert permutation and overwrite c
  for (unsigned int i = 0; i < basis_.n_dofs(); ++i) {
    c[i] = out[perm_[i]];
  }
}

}  // end namespace boltzmann