// ---------------------------------------------------------------------
//
// Copyright (c) 2018 - 2019 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/equation_systems.h>
#include <libmesh/exodusII_io.h>
#include <libmesh/mesh.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/IBFEDirectForcingKinematics.h>
#include <ibamr/IBFEMethod.h>
#include <ibamr/INSStaggeredHierarchyIntegrator.h>

#include <ibtk/AppInitializer.h>
#include <ibtk/IBTKInit.h>
#include <ibtk/IBTK_MPI.h>
#include <ibtk/LEInteractor.h>
#include <ibtk/ibtk_utilities.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>

inline double
kernel(double x)
{
    x += 4.;
    const double x2 = x * x;
    const double x3 = x * x2;
    const double x4 = x * x3;
    const double x5 = x * x4;
    const double x6 = x * x5;
    const double x7 = x * x6;
    if (x <= 0.)
        return 0.;
    else if (x <= 1.)
        return .1984126984126984e-3 * x7;
    else if (x <= 2.)
        return .1111111111111111e-1 * x6 - .1388888888888889e-2 * x7 - .3333333333333333e-1 * x5 +
               .5555555555555556e-1 * x4 - .5555555555555556e-1 * x3 + .3333333333333333e-1 * x2 -
               .1111111111111111e-1 * x + .1587301587301587e-2;
    else if (x <= 3.)
        return .4333333333333333 * x5 - .6666666666666667e-1 * x6 + .4166666666666667e-2 * x7 - 1.500000000000000 * x4 +
               3.055555555555556 * x3 - 3.700000000000000 * x2 + 2.477777777777778 * x - .7095238095238095;
    else if (x <= 4.)
        return 9. * x4 - 1.666666666666667 * x5 + .1666666666666667 * x6 - .6944444444444444e-2 * x7 -
               28.44444444444444 * x3 + 53. * x2 - 54.22222222222222 * x + 23.59047619047619;
    else if (x <= 5.)
        return 96. * x3 - 22.11111111111111 * x4 + 3. * x5 - .2222222222222222 * x6 + .6944444444444444e-2 * x7 -
               245.6666666666667 * x2 + 344. * x - 203.9650793650794;
    else if (x <= 6.)
        return 483.5000000000000 * x2 - 147.0555555555556 * x3 + 26.50000000000000 * x4 - 2.833333333333333 * x5 +
               .1666666666666667 * x6 - .4166666666666667e-2 * x7 - 871.2777777777778 * x + 664.0904761904762;
    else if (x <= 7.)
        return 943.1222222222222 * x - 423.7000000000000 * x2 + 104.9444444444444 * x3 - 15.50000000000000 * x4 +
               1.366666666666667 * x5 - .6666666666666667e-1 * x6 + .1388888888888889e-2 * x7 - 891.1095238095238;
    else if (x <= 8.)
        return 416.1015873015873 - 364.0888888888889 * x + 136.5333333333333 * x2 - 28.44444444444444 * x3 +
               3.555555555555556 * x4 - .2666666666666667 * x5 + .1111111111111111e-1 * x6 - .1984126984126984e-3 * x7;
    else
        return 0.;
} // kernel

void
cylinder_kinematics(double /*data_time*/, Eigen::Vector3d& U_com, Eigen::Vector3d& W_com, void* /*ctx*/)
{
    U_com.setZero();
    W_com.setZero();

    return;
} // cylinder_kinematics

// Function prototypes
static ofstream drag_stream, lift_stream;
static double R;
void postprocess_data(Pointer<PatchHierarchy<NDIM> > patch_hierarchy,
                      Pointer<INSHierarchyIntegrator> navier_stokes_integrator,
                      Mesh& mesh,
                      EquationSystems* 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();

        // Setup user-defined kernel function.
        LEInteractor::s_kernel_fcn = &kernel;
        LEInteractor::s_kernel_fcn_stencil_size = 8;

        // 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 exodus_filename = app_initializer->getExodusIIFilename();

        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.
        Mesh solid_mesh(init.comm(), NDIM);
        const double dx = input_db->getDouble("DX");
        const double ds = input_db->getDouble("MFAC") * dx;
        string elem_type = input_db->getString("ELEM_TYPE");
        R = input_db->getDouble("R");
        if (NDIM == 2 && (elem_type == "TRI3" || 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);
                solid_mesh.add_point(libMesh::Point(R * cos(theta), R * sin(theta)));
            }
            TriangleInterface triangle(solid_mesh);
            triangle.triangulation_type() = TriangleInterface::GENERATE_CONVEX_HULL;
            triangle.elem_type() = Utility::string_to_enum<ElemType>(elem_type);
            triangle.desired_area() = 1.5 * sqrt(3.0) / 4.0 * ds * ds;
            triangle.insert_extra_points() = true;
            triangle.smooth_after_generating() = true;
            triangle.triangulate();
#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 = 2.0 * M_PI * R / ds;
            const int r = log2(0.25 * num_circum_segments);
            MeshTools::Generation::build_sphere(solid_mesh, R, r, Utility::string_to_enum<ElemType>(elem_type));
        }

        // Ensure nodes on the surface are on the analytic boundary.
        MeshBase::element_iterator el_end = solid_mesh.elements_end();
        for (MeshBase::element_iterator el = solid_mesh.elements_begin(); el != el_end; ++el)
        {
            Elem* const elem = *el;
            for (unsigned int side = 0; side < elem->n_sides(); ++side)
            {
                const bool at_mesh_bdry = !elem->neighbor_ptr(side);
                if (!at_mesh_bdry) continue;
                for (unsigned int k = 0; k < elem->n_nodes(); ++k)
                {
                    if (!elem->is_node_on_side(k, side)) continue;
                    Node& n = elem->node_ref(k);
                    n = R * n.unit();
                }
            }
        }
        solid_mesh.prepare_for_use();
        Mesh& mesh = solid_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 = new INSStaggeredHierarchyIntegrator(
            "INSStaggeredHierarchyIntegrator",
            app_initializer->getComponentDatabase("INSStaggeredHierarchyIntegrator"));

        Pointer<IBFEMethod> ib_method_ops =
            new IBFEMethod("IBFEMethod",
                           app_initializer->getComponentDatabase("IBFEMethod"),
                           &mesh,
                           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);

        // Create IBFE direct forcing kinematics object.
        Pointer<IBFEDirectForcingKinematics> df_kinematics_ops = new IBFEDirectForcingKinematics(
            "cylinder_dfk",
            app_initializer->getComponentDatabase("CylinderIBFEDirectForcingKinematics"),
            ib_method_ops,
            /*part*/ 0,
            /*register_for_restart*/ true);
        ib_method_ops->registerDirectForcingKinematics(df_kinematics_ops, /*part*/ 0);

        // Specify structure kinematics
        FreeRigidDOFVector solve_dofs;
        solve_dofs.setZero();
        df_kinematics_ops->setSolveRigidBodyVelocity(solve_dofs);
        df_kinematics_ops->registerKinematicsFunction(&cylinder_kinematics, NULL);

        // Configure the IBFE solver.
        ib_method_ops->initializeFEEquationSystems();
        EquationSystems* equation_systems = ib_method_ops->getFEDataManager()->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> exodus_io(uses_exodus ? new ExodusII_IO(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();
            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)
            {
                exodus_io->write_timestep(
                    exodus_filename, *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);

            drag_stream.precision(10);
            lift_stream.precision(10);
        }

        // 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)
                {
                    exodus_io->write_timestep(
                        exodus_filename, *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))
            {
                postprocess_data(patch_hierarchy,
                                 navier_stokes_integrator,
                                 mesh,
                                 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();
        }

        // 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& mesh,
                 EquationSystems* equation_systems,
                 const int /*iteration_num*/,
                 const double loop_time,
                 const string& /*data_dump_dirname*/)
{
    const unsigned int dim = mesh.mesh_dimension();
    double F_integral[NDIM];
    for (unsigned int d = 0; d < NDIM; ++d) F_integral[d] = 0.0;

    System& F_system = equation_systems->get_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);
    const DofMap& dof_map = F_system.get_dof_map();
    FEType fe_type = dof_map.variable_type(0);
    std::unique_ptr<FEBase> fe(FEBase::build(dim, fe_type));
    std::unique_ptr<QBase> qrule = fe_type.default_quadrature_rule(dim);
    fe->attach_quadrature_rule(qrule.get());
    const vector<double>& JxW = fe->get_JxW();
    const vector<vector<double> >& phi = fe->get_phi();

    std::vector<std::vector<unsigned int> > dof_indices(NDIM);
    boost::multi_array<double, 2> F_node;
    VectorValue<double> F;

    const MeshBase::const_element_iterator el_begin = mesh.active_local_elements_begin();
    const MeshBase::const_element_iterator el_end = mesh.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)
        {
            dof_map.dof_indices(elem, dof_indices[d], d);
        }
        get_values_for_interpolation(F_node, *F_ghost_vec, dof_indices);

        const unsigned int n_qp = qrule->n_points();
        for (unsigned int qp = 0; qp < n_qp; ++qp)
        {
            interpolate(F, qp, F_node, phi);
            for (int d = 0; d < NDIM; ++d)
            {
                F_integral[d] += F(d) * JxW[qp];
            }
        }
    }
    IBTK_MPI::sumReduction(F_integral, NDIM);
    static const double rho = 1.0;
    static const double U_max = 1.0;
    static const double D = 2.0 * R;
    if (IBTK_MPI::getRank() == 0)
    {
        drag_stream << loop_time << " " << -F_integral[0] / (0.5 * rho * U_max * U_max * D) << endl;
        lift_stream << loop_time << " " << -F_integral[1] / (0.5 * rho * U_max * U_max * D) << endl;
    }
    return;
} // postprocess_data


/*
 * Copyright 2020 Robert Bosch GmbH
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 *
 * SPDX-License-Identifier: Apache-2.0
 */
/**
 * \file cloe/utility/geometry.hpp
 */

#pragma once
#ifndef CLOE_UTILITY_GEOMETRY_HPP_
#define CLOE_UTILITY_GEOMETRY_HPP_

#include <Eigen/Geometry>

namespace cloe {
namespace utility {

/**
 * QuaternionFromRPY calculates a quaternion from roll, pitch and yaw.
 */
inline Eigen::Quaterniond quaternion_from_rpy(double roll, double pitch, double yaw) {
  // ZYX body flxed rotations
  Eigen::Quaterniond qt = Eigen::AngleAxisd(yaw, Eigen::Vector3d::UnitZ()) *
                          Eigen::AngleAxisd(pitch, Eigen::Vector3d::UnitY()) *
                          Eigen::AngleAxisd(roll, Eigen::Vector3d::UnitX());
  return qt;
}

/**
 * PoseFromRotationTranslation calculates the pose from rotation and translation.
 */
inline Eigen::Isometry3d pose_from_rotation_translation(const Eigen::Quaterniond& quaternion,
                                                        const Eigen::Vector3d& trans) {
  Eigen::Isometry3d pose;
  pose.setIdentity();
  pose.linear() = quaternion.matrix();
  pose.translation() = trans;
  return pose;
}

/**
 * Change a point's frame of reference.
 * Both the point and the child frame need to have the same parent frame.
 */
inline void transform_to_child_frame(const Eigen::Isometry3d& child_frame, Eigen::Vector3d* point) {
  *point = child_frame.inverse() * (*point);
}

}  // namespace utility
}  // namespace cloe

#endif  // CLOE_UTILITY_GEOMETRY_HPP_


// Copyright (C) 2014 The Regents of the University of California (Regents).
// All rights reserved.
//
// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions are
// met:
//
//     * Redistributions of source code must retain the above copyright
//       notice, this list of conditions and the following disclaimer.
//
//     * Redistributions in binary form must reproduce the above
//       copyright notice, this list of conditions and the following
//       disclaimer in the documentation and/or other materials provided
//       with the distribution.
//
//     * Neither the name of The Regents or University of California nor the
//       names of its contributors may be used to endorse or promote products
//       derived from this software without specific prior written permission.
//
// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
// AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
// ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDERS OR CONTRIBUTORS BE
// LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
// CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
// SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
// INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
// CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
// ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
// POSSIBILITY OF SUCH DAMAGE.
//
// Please contact the author of this library if you have any questions.
// Author: Chris Sweeney (cmsweeney@cs.ucsb.edu)
//         Torsten Sattler (sattlert@inf.ethz.ch)

#include "theia/vision/sfm/camera/camera_pose.h"

#include <Eigen/Core>
#include <Eigen/Geometry>
#include <glog/logging.h>

#include "theia/alignment/alignment.h"

namespace theia {

using Eigen::Map;
using Eigen::Matrix3d;
using Eigen::Matrix;
using Eigen::MatrixXd;
using Eigen::Vector2d;
using Eigen::Vector3d;

CameraPose::CameraPose() : k1_(0.0), k2_(0.0), k3_(0.0), k4_(0.0) {
  calibration_matrix_.setIdentity();
  transformation_matrix_.setIdentity();
  projection_matrix_.setIdentity();
}

CameraPose::~CameraPose() {}

void CameraPose::InitializePose(const Matrix3d& rotation,
                                const Vector3d& translation,
                                const Matrix3d& calibration, const double k1,
                                const double k2, const double k3,
                                const double k4) {
  transformation_matrix_.linear() = rotation;
  transformation_matrix_.translation() = translation;
  calibration_matrix_ = calibration;
  projection_matrix_ = calibration_matrix_ * transformation_matrix_.matrix();
  k1_ = k1;
  k2_ = k2;
  k3_ = k3;
  k4_ = k4;
}

// The projection matrix is defined such that P = K * [R | t], where K = diag(f,
// f, 1) for focal length f. The rows of the rotation matrix must have the same
// norm, so we can recover the focal length by comparing the norms of the first
// and third projection matrix row (since only one of them will be affected by
// the focal length).
void CameraPose::InitializePose(const Matrix<double, 3, 4>& projection_matrix,
                                const double k1, const double k2,
                                const double k3, const double k4) {
  projection_matrix_ = projection_matrix;
  k1_ = k1;
  k2_ = k2;
  k3_ = k3;
  k4_ = k4;

  // Recover focal length and construct the calibration matrix. NOTE: setting
  // the calibration matrix to diag(1, 1, f) is OK since we only recover the
  // focal length as the ratio of diagonal elements (see focal_length() method).
  const double inv_focal_length = sqrt(projection_matrix_.block<1, 3>(
      2, 0).squaredNorm() / projection_matrix.block<1, 3>(0, 0).squaredNorm());
  calibration_matrix_ =
      Eigen::DiagonalMatrix<double, 3>(1.0, 1.0, inv_focal_length);

  // Since P = K * T, where T = [R | t], we can recover T with T = K^-1 * P. We
  // must use a temp matrix so that the AffineCompact3d type of Eigen does not
  // fix the det of the rotation. We use the det of the rotation to recover the
  // sign of the focal length.
  const Eigen::Matrix<double, 3, 4> temp_transformation =
      calibration_matrix_.inverse() * projection_matrix_;

  // Recover the sign of the focal length using the fact that a rotation matrix
  // should have a determinant of 1.
  double det = temp_transformation.block<3, 3>(0, 0).determinant();
  if (det < 0.0) {
    calibration_matrix_(2, 2) *= -1.0;
  }
  transformation_matrix_ = calibration_matrix_.inverse() * projection_matrix_;
}

void CameraPose::InitializePose(
    const Matrix<double, 3, 4>& transformation_matrix,
    const Matrix3d& calibration, const double k1, const double k2,
    const double k3, const double k4) {
  transformation_matrix_ = transformation_matrix;
  calibration_matrix_ = calibration;
  projection_matrix_ = calibration_matrix_ * transformation_matrix_.matrix();
  k1_ = k1;
  k2_ = k2;
  k3_ = k3;
  k4_ = k4;
}

void CameraPose::InitializePose(const CameraPose& pose) {
  InitializePose(pose.transformation_matrix_.matrix(), pose.calibration_matrix_,
                 pose.k1_, pose.k2_, pose.k3_, pose.k4_);
}

void CameraPose::WorldToCamera(const Vector3d& world_point,
                               Vector3d* camera_point) const {
  *camera_point = transformation_matrix_ * world_point;
}

void CameraPose::WorldToCamera(const std::vector<Vector3d>& world_point,
                               std::vector<Vector3d>* camera_point) const {
  Map<const Matrix<double, 3, Eigen::Dynamic> > world_point_matrix(
      world_point[0].data(), 3, world_point.size());
  camera_point->clear();
  camera_point->resize(world_point.size());
  Map<Matrix<double, 3, Eigen::Dynamic> > camera_point_matrix(
      (*camera_point)[0].data(), 3, camera_point->size());
  camera_point_matrix = transformation_matrix_ * world_point_matrix;
}

void CameraPose::CameraToWorld(const Vector3d& camera_point,
                               Vector3d* world_point) const {
  *world_point = transformation_matrix_.inverse() * camera_point;
}

void CameraPose::CameraToWorld(const std::vector<Vector3d>& camera_point,
                               std::vector<Vector3d>* world_point) const {
  Map<const Matrix<double, 3, Eigen::Dynamic> > camera_point_matrix(
      camera_point[0].data(), 3, camera_point.size());
  world_point->clear();
  world_point->resize(camera_point.size());
  Map<Matrix<double, 3, Eigen::Dynamic> > world_point_matrix(
      (*world_point)[0].data(), 3, world_point->size());
  world_point_matrix = transformation_matrix_.inverse() * camera_point_matrix;
}

void CameraPose::CameraToImage(const Vector3d& camera_point,
                               Vector2d* image_point) const {
  *image_point = (calibration_matrix_ * camera_point).hnormalized();
}

void CameraPose::CameraToImage(const std::vector<Vector3d>& camera_point,
                               std::vector<Vector2d>* image_point) const {
  Map<const Matrix<double, 3, Eigen::Dynamic> > camera_point_matrix(
      camera_point[0].data(), 3, camera_point.size());
  image_point->clear();
  image_point->resize(camera_point.size());
  Map<Matrix<double, 2, Eigen::Dynamic> > image_point_matrix(
      (*image_point)[0].data(), 2, image_point->size());
  image_point_matrix = (calibration_matrix_.inverse() * camera_point_matrix)
      .colwise().hnormalized();
}

bool CameraPose::WorldToImage(const Vector3d& world_point,
                              Vector2d* image_point) const {
  const Vector3d proj_point = projection_matrix_ * world_point.homogeneous();
  // Return false if the point is behind the camera.
  if (calibration_matrix_(0, 0) * proj_point[2] <= 0.0) {
    return false;
  }
  *image_point = proj_point.hnormalized();
  return true;
}

void CameraPose::WorldToImage(const std::vector<Vector3d>& world_point,
                              std::vector<Vector2d>* image_point) const {
  Map<const Matrix<double, 3, Eigen::Dynamic> > world_point_matrix(
      world_point[0].data(), 3, world_point.size());

  // Create a temp 3xN matrix of the projected points... eigen tends to struggle
  // when using colwise().homogeneous() and colwise().hnormalized() in the same
  // expression.
  Matrix<double, 3, Eigen::Dynamic> temp_proj_mat =
      projection_matrix_ * (world_point_matrix.colwise().homogeneous());

  image_point->clear();
  image_point->resize(world_point.size());
  Map<Matrix<double, 2, Eigen::Dynamic> > image_point_matrix(
      (*image_point)[0].data(), 2, image_point->size());
  image_point_matrix = temp_proj_mat.colwise().hnormalized();
}

void CameraPose::UndistortImagePoint(const Vector2d& distorted_point,
                                     Vector2d* undistorted_point) const {
  const double r = distorted_point.squaredNorm();
  const double w_term =
      1.0 + k1_ * r + k2_ * r * r + k3_ * r * r * r + k4_ * r * r * r * r;
  *undistorted_point = distorted_point / w_term;
}

void CameraPose::UndistortImagePoint(
    const std::vector<Vector2d>& distorted_point,
    std::vector<Vector2d>* undistorted_point) const {
  undistorted_point->clear();
  undistorted_point->resize(distorted_point.size());

  Eigen::Map<const Eigen::Matrix<double, 2, Eigen::Dynamic> > distorted_map(
      distorted_point[0].data(), 2, distorted_point.size());
  Eigen::ArrayXd radius = distorted_map.colwise().squaredNorm();
  radius =
      1.0 + radius * k1_ + radius * radius * k2_ +
      radius * radius * radius * k3_ + radius * radius * radius * radius * k4_;

  Eigen::Map<Eigen::Matrix<double, 2, Eigen::Dynamic> > undistorted_map(
      (*undistorted_point)[0].data(), 2, undistorted_point->size());
  undistorted_map = distorted_map;
  undistorted_map = undistorted_map.array().rowwise() / radius.transpose();
}

}  // namespace theia


// ----------------------------------------------------------------------------
// -                        Open3D: www.open3d.org                            -
// ----------------------------------------------------------------------------
// The MIT License (MIT)
//
// Copyright (c) 2018 www.open3d.org
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
// FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS
// IN THE SOFTWARE.
// ----------------------------------------------------------------------------

#include "Registration.h"

#include <cstdlib>
#include <ctime>
#include <iostream>

#include <Eigen/Dense>
#include <Core/Utility/Console.h>
#include <Core/Geometry/PointCloud.h>
#include <Core/Geometry/KDTreeFlann.h>
#include <Core/Registration/Feature.h>

namespace open3d {

namespace {

double compute_distance_squared(const Eigen::Vector3d p1, const Eigen::Vector3d p2){
    Eigen::Vector3d diff = p1 - p2;
    return diff[0]*diff[0] + diff[1]*diff[1] + diff[2]*diff[2];
}

int compute_closest_vertex(const Eigen::Vector3d best_intpoint, int f, const TriangleMesh target){
    int vertex0 = target.triangles_[f][0];
    int vertex1 = target.triangles_[f][1];
    int vertex2 = target.triangles_[f][2];
    double d0 = compute_distance_squared(best_intpoint, target.vertices_[vertex0]);
    double d1 = compute_distance_squared(best_intpoint, target.vertices_[vertex1]);
    double d2 = compute_distance_squared(best_intpoint, target.vertices_[vertex2]);
    //PrintInfo("d0: %.2f\n", std::sqrt(d0));
    //PrintInfo("d1: %.2f\n", std::sqrt(d1));
    //PrintInfo("d2: %.2f\n", std::sqrt(d2));
    if(d0 <= d1 && d0 <= d2){
        //PrintInfo("picking d0\n");
        return vertex0;
    }
    else{
        if(d1 <= d0 && d1 <= d2){
            //PrintInfo("picking d1\n");
            return vertex1;
        }
        else{
            if(d2 <= d0 && d2 <= d1){
                //PrintInfo("picking d2\n");
                return vertex2;
            }
            else{
                //PrintInfo("Big mistake\n");
                return -1;
            }
        }
    }
}

// Returns true if a ray intersects a triangle.
bool RayIntersectsTriangle(const Eigen::Vector3d rayOrigin,
                        const Eigen::Vector3d rayVector,
                        const Eigen::Vector3d vertex0,
                        const Eigen::Vector3d vertex1,
                        const Eigen::Vector3d vertex2,
                        Eigen::Vector3d& outIntersectionPoint)
{
    const float EPSILON = 0.0000001;
    Eigen::Vector3d edge1, edge2, h, s, q;
    float a,f,u,v;
    edge1 = vertex1 - vertex0;
    edge2 = vertex2 - vertex0;
    h = rayVector.cross(edge2);
    a = edge1.dot(h);
    if (a > -EPSILON && a < EPSILON)
        return false;    // This ray is parallel to this triangle.
    f = 1.0/a;
    s = rayOrigin - vertex0;
    u = f * (s.dot(h));
    if (u < 0.0 || u > 1.0)
        return false;
    q = s.cross(edge1);
    v = f * rayVector.dot(q);
    if (v < 0.0 || u + v > 1.0)
        return false;
    // At this stage we can compute t to find out where the intersection point is on the line.
    float t = f * edge2.dot(q);
    if (t > EPSILON) // ray intersection
    {
        outIntersectionPoint = rayOrigin + rayVector * t;
        return true;
    }
    else // This means that there is a line intersection but not a ray intersection.
        return false;
}

// For each of the neighboring triangles to each point, find which one is intersected by the line.
void find_sc_corr(const TriangleMesh &target, const std::vector<int> &neigh_triangles,\
        const Eigen::Vector3d &source_point, const Eigen::Vector3d &source_normal,\
        double max_correspondence_distance, \
        std::vector<int> &target_corr, std::vector<double> & sc){

    // For each triangle in the set, find if it is intersecting.
    Eigen::Vector3d best_intpoint(0.0, 0.0, 0.0);
    double best_intpoint_dist = 100000000;
    int intpoint_face_ix = -1;
    for (int i = 0; i < (int) neigh_triangles.size(); i++){
        Eigen::Vector3i f = target.triangles_[neigh_triangles[i]];

        Eigen::Vector3d intPoint;
        bool intersect = RayIntersectsTriangle(source_point, source_normal, \
                                                target.vertices_[f[0]], target.vertices_[f[1]],\
                                                target.vertices_[f[2]], intPoint);
        if(!intersect){
            continue;
        }

        double dist2 = compute_distance_squared(intPoint, source_point);
        if(dist2 < best_intpoint_dist){
            best_intpoint_dist = dist2;
            best_intpoint = intPoint;
            intpoint_face_ix = neigh_triangles[i];
        }
    }
    if(intpoint_face_ix >= 0){
        int f = intpoint_face_ix;

        // Pick the vertex closest to the intersection point.
        target_corr[0] = compute_closest_vertex(best_intpoint, f, target);
        double dist2 = compute_distance_squared(target.vertices_[target_corr[0]], best_intpoint);
        Eigen::Vector3d v = best_intpoint;
        //PrintInfo("--------------\n");
        //PrintInfo("Intpoint: %.2f %.2f %.2f\n", v[0], v[1], v[2]);
        v = target.vertices_[target_corr[0]];
        //PrintInfo("vrtpoint: %.2f %.2f %.2f\n", v[0], v[1], v[2]);
        //PrintInfo("vrtpoint to intpoint: %.2f\n", std::sqrt(dist2));
        v = source_point;
        //PrintInfo("Srcpoint: %.2f %.2f %.2f\n", v[0], v[1], v[2]);
        dist2 = compute_distance_squared(target.vertices_[target_corr[0]], source_point);

        // Compute shape complementarity
        sc[0] = target.vertex_normals_[target_corr[0]].dot(-source_normal);
        //PrintInfo("normal mult: %f\n", sc[0]);
        sc[0] = sc[0]*std::exp(-0.5*(dist2));
        //PrintInfo("shape comp: %f \n", sc[0]);
    }
    else{
        target_corr[0] = -1;
        sc[0] = 0.0;
    }
}

    //
// PGC 2019: Get registration results based on shape complementarity
RegistrationResult GetRegistrationResultAndCorrespondencesShapeComplementarity(
        const PointCloud &source,
        const TriangleMesh &target,
        const KDTreeFlann &target_faces_kdtree,
        double max_correspondence_distance,
        const Eigen::Matrix4d &transformation,
        const Feature &source_feature,
        const Feature &target_feature,
        const int fitness_type
        ) {

    RegistrationResult result(transformation);
    if (max_correspondence_distance <= 0.0) {
        return std::move(result);
    }

    double error2 = 0.0;

#ifdef _OPENMP
#pragma omp parallel
    {
#endif
        double error2_private = 0.0;
        CorrespondenceSet correspondence_set_private;
#ifdef _OPENMP
#pragma omp for nowait
#endif
        // Compute the correspondences based on shape complementarity.
        // First find all target points within X of each source point. X=max_correspondence_distance
        // Then get all the triangles on which each target point participates.
        // Find the intersection of the line between source point/normal  and each of the triangles.
        // If the intersection is less than 2, we are good.
        // Find the nearest neighbor to the intersection point.
        for (int i = 0; i < (int)source.points_.size(); i++) {
            std::vector<int> indices(20);
            std::vector<double> dists(20);
            std::vector<int> target_corr(1);
            std::vector<double> target_sc(1);
            const auto &point = source.points_[i];
            const auto &normal= source.normals_[i];
            // Identify all faces within X of the target, up to 20
            if (target_faces_kdtree.SearchHybrid(point, max_correspondence_distance,
                                           20, indices, dists) > 0) {
                // Find the intersection of the line between source point/normal
                //      and each of the triangles.
                find_sc_corr(target, indices, point, normal, max_correspondence_distance, \
                        target_corr, target_sc);

                if(target_sc[0] > 0){
                    //PrintInfo("Adding correspondences: %d %d\n", i, target_corr[0]);
                    error2_private += target_sc[0];
                    correspondence_set_private.push_back(
                        Eigen::Vector2i(i, target_corr[0]));
                }
            }
        }
#ifdef _OPENMP
#pragma omp critical
#endif
        {
            for (int i = 0; i < (int)correspondence_set_private.size(); i++) {
                result.correspondence_set_.push_back(
                        correspondence_set_private[i]);
            }
            error2 += error2_private;
        }
#ifdef _OPENMP
    }
#endif

    if (result.correspondence_set_.empty()) {
        result.fitness_ = 0.0;
        result.inlier_rmse_ = 1000.0;
    } else {
        size_t corres_number = result.correspondence_set_.size();
        //result.fitness_ = (double)corres_number / (double)source.points_.size();
        if(fitness_type == 1){
            // PGC 2019: fitness: the number of correspondences
            result.fitness_ = (double)corres_number;
        }
        else{
            if(fitness_type == 2){
                // Compute the descriptor distances for all pairs of correspondences.
                double myfitness = 0;
                for (int i = 0; i < (int)result.correspondence_set_.size(); i++) {
                    int s_vix =  result.correspondence_set_[i][0];
                    int t_vix = result.correspondence_set_[i][1];
                    Eigen::VectorXd feat_s = Eigen::VectorXd(source_feature.data_.col(s_vix));
                    Eigen::VectorXd feat_t = Eigen::VectorXd(target_feature.data_.col(t_vix));
                    double desc_dist = 0.0;
//                    std::cout << "Dimension = " << source_feature.Dimension() << std::endl;
//                    std::cout << "Size = " << source_feature.Num() << std::endl;
                    for (int j = 0; j < source_feature.Dimension(); j++){
                        double dist = feat_t[j] - feat_s[j];
                        dist = dist*dist;
                        desc_dist += dist;
                    }

                    myfitness += 1.0/desc_dist;
                }

                result.fitness_ = myfitness;

            }
            else{
                if (fitness_type == 3.0){ // Use shape complementarity (which is in the error).
                    result.fitness_ = error2 / (double)source.points_.size();
                }
                else{
                    result.fitness_ = (double)corres_number / (double)source.points_.size();
                }
            }
        }
        result.inlier_rmse_ = std::sqrt(error2 / (double)corres_number);
    }
    return std::move(result);
}

// PGC 2019: include the fitness type and the features for a custom fitness function.
RegistrationResult GetRegistrationResultAndCorrespondencesCustom(
        const PointCloud &source,
        const PointCloud &target,
        const KDTreeFlann &target_kdtree,
        double max_correspondence_distance,
        const Eigen::Matrix4d &transformation,
        const Feature &source_feature,
        const Feature &target_feature,
        const int fitness_type
        ) {
    RegistrationResult result(transformation);
    if (max_correspondence_distance <= 0.0) {
        return std::move(result);
    }

    double error2 = 0.0;

#ifdef _OPENMP
#pragma omp parallel
    {
#endif
        double error2_private = 0.0;
        CorrespondenceSet correspondence_set_private;
#ifdef _OPENMP
#pragma omp for nowait
#endif
        for (int i = 0; i < (int)source.points_.size(); i++) {
            std::vector<int> indices(1);
            std::vector<double> dists(1);
            const auto &point = source.points_[i];
            if (target_kdtree.SearchHybrid(point, max_correspondence_distance,
                                           1, indices, dists) > 0) {
                error2_private += dists[0];
                correspondence_set_private.push_back(
                        Eigen::Vector2i(i, indices[0]));
            }
        }
#ifdef _OPENMP
#pragma omp critical
#endif
        {
            for (int i = 0; i < (int)correspondence_set_private.size(); i++) {
                result.correspondence_set_.push_back(
                        correspondence_set_private[i]);
            }
            error2 += error2_private;
        }
#ifdef _OPENMP
    }
#endif

    if (result.correspondence_set_.empty()) {
        result.fitness_ = 0.0;
        result.inlier_rmse_ = 0.0;
    } else {
        size_t corres_number = result.correspondence_set_.size();
        //result.fitness_ = (double)corres_number / (double)source.points_.size();
        if(fitness_type == 1){
            // PGC 2019: fitness: the number of correspondences
            result.fitness_ = (double)corres_number;
        }
        else{
            if(fitness_type == 2){
                // Compute the descriptor distances for all pairs of correspondences.
                double myfitness = 0;
                for (int i = 0; i < (int)result.correspondence_set_.size(); i++) {
                    int s_vix =  result.correspondence_set_[i][0];
                    int t_vix = result.correspondence_set_[i][1];
                    Eigen::VectorXd feat_s = Eigen::VectorXd(source_feature.data_.col(s_vix));
                    Eigen::VectorXd feat_t = Eigen::VectorXd(target_feature.data_.col(t_vix));
                    double desc_dist = 0.0;
//                    std::cout << "Dimension = " << source_feature.Dimension() << std::endl;
//                    std::cout << "Size = " << source_feature.Num() << std::endl;
                    for (int j = 0; j < source_feature.Dimension(); j++){
                        double dist = feat_t[j] - feat_s[j];
                        dist = dist*dist;
                        desc_dist += dist;
                    }

                    myfitness += 1.0/desc_dist;
                }

                result.fitness_ = myfitness;

            }
            else{
                result.fitness_ = (double)corres_number / (double)source.points_.size();
            }
        }
        result.inlier_rmse_ = std::sqrt(error2 / (double)corres_number);
    }
    return std::move(result);
}

RegistrationResult GetRegistrationResultAndCorrespondences(
        const PointCloud &source,
        const PointCloud &target,
        const KDTreeFlann &target_kdtree,
        double max_correspondence_distance,
        const Eigen::Matrix4d &transformation) {
    RegistrationResult result(transformation);
    if (max_correspondence_distance <= 0.0) {
        return std::move(result);
    }

    double error2 = 0.0;

#ifdef _OPENMP
#pragma omp parallel
    {
#endif
        double error2_private = 0.0;
        CorrespondenceSet correspondence_set_private;
#ifdef _OPENMP
#pragma omp for nowait
#endif
        for (int i = 0; i < (int)source.points_.size(); i++) {
            std::vector<int> indices(1);
            std::vector<double> dists(1);
            const auto &point = source.points_[i];
            if (target_kdtree.SearchHybrid(point, max_correspondence_distance,
                                           1, indices, dists) > 0) {
                error2_private += dists[0];
                correspondence_set_private.push_back(
                        Eigen::Vector2i(i, indices[0]));
            }
        }
#ifdef _OPENMP
#pragma omp critical
#endif
        {
            for (int i = 0; i < (int)correspondence_set_private.size(); i++) {
                result.correspondence_set_.push_back(
                        correspondence_set_private[i]);
            }
            error2 += error2_private;
        }
#ifdef _OPENMP
    }
#endif

    if (result.correspondence_set_.empty()) {
        result.fitness_ = 0.0;
        result.inlier_rmse_ = 0.0;
    } else {
        size_t corres_number = result.correspondence_set_.size();
        //result.fitness_ = (double)corres_number / (double)source.points_.size();
        // PGC 2019: fitness is the number of correspondences, not the fraction.
        result.fitness_ = (double)corres_number;
        result.inlier_rmse_ = std::sqrt(error2 / (double)corres_number);
    }
    return std::move(result);
}

RegistrationResult EvaluateRANSACBasedOnCorrespondence(
        const PointCloud &source,
        const PointCloud &target,
        const CorrespondenceSet &corres,
        double max_correspondence_distance,
        const Eigen::Matrix4d &transformation) {
    RegistrationResult result(transformation);
    double error2 = 0.0;
    int good = 0;
    double max_dis2 = max_correspondence_distance * max_correspondence_distance;
    for (const auto &c : corres) {
        double dis2 =
                (source.points_[c[0]] - target.points_[c[1]]).squaredNorm();
        if (dis2 < max_dis2) {
            good++;
            error2 += dis2;
        }
    }
    if (good == 0) {
        result.fitness_ = 0.0;
        result.inlier_rmse_ = 0.0;
    } else {
        result.fitness_ = (double)good / (double)corres.size();
        result.inlier_rmse_ = std::sqrt(error2 / (double)good);
    }
    return result;
}

}  // unnamed namespace

RegistrationResult EvaluateRegistration(
        const PointCloud &source,
        const PointCloud &target,
        double max_correspondence_distance,
        const Eigen::Matrix4d
                &transformation /* = Eigen::Matrix4d::Identity()*/) {
    KDTreeFlann kdtree;
    kdtree.SetGeometry(target);
    PointCloud pcd = source;
    if (transformation.isIdentity() == false) {
        pcd.Transform(transformation);
    }
    return GetRegistrationResultAndCorrespondences(
            pcd, target, kdtree, max_correspondence_distance, transformation);
}

RegistrationResult RegistrationICP(
        const PointCloud &source,
        const PointCloud &target,
        double max_correspondence_distance,
        const Eigen::Matrix4d &init /* = Eigen::Matrix4d::Identity()*/,
        const TransformationEstimation &estimation
        /* = TransformationEstimationPointToPoint(false)*/,
        const ICPConvergenceCriteria
                &criteria /* = ICPConvergenceCriteria()*/) {
    if (max_correspondence_distance <= 0.0) {
        PrintError("Error: Invalid max_correspondence_distance.\n");
        return RegistrationResult(init);
    }
    if (estimation.GetTransformationEstimationType() ==
                TransformationEstimationType::PointToPlane &&
        (!source.HasNormals() || !target.HasNormals())) {
        PrintError(
                "Error: TransformationEstimationPointToPlane requires "
                "pre-computed normal vectors.\n");
        return RegistrationResult(init);
    }

    Eigen::Matrix4d transformation = init;
    KDTreeFlann kdtree;
    kdtree.SetGeometry(target);
    PointCloud pcd = source;
    if (init.isIdentity() == false) {
        pcd.Transform(init);
    }
    RegistrationResult result;
    result = GetRegistrationResultAndCorrespondences(
            pcd, target, kdtree, max_correspondence_distance, transformation);
    for (int i = 0; i < criteria.max_iteration_; i++) {
        PrintDebug("ICP Iteration #%d: Fitness %.4f, RMSE %.4f\n", i,
                   result.fitness_, result.inlier_rmse_);
        Eigen::Matrix4d update = estimation.ComputeTransformation(
                pcd, target, result.correspondence_set_);
        transformation = update * transformation;
        pcd.Transform(update);
        RegistrationResult backup = result;
        result = GetRegistrationResultAndCorrespondences(
                pcd, target, kdtree, max_correspondence_distance,
                transformation);
        if (std::abs(backup.fitness_ - result.fitness_) <
                    criteria.relative_fitness_ &&
            std::abs(backup.inlier_rmse_ - result.inlier_rmse_) <
                    criteria.relative_rmse_) {
            break;
        }
    }
    return result;
}

RegistrationResult RegistrationRANSACBasedOnCorrespondence(
        const PointCloud &source,
        const PointCloud &target,
        const CorrespondenceSet &corres,
        double max_correspondence_distance,
        const TransformationEstimation &estimation
        /* = TransformationEstimationPointToPoint(false)*/,
        int ransac_n /* = 6*/,
        const RANSACConvergenceCriteria &criteria
        /* = RANSACConvergenceCriteria()*/) {
    if (ransac_n < 3 || (int)corres.size() < ransac_n ||
        max_correspondence_distance <= 0.0) {
        return RegistrationResult();
    }
    std::srand((unsigned int)std::time(0));
    Eigen::Matrix4d transformation;
    CorrespondenceSet ransac_corres(ransac_n);
    RegistrationResult result;
    for (int itr = 0;
         itr < criteria.max_iteration_ && itr < criteria.max_validation_;
         itr++) {
        for (int j = 0; j < ransac_n; j++) {
            ransac_corres[j] = corres[std::rand() % (int)corres.size()];
        }
        transformation =
                estimation.ComputeTransformation(source, target, ransac_corres);
        PointCloud pcd = source;
        pcd.Transform(transformation);
        auto this_result = EvaluateRANSACBasedOnCorrespondence(
                pcd, target, corres, max_correspondence_distance,
                transformation);
        if (this_result.fitness_ > result.fitness_ ||
            (this_result.fitness_ == result.fitness_ &&
             this_result.inlier_rmse_ < result.inlier_rmse_)) {
            result = this_result;
        }
    }
    PrintDebug("RANSAC: Fitness %.4f, RMSE %.4f\n", result.fitness_,
               result.inlier_rmse_);
    return result;
}

RegistrationResult RegistrationRANSACBasedOnFeatureMatching(
        const PointCloud &source,
        const PointCloud &target,
        const Feature &source_feature,
        const Feature &target_feature,
        double max_correspondence_distance,
        const TransformationEstimation &estimation
        /* = TransformationEstimationPointToPoint(false)*/,
        int ransac_n /* = 4*/,
        const std::vector<std::reference_wrapper<const CorrespondenceChecker>>
                &checkers /* = {}*/,
        const RANSACConvergenceCriteria &criteria,
        const double ransac_random_seed, /* default: 0*/
        const int fitness_type /* 0: standard ransac fitness function. 1: use the number of inliers (instead of the ration); 2: use 1/d^2 for inliers, where d is descriptor distance.*/

        /* = RANSACConvergenceCriteria()*/) {
    if (ransac_n < 3 || max_correspondence_distance <= 0.0) {
        return RegistrationResult();
    }

    RegistrationResult result;
    int total_validation = 0;
    bool finished_validation = false;
    int num_similar_features = 1;
    std::vector<std::vector<int>> similar_features(source.points_.size());

#ifdef _OPENMP
#pragma omp parallel
    {
#endif
        CorrespondenceSet ransac_corres(ransac_n);
        KDTreeFlann kdtree(target);
        KDTreeFlann kdtree_feature(target_feature);
        RegistrationResult result_private;
        unsigned int seed_number;
#ifdef _OPENMP
        // each thread has different seed_number
        //seed_number = (unsigned int)std::time(0) * (omp_get_thread_num() + 1);
        seed_number = (unsigned int)ransac_random_seed * (omp_get_thread_num() + 1);
#else
        //seed_number = (unsigned int)std::time(0);
        seed_number = (unsigned int)ransac_random_seed;//std::time(0);
#endif
        std::srand(seed_number);

#ifdef _OPENMP
#pragma omp for nowait
#endif
        for (int itr = 0; itr < criteria.max_iteration_; itr++) {
            if (!finished_validation) {
                std::vector<double> dists(num_similar_features);
                Eigen::Matrix4d transformation;
                for (int j = 0; j < ransac_n; j++) {
                    int source_sample_id =
                            std::rand() % (int)source.points_.size();
                    if (similar_features[source_sample_id].empty()) {
                        std::vector<int> indices(num_similar_features);
                        kdtree_feature.SearchKNN(
                                Eigen::VectorXd(source_feature.data_.col(
                                        source_sample_id)),
                                num_similar_features, indices, dists);
#ifdef _OPENMP
#pragma omp critical
#endif
                        { similar_features[source_sample_id] = indices; }
                    }
                    ransac_corres[j](0) = source_sample_id;
                    if (num_similar_features == 1)
                        ransac_corres[j](1) =
                                similar_features[source_sample_id][0];
                    else
                        ransac_corres[j](1) =
                                similar_features[source_sample_id]
                                                [std::rand() %
                                                 num_similar_features];
                }
                bool check = true;
                for (const auto &checker : checkers) {
                    if (checker.get().require_pointcloud_alignment_ == false &&
                        checker.get().Check(source, target, ransac_corres,
                                            transformation) == false) {
                        check = false;
                        break;
                    }
                }
                if (check == false) continue;
                transformation = estimation.ComputeTransformation(
                        source, target, ransac_corres);
                check = true;
                for (const auto &checker : checkers) {
                    if (checker.get().require_pointcloud_alignment_ == true &&
                        checker.get().Check(source, target, ransac_corres,
                                            transformation) == false) {
                        check = false;
                        break;
                    }
                }
                if (check == false) continue;
                PointCloud pcd = source;
                pcd.Transform(transformation);
                auto this_result = GetRegistrationResultAndCorrespondencesCustom(
                        pcd, target, kdtree, max_correspondence_distance,
                        transformation, source_feature, target_feature, fitness_type);
                if (this_result.fitness_ > result_private.fitness_ ||
                    (this_result.fitness_ == result_private.fitness_ &&
                     this_result.inlier_rmse_ < result_private.inlier_rmse_)) {
                    result_private = this_result;
                }
#ifdef _OPENMP
#pragma omp critical
#endif
                {
                    total_validation = total_validation + 1;
                    if (total_validation >= criteria.max_validation_)
                        finished_validation = true;
                }
            }  // end of if statement
        }      // end of for-loop
#ifdef _OPENMP
#pragma omp critical
#endif
        {
            if (result_private.fitness_ > result.fitness_ ||
                (result_private.fitness_ == result.fitness_ &&
                 result_private.inlier_rmse_ < result.inlier_rmse_)) {
                result = result_private;
            }
        }
#ifdef _OPENMP
    }
#endif
    PrintDebug("total_validation : %d\n", total_validation);
    PrintDebug("RANSAC: Fitness %.4f, RMSE %.4f\n", result.fitness_,
               result.inlier_rmse_);
    return result;
}

RegistrationResult RegistrationRANSACBasedOnShapeComplementarity(
        const PointCloud &source,
        const PointCloud &target,
        const TriangleMesh &target_mesh,
        const PointCloud &target_face_centroids_pcd,
        const Feature &source_feature,
        const Feature &target_feature,
        double max_correspondence_distance,
        const TransformationEstimation &estimation
        /* = TransformationEstimationPointToPoint(false)*/,
        int ransac_n /* = 4*/,
        const std::vector<std::reference_wrapper<const CorrespondenceChecker>>
                &checkers /* = {}*/,
        const RANSACConvergenceCriteria &criteria,
        const double ransac_random_seed, /* default: 0*/
        const int fitness_type /* 0: standard ransac fitness function. 1: use the number of inliers (instead of the ration); 2: use 1/d^2 for inliers, where d is descriptor distance.*/

        /* = RANSACConvergenceCriteria()*/) {
    if (ransac_n < 3 || max_correspondence_distance <= 0.0) {
        return RegistrationResult();
    }

    RegistrationResult result;
    int total_validation = 0;
    bool finished_validation = false;
    int num_similar_features = 1;
    std::vector<std::vector<int>> similar_features(source.points_.size());

#ifdef _OPENMP
#pragma omp parallel
    {
#endif
        CorrespondenceSet ransac_corres(ransac_n);
        // Do a KDTree for the triangles.
        PointCloud target_triangle_centroids();
        KDTreeFlann kdtree(target_face_centroids_pcd);
        KDTreeFlann kdtree_feature(target_feature);
        RegistrationResult result_private;
        unsigned int seed_number;
#ifdef _OPENMP
        // each thread has different seed_number
        //seed_number = (unsigned int)std::time(0) * (omp_get_thread_num() + 1);
        seed_number = (unsigned int)ransac_random_seed * (omp_get_thread_num() + 1);
#else
        //seed_number = (unsigned int)std::time(0);
        seed_number = (unsigned int)ransac_random_seed;//std::time(0);
#endif
        std::srand(seed_number);

#ifdef _OPENMP
#pragma omp for nowait
#endif
        for (int itr = 0; itr < criteria.max_iteration_; itr++) {
            if (!finished_validation) {
                std::vector<double> dists(num_similar_features);
                Eigen::Matrix4d transformation;
                for (int j = 0; j < ransac_n; j++) {
                    int source_sample_id =
                            std::rand() % (int)source.points_.size();
                    if (similar_features[source_sample_id].empty()) {
                        std::vector<int> indices(num_similar_features);
                        kdtree_feature.SearchKNN(
                                Eigen::VectorXd(source_feature.data_.col(
                                        source_sample_id)),
                                num_similar_features, indices, dists);
#ifdef _OPENMP
#pragma omp critical
#endif
                        { similar_features[source_sample_id] = indices; }
                    }
                    ransac_corres[j](0) = source_sample_id;
                    if (num_similar_features == 1)
                        ransac_corres[j](1) =
                                similar_features[source_sample_id][0];
                    else
                        ransac_corres[j](1) =
                                similar_features[source_sample_id]
                                                [std::rand() %
                                                 num_similar_features];
                }
                bool check = true;
                for (const auto &checker : checkers) {
                    if (checker.get().require_pointcloud_alignment_ == false &&
                        checker.get().Check(source, target, ransac_corres,
                                            transformation) == false) {
                        check = false;
                        break;
                    }
                }
                if (check == false) continue;
                transformation = estimation.ComputeTransformation(
                        source, target, ransac_corres);
                check = true;
                for (const auto &checker : checkers) {
                    if (checker.get().require_pointcloud_alignment_ == true &&
                        checker.get().Check(source, target, ransac_corres,
                                            transformation) == false) {
                        check = false;
                        break;
                    }
                }
                if (check == false) continue;
                PointCloud pcd = source;
                pcd.Transform(transformation);
                //PrintInfo("###################\n");
                //PrintInfo("###################\n");
                //PrintInfo("Evaluating transformation\n");
                auto this_result = GetRegistrationResultAndCorrespondencesShapeComplementarity(
                        pcd, target_mesh, kdtree, max_correspondence_distance,
                        transformation, source_feature, target_feature, fitness_type);
                if (this_result.fitness_ > result_private.fitness_ ||
                    (this_result.fitness_ == result_private.fitness_ &&
                     this_result.inlier_rmse_ < result_private.inlier_rmse_)) {
                    result_private = this_result;
                }
#ifdef _OPENMP
#pragma omp critical
#endif
                {
                    total_validation = total_validation + 1;
                    if (total_validation >= criteria.max_validation_)
                        finished_validation = true;
                }
            }  // end of if statement
        }      // end of for-loop
#ifdef _OPENMP
#pragma omp critical
#endif
        {
            if (result_private.fitness_ > result.fitness_ ||
                (result_private.fitness_ == result.fitness_ &&
                 result_private.inlier_rmse_ < result.inlier_rmse_)) {
                result = result_private;
            }
        }
#ifdef _OPENMP
    }
#endif
    PrintDebug("total_validation : %d\n", total_validation);
    PrintDebug("RANSAC: Fitness %.4f, RMSE %.4f\n", result.fitness_,
               result.inlier_rmse_);
    return result;
}

Eigen::Matrix6d GetInformationMatrixFromPointClouds(
        const PointCloud &source,
        const PointCloud &target,
        double max_correspondence_distance,
        const Eigen::Matrix4d &transformation) {
    PointCloud pcd = source;
    if (transformation.isIdentity() == false) {
        pcd.Transform(transformation);
    }
    RegistrationResult result;
    KDTreeFlann target_kdtree(target);
    result = GetRegistrationResultAndCorrespondences(
            pcd, target, target_kdtree, max_correspondence_distance,
            transformation);

    // write q^*
    // see http://redwood-data.org/indoor/registration.html
    // note: I comes first in this implementation
    Eigen::Matrix6d GTG = Eigen::Matrix6d::Identity();
#ifdef _OPENMP
#pragma omp parallel
    {
#endif
        Eigen::Matrix6d GTG_private = Eigen::Matrix6d::Identity();
        Eigen::Vector6d G_r_private = Eigen::Vector6d::Zero();
#ifdef _OPENMP
#pragma omp for nowait
#endif
        for (auto c = 0; c < result.correspondence_set_.size(); c++) {
            int t = result.correspondence_set_[c](1);
            double x = target.points_[t](0);
            double y = target.points_[t](1);
            double z = target.points_[t](2);
            G_r_private.setZero();
            G_r_private(1) = z;
            G_r_private(2) = -y;
            G_r_private(3) = 1.0;
            GTG_private.noalias() += G_r_private * G_r_private.transpose();
            G_r_private.setZero();
            G_r_private(0) = -z;
            G_r_private(2) = x;
            G_r_private(4) = 1.0;
            GTG_private.noalias() += G_r_private * G_r_private.transpose();
            G_r_private.setZero();
            G_r_private(0) = y;
            G_r_private(1) = -x;
            G_r_private(5) = 1.0;
            GTG_private.noalias() += G_r_private * G_r_private.transpose();
        }
#ifdef _OPENMP
#pragma omp critical
#endif
        { GTG += GTG_private; }
#ifdef _OPENMP
    }
#endif
    return std::move(GTG);
}

}  // namespace open3d


/*
 * DynamicalSystem.hpp
 *
 *  Created on: July 7, 2019
 *      Author: Quincy Jones
 *
 * Copyright (c) <2019> <Quincy Jones - quincy@implementedrobotics.com/>
 * Permission is hereby granted, free of charge, to any person obtaining a
 * copy of this software and associated documentation files (the "Software"),
 * to deal in the Software without restriction, including without limitation
 * the rights to use, copy, modify, merge, publish, distribute, sublicense,
 * and/or sell copies of the Software, and to permit persons to whom the Software
 * is furnished to do so, subject to the following conditions:
 * The above copyright notice and this permission notice shall be included in all
 * copies or substantial portions of the Software.
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY#endif // NOMAD_CORE_NOMAD_NOMAD_H,
 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
 * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY,
 * WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
 * CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
 */

#ifndef NOMAD_CORE_SYSTEMS_DYNAMICALSYSTEM_H_
#define NOMAD_CORE_SYSTEMS_DYNAMICALSYSTEM_H_

#include <Eigen/Dense>
#include <unsupported/Eigen/MatrixFunctions>
#include <vector>

// TODO: Need to split all of this into Discrete Time and Continuous Time Systems
// TODO: Also Maybe a SS Container?
class DynamicalSystem
{

public:
    DynamicalSystem(const int& num_states, const int& num_inputs, const double& Ts = 1e-1);

    // Step System
    virtual void Step(const Eigen::VectorXd& u) = 0;

    // Update
    virtual void Update() = 0;

    // Get Current State
    Eigen::VectorXd GetState() const { return x_; }

    // Set Current State
    void SetState(const Eigen::VectorXd &x) { x_ = x; }

    int NumStates() const { return num_states_; }
    int NumInputs() const { return num_inputs_; }

protected:

    Eigen::VectorXd x_;  // State Vector

    int num_states_; // Number of System States
    int num_inputs_;     // Number of inputs

    double T_s_; // Sample Time
    double t_; // Current Time
};

class NonLinearDynamicalSystem : public DynamicalSystem
{

protected:

    virtual void f(const Eigen::MatrixXd& x, const Eigen::MatrixXd& u) = 0;

};

class LinearDynamicalSystem : public DynamicalSystem
{

public:

    LinearDynamicalSystem(const int& num_states, const int& num_inputs, const double& T_s = 1e-1);

    // Model Matrices
    Eigen::MatrixXd A() const { return A_; };
    Eigen::MatrixXd B() const { return B_; };

    // Discrete Time Model Matrices
    Eigen::MatrixXd A_d() const { return A_d_; };
    Eigen::MatrixXd B_d() const { return B_d_; };

protected:

    Eigen::MatrixXd A_;  // State Transition Mat rix
    Eigen::MatrixXd B_;  // Control Input Matrix

    Eigen::MatrixXd A_d_;  // Discrete Time State Transition Mat rix
    Eigen::MatrixXd B_d_;  // Discrete Time Input Matrix

};

class LinearTimeVaryingDynamicalSystem : public LinearDynamicalSystem
{

public:

    // Constructor
    LinearTimeVaryingDynamicalSystem(const int& num_states, const int& num_inputs, const double& T_s = 1e-1);

    // Get the Time-Varying Matrices from Model
    virtual void GetModelMatrices(const int& N, Eigen::MatrixXd& A, Eigen::MatrixXd& B) = 0;

    // Model Matrices (Time Varying)
    std::vector<Eigen::MatrixXd> A_TV() const { return A_tv_; };
    std::vector<Eigen::MatrixXd> B_TV() const { return B_tv_; };

    // Discrete Time Model Matrices (Time Varying)
    std::vector<Eigen::MatrixXd> A_d_TV() const { return A_d_tv_; };
    std::vector<Eigen::MatrixXd> B_d_TV() const { return B_d_tv_; };

protected:

    // Model Matrices (Time Varying)
    std::vector<Eigen::MatrixXd> A_tv_;
    std::vector<Eigen::MatrixXd> B_tv_;

    // Discrete Time Model Matrices (Time Varying)
    std::vector<Eigen::MatrixXd> A_d_tv_;
    std::vector<Eigen::MatrixXd> B_d_tv_;
};

#endif // NOMAD_CORE_SYSTEMS_DYNAMICALSYSTEM_H_


/*
 boost/numeric/odeint/stepper/detail/controlled_adams_bashforth_moulton.hpp

 [begin_description]
 Implemetation of an controlled adams bashforth moulton stepper.
 [end_description]

 Copyright 2017 Valentin Noah Hartmann

 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)
*/

#ifndef BOOST_NUMERIC_ODEINT_STEPPER_CONTROLLED_ADAMS_BASHFORTH_MOULTON_HPP_INCLUDED
#define BOOST_NUMERIC_ODEINT_STEPPER_CONTROLLED_ADAMS_BASHFORTH_MOULTON_HPP_INCLUDED

#include <boost/numeric/odeint/stepper/stepper_categories.hpp>
#include <boost/numeric/odeint/stepper/controlled_step_result.hpp>

#include <boost/numeric/odeint/stepper/adaptive_adams_bashforth_moulton.hpp>
#include <boost/numeric/odeint/stepper/detail/pid_step_adjuster.hpp>

#include <boost/numeric/odeint/util/unwrap_reference.hpp>
#include <boost/numeric/odeint/util/is_resizeable.hpp>
#include <boost/numeric/odeint/util/resizer.hpp>

#include <boost/numeric/odeint/util/copy.hpp>
#include <boost/numeric/odeint/util/bind.hpp>

#include <iostream>

namespace boost {
namespace numeric {
namespace odeint {

template<
size_t MaxOrder,
class State,
class Value = double,
class Algebra = typename algebra_dispatcher< State >::algebra_type
>
class default_order_adjuster
{
public:
    typedef State state_type;
    typedef Value value_type;
    typedef state_wrapper< state_type > wrapped_state_type;

    typedef Algebra algebra_type;

    default_order_adjuster( const algebra_type &algebra = algebra_type() )
    : m_algebra( algebra )
    {};

    size_t adjust_order(size_t order, size_t init, boost::array<wrapped_state_type, 4> &xerr)
    {
        using std::abs;

        value_type errc = abs(m_algebra.norm_inf(xerr[2].m_v));

        value_type errm1 = 3*errc;
        value_type errm2 = 3*errc;

        if(order > 2)
        {
            errm2 = abs(m_algebra.norm_inf(xerr[0].m_v));
        }
        if(order >= 2)
        {
            errm1 = abs(m_algebra.norm_inf(xerr[1].m_v));
        }

        size_t o_new = order;

        if(order == 2 && errm1 <= 0.5*errc)
        {
            o_new = order - 1;
        }
        else if(order > 2 && errm2 < errc && errm1 < errc)
        {
            o_new = order - 1;
        }

        if(init < order)
        {
            return order+1;
        }
        else if(o_new == order - 1)
        {
            return order-1;
        }
        else if(order <= MaxOrder)
        {
            value_type errp = abs(m_algebra.norm_inf(xerr[3].m_v));

            if(order > 1 && errm1 < errc && errp)
            {
                return order-1;
            }
            else if(order < MaxOrder && errp < (0.5-0.25*order/MaxOrder) * errc)
            {
                return order+1;
            }
        }

        return order;
    };
private:
    algebra_type m_algebra;
};

template<
class ErrorStepper,
class StepAdjuster = detail::pid_step_adjuster< typename ErrorStepper::state_type,
    typename ErrorStepper::value_type,
    typename ErrorStepper::deriv_type,
    typename ErrorStepper::time_type,
    typename ErrorStepper::algebra_type,
    typename ErrorStepper::operations_type,
    detail::H211PI
    >,
class OrderAdjuster = default_order_adjuster< ErrorStepper::order_value,
    typename ErrorStepper::state_type,
    typename ErrorStepper::value_type,
    typename ErrorStepper::algebra_type
    >,
class Resizer = initially_resizer
>
class controlled_adams_bashforth_moulton
{
public:
    typedef ErrorStepper stepper_type;

    static const typename stepper_type::order_type order_value = stepper_type::order_value;

    typedef typename stepper_type::state_type state_type;
    typedef typename stepper_type::value_type value_type;
    typedef typename stepper_type::deriv_type deriv_type;
    typedef typename stepper_type::time_type time_type;

    typedef typename stepper_type::algebra_type algebra_type;
    typedef typename stepper_type::operations_type operations_type;
    typedef Resizer resizer_type;

    typedef StepAdjuster step_adjuster_type;
    typedef OrderAdjuster order_adjuster_type;
    typedef controlled_stepper_tag stepper_category;

    typedef typename stepper_type::wrapped_state_type wrapped_state_type;
    typedef typename stepper_type::wrapped_deriv_type wrapped_deriv_type;
    typedef boost::array< wrapped_state_type , 4 > error_storage_type;

    typedef typename stepper_type::coeff_type coeff_type;
    typedef controlled_adams_bashforth_moulton< ErrorStepper , StepAdjuster , OrderAdjuster , Resizer > controlled_stepper_type;

    controlled_adams_bashforth_moulton(step_adjuster_type step_adjuster = step_adjuster_type())
    :m_stepper(),
    m_dxdt_resizer(), m_xerr_resizer(), m_xnew_resizer(),
    m_step_adjuster( step_adjuster ), m_order_adjuster()
    {};

    template< class ExplicitStepper, class System >
    void initialize(ExplicitStepper stepper, System system, state_type &inOut, time_type &t, time_type dt)
    {
        m_stepper.initialize(stepper, system, inOut, t, dt);
    };

    template< class System >
    void initialize(System system, state_type &inOut, time_type &t, time_type dt)
    {
        m_stepper.initialize(system, inOut, t, dt);
    };

    template< class ExplicitStepper, class System >
    void initialize_controlled(ExplicitStepper stepper, System system, state_type &inOut, time_type &t, time_type &dt)
    {
        reset();
        coeff_type &coeff = m_stepper.coeff();

        m_dxdt_resizer.adjust_size( inOut , detail::bind( &controlled_stepper_type::template resize_dxdt_impl< state_type > , detail::ref( *this ) , detail::_1 ) );

        controlled_step_result res = fail;

        for( size_t i=0 ; i<order_value; ++i )
        {
            do
            {
                res = stepper.try_step( system, inOut, t, dt );
            }
            while(res != success);

            system( inOut , m_dxdt.m_v , t );

            coeff.predict(t-dt, dt);
            coeff.do_step(m_dxdt.m_v);
            coeff.confirm();

            if(coeff.m_eo < order_value)
            {
                ++coeff.m_eo;
            }
        }
    }

    template< class System >
    controlled_step_result try_step(System system, state_type & inOut, time_type &t, time_type &dt)
    {
        m_xnew_resizer.adjust_size( inOut , detail::bind( &controlled_stepper_type::template resize_xnew_impl< state_type > , detail::ref( *this ) , detail::_1 ) );

        controlled_step_result res = try_step(system, inOut, t, m_xnew.m_v, dt);

        if(res == success)
        {
            boost::numeric::odeint::copy( m_xnew.m_v , inOut);
        }

        return res;
    };

    template< class System >
    controlled_step_result try_step(System system, const state_type & in, time_type &t, state_type & out, time_type &dt)
    {
        m_xerr_resizer.adjust_size( in , detail::bind( &controlled_stepper_type::template resize_xerr_impl< state_type > , detail::ref( *this ) , detail::_1 ) );
        m_dxdt_resizer.adjust_size( in , detail::bind( &controlled_stepper_type::template resize_dxdt_impl< state_type > , detail::ref( *this ) , detail::_1 ) );

        m_stepper.do_step_impl(system, in, t, out, dt, m_xerr[2].m_v);

        coeff_type &coeff = m_stepper.coeff();

        time_type dtPrev = dt;
        dt = m_step_adjuster.adjust_stepsize(coeff.m_eo, dt, m_xerr[2].m_v, out, m_stepper.dxdt() );

        if(dt / dtPrev >= step_adjuster_type::threshold())
        {
            system(out, m_dxdt.m_v, t+dtPrev);

            coeff.do_step(m_dxdt.m_v);
            coeff.confirm();

            t += dtPrev;

            size_t eo = coeff.m_eo;

            // estimate errors for next step
            double factor = 1;
            algebra_type m_algebra;

            m_algebra.for_each2(m_xerr[2].m_v, coeff.phi[1][eo].m_v,
                typename operations_type::template scale_sum1<double>(factor*dt*(coeff.gs[eo])));

            if(eo > 1)
            {
                m_algebra.for_each2(m_xerr[1].m_v, coeff.phi[1][eo-1].m_v,
                    typename operations_type::template scale_sum1<double>(factor*dt*(coeff.gs[eo-1])));
            }
            if(eo > 2)
            {
                m_algebra.for_each2(m_xerr[0].m_v, coeff.phi[1][eo-2].m_v,
                    typename operations_type::template scale_sum1<double>(factor*dt*(coeff.gs[eo-2])));
            }
            if(eo < order_value && coeff.m_eo < coeff.m_steps_init-1)
            {
                m_algebra.for_each2(m_xerr[3].m_v, coeff.phi[1][eo+1].m_v,
                    typename operations_type::template scale_sum1<double>(factor*dt*(coeff.gs[eo+1])));
            }

            // adjust order
            coeff.m_eo = m_order_adjuster.adjust_order(coeff.m_eo, coeff.m_steps_init-1, m_xerr);

            return success;
        }
        else
        {
            return fail;
        }
    };

    void reset() { m_stepper.reset(); };

private:
    template< class StateType >
    bool resize_dxdt_impl( const StateType &x )
    {
        return adjust_size_by_resizeability( m_dxdt, x, typename is_resizeable<deriv_type>::type() );
    };
    template< class StateType >
    bool resize_xerr_impl( const StateType &x )
    {
        bool resized( false );

        for(size_t i=0; i<m_xerr.size(); ++i)
        {
            resized |= adjust_size_by_resizeability( m_xerr[i], x, typename is_resizeable<state_type>::type() );
        }
        return resized;
    };
    template< class StateType >
    bool resize_xnew_impl( const StateType &x )
    {
        return adjust_size_by_resizeability( m_xnew, x, typename is_resizeable<state_type>::type() );
    };

    stepper_type m_stepper;

    wrapped_deriv_type m_dxdt;
    error_storage_type m_xerr;
    wrapped_state_type m_xnew;

    resizer_type m_dxdt_resizer;
    resizer_type m_xerr_resizer;
    resizer_type m_xnew_resizer;

    step_adjuster_type m_step_adjuster;
    order_adjuster_type m_order_adjuster;
};

} // odeint
} // numeric
} // boost

#endif


/*****************************************************************************
*
* Functions for calculating the gap
*
* Copyright (C) 2018 by Hidemaro Suwa
* e-mail:suwamaro@phys.s.u-tokyo.ac.jp
*
*****************************************************************************/

// #include <functional>
#include <armadillo>
#include <cuba.h>
#include "calc_gap.h"
#include "rpa_util.h"
#include "calc_intensity.h"

/* Index: NCOMP = (ope * nelem + elem) | (real=0 or imag=1) */
std::tuple<int, int> comp_to_ope_elem(int comp, int nelem){
  comp >>= 1;
  int ope = comp / nelem;
  int elem = comp % nelem;
  return std::make_tuple(ope, elem);
}

/* Member functions of ResponseFuncIntegrand */
void ResponseFuncIntegrand::update_parameters(double _T, double _delta, double _mu, cx_double _omega, Polarization const& _Pz){
  T_ = _T;
  delta_ = _delta;
  mu_ = _mu;
  omega_ = _omega;
  Pz_ = _Pz;  // without copying the tables
  Pz_.is_table_set_ = false;
};

/* Member functions of ResponseFuncIntegrandBilayer */
void ResponseFuncIntegrandBilayer::set_parameters(hoppings_bilayer2 const& h, double T, double delta, double mu, cx_double omega, Polarization const& Pz){
  hb_ = h;
  ResponseFuncIntegrand::update_parameters(T, delta, mu, omega, Pz);
}

int ResponseFuncIntegrandBilayer::calc(const int *ndim, const cubareal xx[], const int *ncomp, cubareal ff[], void *userdata) const {
   /* Reset */
   for(int comp=0; comp<*ncomp; comp++){ ff[comp] = 0; }

   /* Wavenumbers */
   double k1 = xx[0] * 2 * M_PI;
   double k2 = xx[1] * 2 * M_PI;

   double kx = 0.5 * (k2 + k1);
   double ky = 0.5 * (k2 - k1);

   /* Polarizaiton */
   cx_double Ppmk[NSUBL*NSUBL];
   cx_double Pzzk[NSUBL*NSUBL];

   /* Sum over kz */
   for(int z=0; z < 2; z++){
     double kz = M_PI * z;

     /* Sum over bands (signs) */
     for(int sg1=-1; sg1<=1; sg1+=2){
       for(int sg2=-1; sg2<=1; sg2+=2){
	 // int sg2 = - sg1; /* Opposite sign */
	 cx_double prefactor = calc_prefactor_bare_res_func_bilayer(sg1, sg2, *ts(), T(), kx, ky, kz, Pz()->qx(), Pz()->qy(), Pz()->qz(), omega(), delta(), mu());

	 /* Getting the polarization */
	 Pz()->calc_polarization(hb_, delta(), kx, ky, kz, sg1, sg2, Ppmk, Pzzk);

	 /* Integrand */
	 for(int comp=0; comp<*ncomp; comp+=2){
	   int ope, elem;
	   std::tie(ope, elem) = comp_to_ope_elem(comp, NSUBL*NSUBL);
	   if ( ope == 0 ) {  /* Ppmk */
	     ff[comp] += std::real(prefactor * Ppmk[elem]);
	     ff[comp+1] += std::imag(prefactor * Ppmk[elem]);
	   } else /* ope == 1 */ { /* Pzzk */
	     ff[comp] += std::real(prefactor * Pzzk[elem]);
	     ff[comp+1] += std::imag(prefactor * Pzzk[elem]);
	   }
	 }
       }
     }
   }

   return 0;
}

cx_double calc_intensity_bilayer(int L, hoppings const& ts, double mu, double U, double delta, double qx, double qy, double qz, cx_double omega, bool zz){
  double k1 = 2. * M_PI / (double)L;

  cx_double A = 0, B = 0, C = 0, D = 0;
  for(int z=0; z < 2; z++){
    double kz = M_PI * z;
    for(int x=-L/2; x < L/2; x++){
      double kx = k1 * x;
      for(int y=-L/2; y < L/2; y++){
	double ky = k1 * y;
	add_to_sus_mat2( ts, mu, A, B, C, D, qx, qy, qz, kx, ky, kz, delta, omega, zz );
      }
    }
  }

  int n_sites = L * L * 2;
  double norm = 2. / (double)n_sites;
  A *= norm;
  B *= norm;
  C *= norm;
  D *= norm;

  /* RPA */
  arma::cx_mat chi0_mat(2,2);
  chi0_mat(0,0) = A;   // (A, A) correlation
  chi0_mat(0,1) = B;   // (A, B)
  chi0_mat(1,0) = C;   // (B, A)
  chi0_mat(1,1) = D;   // (B, B)

  /* Transverse = < \sigma^- \sigma^+ >; Longitudinal (zz) = < \sigma^z \sigma^z > */
  /* Note that 2 < \sigma^- \sigma^+ > = < \sigma^z \sigma^z > (U -> 0 for the SU(2) case) */
  double factor_channel = 1.0;
  if ( zz ) { factor_channel = 0.5; }

  arma::cx_mat denom = arma::eye<arma::cx_mat>(NSUBL,NSUBL) - factor_channel * U * chi0_mat;
  arma::cx_mat chi_mat = chi0_mat * arma::inv(denom);

  // sigma-to-spin factor
  double factor_operator = 0.5;
  cx_double chi = factor_operator * factor_operator * ( chi_mat(0,0) - chi_mat(1,0) - chi_mat(0,1) + chi_mat(1,1) );

  // // for check
  // std::cout << chi0_mat << std::endl;
  // std::cout << chi_mat << std::endl;

  return chi;
}

/* Instantiation */
ResponseFuncIntegrandBilayer rfib;

int integrand_wrapper(const int *ndim, const cubareal xx[], const int *ncomp, cubareal ff[], void *userdata){
  return rfib.calc(ndim, xx, ncomp, ff, userdata);
}

std::tuple<arma::cx_mat, arma::cx_mat> calc_bare_response_bilayer(int L, hoppings_bilayer2 const& ts, double mu, double U, double T, double delta, CubaParam const& cbp, Polarization const& Pz, cx_double omega, bool continuous_k){
  arma::cx_mat chi0_pm(NSUBL,NSUBL,arma::fill::zeros);
  arma::cx_mat chi0_zz_u(NSUBL,NSUBL,arma::fill::zeros);

  if ( continuous_k ) {
    /* Changing the parameters */
    rfib.set_parameters(ts, T, delta, mu, omega, Pz);

    /* For Cuba */
    int nregions, neval, fail;
    cubareal integral[cbp.NCOMP], error[cbp.NCOMP], prob[cbp.NCOMP];

    /* Cuhre */
    Cuhre(cbp.NDIM, cbp.NCOMP, integrand_wrapper, cbp.userdata, cbp.nvec, cbp.epsrel, cbp.epsabs, cbp.flags, cbp.mineval, cbp.maxeval, cbp.key, cbp.statefile, cbp.spin, &nregions, &neval, &fail, integral, error, prob);

    // for check
    printf("CUHRE RESULT:\tnregions %d\tneval %d\tfail %d\n", nregions, neval, fail);
    // for(int comp = 0; comp < cbp.NCOMP; comp++ )
    //   printf("CUHRE RESULT:\t%.8f +- %.8f\tp = %.3f\n",
    // 	     (double)integral[comp], (double)error[comp], (double)prob[comp]);

    /* Extracting the integral results */
    for(int comp=0; comp<cbp.NCOMP; comp+=2){
      int ope, elem;
      std::tie(ope, elem) = comp_to_ope_elem(comp, NSUBL*NSUBL);
      int g1 = elem / NSUBL;
      int g2 = elem % NSUBL;
      cx_double int_ope_elem(integral[comp], integral[comp+1]);
      if ( ope == 0 ) {  /* Ppmk */
	chi0_pm(g1,g2) = int_ope_elem;
      } else /* ope == 1 */ { /* Pzzk */
	chi0_zz_u(g1,g2) = int_ope_elem;
      }
    }

    chi0_pm /= 2; /* A factor (2*M_PI) cancels because of the scale change of the integration variables. */
    chi0_zz_u /= 2;
  } else { /* Integral for a finite-size system */
    double k1 = 2. * M_PI / L;
    for(int z=0; z < 2; z++){
      double kz = M_PI * z;
      for(int x=-L/2; x < L/2; x++){
	double kx = k1 * x;
	for(int y=-L/2; y < L/2; y++){
	  double ky = k1 * y;
	  add_to_sus_mat4( ts, T, mu, chi0_pm, chi0_zz_u, kx, ky, kz, Pz, delta, omega );
	}
      }
    }

    int n_units = L * L;  // Number of unit cells
    chi0_pm /= (double)(n_units);
    chi0_zz_u /= (double)(n_units);
  }

  /* Adding the contribution from the down spin */
  arma::cx_mat chi0_zz_d(chi0_zz_u);
  /* Assume NSUBL == 2. */
  if ( NSUBL == 2 ) {
    chi0_zz_d.swap_rows( 0, 1 );
    chi0_zz_d.swap_cols( 0, 1 );
  } else {
    std::cerr << "NSUBL is not 2.\n";
    std::exit(EXIT_FAILURE);
  }
  arma::cx_mat chi0_zz = chi0_zz_u + chi0_zz_d;
  return std::make_tuple(chi0_pm, chi0_zz);
}

std::tuple<cx_double, cx_double> calc_intensity_bilayer2(int L, hoppings_bilayer2& ts, double mu, double U, double T, double delta, CubaParam const& cbp, Polarization const& Pz, cx_double omega, bool continuous_k){
  arma::cx_mat chi0_pm(NSUBL,NSUBL,arma::fill::zeros);
  arma::cx_mat chi0_zz(NSUBL,NSUBL,arma::fill::zeros);

  /* Calculating the bare response functions */
  std::tie(chi0_pm, chi0_zz) = calc_bare_response_bilayer(L, ts, mu, U, T, delta, cbp, Pz, omega, continuous_k);

  /* RPA */
  /* Transverse = < \sigma^- \sigma^+ >; Longitudinal (zz) = < \sigma^z \sigma^z > */
  /* Note that 2 < \sigma^- \sigma^+ > = < \sigma^z \sigma^z > (U -> 0 for the SU(2) case) */
  arma::cx_mat denom_pm = arma::eye<arma::cx_mat>(NSUBL,NSUBL) - U * chi0_pm;
  arma::cx_mat denom_zz = arma::eye<arma::cx_mat>(NSUBL,NSUBL) - 0.5 * U * chi0_zz;
  arma::cx_mat chi_pm = arma::inv(denom_pm) * chi0_pm;
  arma::cx_mat chi_zz = arma::inv(denom_zz) * chi0_zz;

  // sigma-to-spin factor
  cx_double chi_xy = 0.5 * arma::accu(chi_pm);
  cx_double chi_z = 0.25 * arma::accu(chi_zz);

  // // for check
  // std::cout << chi0_pm << std::endl;
  // std::cout << chi_pm << std::endl;
  // std::cout << chi0_zz << std::endl;
  // std::cout << chi_zz << std::endl;

  return std::make_tuple(chi_xy, chi_z);
}