#define BOOST_TEST_MODULE "test_afm_fitting_interaction"

#ifdef BOOST_TEST_DYN_LINK
#include <boost/test/unit_test.hpp>
#else
#include <boost/test/included/unit_test.hpp>
#endif

#include <mjolnir/core/BoundaryCondition.hpp>
#include <mjolnir/core/SimulatorTraits.hpp>
#include <mjolnir/forcefield/AFMFit/AFMFitInteraction.hpp>
#include <mjolnir/util/make_unique.hpp>

#include <random>

namespace test
{
template<typename realT, typename coordT>
std::vector<realT>
make_image(const std::vector<realT>& radii, const std::vector<coordT>& positions,
           const realT pixel_width, const realT sigma, const realT gamma,
           const std::size_t n_x, const std::size_t n_y)
{
    std::vector<realT> image; image.reserve(n_x * n_y);
    for(std::size_t i_y=0; i_y < n_y; ++i_y)
    {
        const realT y0 = pixel_width * (i_y + 0.5);
        for(std::size_t i_x=0; i_x < n_x; ++i_x)
        {
            const realT x0 = pixel_width * (i_x + 0.5);

            realT sumexp = 1.0; // exp(0)
            for(std::size_t i=0; i<positions.size(); ++i)
            {
                const auto& p = positions.at(i);
                const auto& r = radii.at(i);
                const realT dx = mjolnir::math::X(p) - x0;
                const realT dy = mjolnir::math::Y(p) - y0;
                const realT dz = mjolnir::math::Z(p) - 0.0;

                const realT gx = -0.5 * dx * dx / (sigma * sigma);
                const realT gy = -0.5 * dy * dy / (sigma * sigma);
                const realT hz = (dz + r) / gamma;

                sumexp += std::exp(gx + gy + hz);
            }
            image.push_back(gamma * std::log(sumexp));
        }
    }
    assert(image.size() == n_x * n_y);
    return image;
}
}

BOOST_AUTO_TEST_CASE(AFMFitting_calc_force)
{
    mjolnir::LoggerManager::set_default_logger("test_afm_fitting_interaction.log");

    using traits_type      = mjolnir::SimulatorTraits<double, mjolnir::UnlimitedBoundary>;
    using real_type        = traits_type::real_type;
    using coord_type       = traits_type::coordinate_type;
    using boundary_type    = traits_type::boundary_type;
    using system_type      = mjolnir::System<traits_type>;
    using interaction_type = mjolnir::AFMFitInteraction<traits_type>;

    const std::size_t num_particles = 5;

    const real_type   k        =  5.0;
    const real_type   gamma    =  1.0;
    const real_type   z0       =  0.0;
    const real_type   cutoff   =  5.01;
    const real_type   margin   =  0.5;
    const real_type   sigma_x  =  2.0;
    const real_type   sigma_y  = sigma_x;
    const real_type   pixel_x  = 10.0;
    const real_type   pixel_y  = pixel_x;
    const std::size_t length_x = 10u;
    const std::size_t length_y = 10u;
    const std::vector<std::pair<std::size_t, real_type>> radii = {
        {0, 1.0}, {1, 2.0}, {2, 3.0}, {3, 4.0}, {4, 5.0}
    };

    std::vector<coord_type> initial_coordinates = {
        coord_type(20.0, 20.0, 10.0),
        coord_type(50.0, 50.0, 15.0),
        coord_type(80.0, 80.0,  5.0),
        coord_type(20.0, 50.0, 20.0),
        coord_type(50.0, 80.0, 10.0)
    };

    const std::vector<real_type> image = test::make_image(
            {1.0, 2.0, 3.0, 4.0, 5.0}, initial_coordinates,
            pixel_x, sigma_x, gamma, length_x, length_y);

    interaction_type interaction(k, gamma, z0, cutoff, margin,
            sigma_x, sigma_y, pixel_x, pixel_y, length_x, length_y,
            radii, image);

    std::mt19937 mt(123456789);
    std::uniform_real_distribution<real_type> uni(-0.01, 0.01);

    for(std::size_t trial = 0; trial < 1000; ++trial)
    {
        system_type sys(num_particles, boundary_type{});

        for(std::size_t i=0; i<sys.size(); ++i)
        {
            sys.mass(i)  = 1.0;
            sys.rmass(i) = 1.0;

            sys.position(i)  = initial_coordinates.at(i);
            sys.position(i) += coord_type(uni(mt), uni(mt), uni(mt));
            sys.velocity(i)  = coord_type(0.0, 0.0, 0.0);
            sys.force(i)     = coord_type(0.0, 0.0, 0.0);
            sys.name(i)      = "X";
            sys.group(i)     = "TEST";
        }

        const auto init = sys;
        interaction.initialize(init);

        constexpr real_type tol = 1e-3;
        constexpr real_type dr  = 1e-4;
        for(std::size_t idx=0; idx<sys.size(); ++idx)
        {
            {
                // ----------------------------------------------------------------
                // reset positions
                sys = init;

                // calc U(x-dx)
                const auto E0 = interaction.calc_energy(sys);

                mjolnir::math::X(sys.position(idx)) += dr;

                // calc F(x)
                interaction.calc_force(sys);

                mjolnir::math::X(sys.position(idx)) += dr;

                // calc U(x+dx)
                const auto E1 = interaction.calc_energy(sys);

                // central difference
                const auto dE = (E1 - E0) * 0.5;

                BOOST_TEST(-dE == dr * mjolnir::math::X(sys.force(idx)),
                           boost::test_tools::tolerance(tol));
            }
            {
                // ----------------------------------------------------------------
                // reset positions
                sys = init;

                // calc U(x-dx)
                const auto E0 = interaction.calc_energy(sys);

                mjolnir::math::Y(sys.position(idx)) += dr;

                // calc F(x)
                interaction.calc_force(sys);

                mjolnir::math::Y(sys.position(idx)) += dr;

                // calc U(x+dx)
                const auto E1 = interaction.calc_energy(sys);

                // central difference
                const auto dE = (E1 - E0) * 0.5;

                BOOST_TEST(-dE == dr * mjolnir::math::Y(sys.force(idx)),
                           boost::test_tools::tolerance(tol));
            }
            {
                // ----------------------------------------------------------------
                // reset positions
                sys = init;

                // calc U(x-dx)
                const auto E0 = interaction.calc_energy(sys);

                mjolnir::math::Z(sys.position(idx)) += dr;

                // calc F(x)
                interaction.calc_force(sys);

                mjolnir::math::Z(sys.position(idx)) += dr;

                // calc U(x+dx)
                const auto E1 = interaction.calc_energy(sys);

                // central difference
                const auto dE = (E1 - E0) * 0.5;

                BOOST_TEST(-dE == dr * mjolnir::math::Z(sys.force(idx)),
                           boost::test_tools::tolerance(tol));
            }
        }
    }
}

BOOST_AUTO_TEST_CASE(AFMFitting_calc_force_and_energy)
{
    mjolnir::LoggerManager::set_default_logger("test_afm_fitting_interaction.log");

    using traits_type      = mjolnir::SimulatorTraits<double, mjolnir::UnlimitedBoundary>;
    using real_type        = traits_type::real_type;
    using coord_type       = traits_type::coordinate_type;
    using boundary_type    = traits_type::boundary_type;
    using system_type      = mjolnir::System<traits_type>;
    using interaction_type = mjolnir::AFMFitInteraction<traits_type>;

    const std::size_t num_particles = 5;

    const real_type   k        =  5.0;
    const real_type   gamma    =  1.0;
    const real_type   z0       =  0.0;
    const real_type   cutoff   =  5.01;
    const real_type   margin   =  0.5;
    const real_type   sigma_x  =  2.0;
    const real_type   sigma_y  = sigma_x;
    const real_type   pixel_x  = 10.0;
    const real_type   pixel_y  = pixel_x;
    const std::size_t length_x = 10u;
    const std::size_t length_y = 10u;
    const std::vector<std::pair<std::size_t, real_type>> radii = {
        {0, 1.0}, {1, 2.0}, {2, 3.0}, {3, 4.0}, {4, 5.0}
    };

    std::vector<coord_type> initial_coordinates = {
        coord_type(20.0, 20.0, 10.0),
        coord_type(50.0, 50.0, 15.0),
        coord_type(80.0, 80.0,  5.0),
        coord_type(20.0, 50.0, 20.0),
        coord_type(50.0, 80.0, 10.0)
    };

    const std::vector<real_type> image = test::make_image(
            {1.0, 2.0, 3.0, 4.0, 5.0}, initial_coordinates,
            pixel_x, sigma_x, gamma, length_x, length_y);

    interaction_type interaction(k, gamma, z0, cutoff, margin,
            sigma_x, sigma_y, pixel_x, pixel_y, length_x, length_y,
            radii, image);

    std::mt19937 mt(123456789);
    std::uniform_real_distribution<real_type> uni(-0.01, 0.01);

    for(std::size_t trial = 0; trial < 1000; ++trial)
    {
        system_type sys(num_particles, boundary_type{});

        for(std::size_t i=0; i<sys.size(); ++i)
        {
            sys.mass(i)  = 1.0;
            sys.rmass(i) = 1.0;

            sys.position(i)  = initial_coordinates.at(i);
            sys.position(i) += coord_type(uni(mt), uni(mt), uni(mt));
            sys.velocity(i)  = coord_type(0.0, 0.0, 0.0);
            sys.force(i)     = coord_type(0.0, 0.0, 0.0);
            sys.name(i)      = "X";
            sys.group(i)     = "TEST";
        }

        interaction.initialize(sys);
        system_type ref_sys = sys;

        constexpr real_type tol = 1e-3;

        const auto energy = interaction.calc_force_and_energy(sys);
        const auto ref_energy = interaction.calc_energy(ref_sys);
        interaction.calc_force(ref_sys);
        BOOST_TEST(ref_energy == energy, boost::test_tools::tolerance(tol));

        for(std::size_t idx=0; idx<sys.size(); ++idx)
        {
            BOOST_TEST(mjolnir::math::X(sys.force(idx)) == mjolnir::math::X(ref_sys.force(idx)), boost::test_tools::tolerance(tol));
            BOOST_TEST(mjolnir::math::Y(sys.force(idx)) == mjolnir::math::Y(ref_sys.force(idx)), boost::test_tools::tolerance(tol));
            BOOST_TEST(mjolnir::math::Z(sys.force(idx)) == mjolnir::math::Z(ref_sys.force(idx)), boost::test_tools::tolerance(tol));
        }
    }
}


#include "gtest/gtest.h"

#include <Configuration.h>
#include <Eigen/Dense>
#include <Utils/IO/IOUtilities.hpp>
#include <PnC/RobotSystem/RobotSystem.hpp>
#include <PnC/WBC/IHWBC/IHWBC.hpp>
#include <PnC/WBC/IHWBC/IHWBC_JointIntegrator.hpp>

TEST(IHWBC_JOINT_INTEGRATION, constructor_test) {
	int num_joints = 10;
	double dt = 1e-3;
	IHWBC_JointIntegrator joint_integrator(num_joints, dt);

	// Leaky Integrator frequency. Much lower than usual cut-offs.
	joint_integrator.setVelocityFrequencyCutOff(2.0);
	joint_integrator.setPositionFrequencyCutOff(1.0);
	joint_integrator.setMaxPositionError(0.2);
	joint_integrator.printIntegrationParams();

	// Initialize current velocity and positions
	Eigen::VectorXd vel = Eigen::VectorXd::Zero(num_joints);
	Eigen::VectorXd pos = M_PI/4.0*Eigen::VectorXd::Ones(num_joints);

	joint_integrator.initializeStates(vel, pos);

	// Initialize output containers
	Eigen::VectorXd vel_out = Eigen::VectorXd::Zero(num_joints);
	Eigen::VectorXd pos_out = Eigen::VectorXd::Zero(num_joints);

	// Initialize Output variables
	// Set joint acceleration
	Eigen::VectorXd acc = 2.0*Eigen::VectorXd::Ones(num_joints);

	myUtils::pretty_print(acc, std::cout, "acc");
	myUtils::pretty_print(vel, std::cout, "vel");
	myUtils::pretty_print(pos, std::cout, "pos");

	joint_integrator.integrate(acc, vel, pos, vel_out, pos_out);

	myUtils::pretty_print(vel_out, std::cout, "vel_out");
	myUtils::pretty_print(pos_out, std::cout, "pos_out");

}


#include <iostream>
#include <iomanip>
#include <boost/filesystem.hpp>
#include <nanikanizer/nanikanizer.hpp>
#include "game.hpp"
#include "enumerate_action.hpp"
#include "apply_state.hpp"
#include "random_state_generator.hpp"

namespace shogi
{

	class shogi_ai
	{
	public:

		shogi_ai()
			: l1_(input_size, 1500)
			, l2_(1500, 1000)
			, l3_(1000, 750)
			, l4_(750, 500)
			, l5_(500, 250)
			, l6_(250, 1)
			, opt_(0.0001)
		{
			opt_.add_parameter(l1_);
			opt_.add_parameter(l2_);
			opt_.add_parameter(l3_);
			opt_.add_parameter(l4_);
			opt_.add_parameter(l5_);
			opt_.add_parameter(l6_);
		}

		float evaluate(const game& g)
		{
			auto e = forward(g);
			nnk::evaluator<float> ev(e);
			ev.forward();

			return e.root()->output()[0];
		}

		action step(const game& g)
		{
			BOOST_ASSERT(!g.state().is_finished());

			std::vector<action> actions;
			enumerate_action(g, std::back_inserter(actions));

			std::vector<float> values(actions.size());

			#pragma omp parallel for schedule(dynamic)
			for (std::ptrdiff_t i = 0; i < static_cast<std::ptrdiff_t>(actions.size()); ++i)
			{
				auto next_g = g;
				next_g.apply(actions[i]);

				values[i] = -evaluate(next_g);
			}

			std::size_t maximum_index = std::max_element(values.begin(), values.end()) - values.begin();
			return actions[maximum_index];
		}

		void zero_grads()
		{
			opt_.zero_grads();
		}

		float train(const game& g1, const game& g2)
		{
			auto e1 = forward(g1);
			auto e2 = forward(g2);

			auto loss = nnk::square(e1 + e2);
			nnk::evaluator<float> ev(loss);
			ev.forward();
			ev.backward();

			return loss.root()->output()[0];
		}

		void update()
		{
			opt_.update();
		}

		void save(const boost::filesystem::path& filename)
		{
			std::ofstream os;
			os.exceptions(std::ios::badbit | std::ios::failbit | std::ios::eofbit);
			os.open(filename.c_str(), std::ios::binary);

			nnk::binary_writer writer(os);
			l1_.save(writer);
			l2_.save(writer);
			l3_.save(writer);
			l4_.save(writer);
			l5_.save(writer);
			l6_.save(writer);
		}

		void load(const boost::filesystem::path& filename)
		{
			std::ifstream is;
			is.exceptions(std::ios::badbit | std::ios::failbit | std::ios::eofbit);
			is.open(filename.c_str(), std::ios::binary);

			nnk::binary_reader reader(is);
			l1_.load(reader);
			l2_.load(reader);
			l3_.load(reader);
			l4_.load(reader);
			l5_.load(reader);
			l6_.load(reader);
		}

	private:

		nnk::expression<float> forward(const game& g)
		{
			if (g.state().is_finished())
			{
				if (g.turn() ^ g.state().is_win())
					return { 1.0 };
				else
					return { -1.0 };
			}

			nnk::variable<float> input;

			input.value().resize(input_size);
			apply_state(&input.value()[0], g.state(), g.turn());

			auto e = input.expr();
			e = nnk::relu(l1_.forward(e));
			e = nnk::relu(l2_.forward(e));
			e = nnk::relu(l3_.forward(e));
			e = nnk::relu(l4_.forward(e));
			e = nnk::relu(l5_.forward(e));
			e = nnk::tanh(l6_.forward(e));

			return e;
		}

		nnk::linear_layer<float> l1_;
		nnk::linear_layer<float> l2_;
		nnk::linear_layer<float> l3_;
		nnk::linear_layer<float> l4_;
		nnk::linear_layer<float> l5_;
		nnk::linear_layer<float> l6_;

		nnk::adam_optimizer opt_;

	};

}

int main(int /*argc*/, char* /*argv*/[])
{
	try
	{
		std::cout << std::fixed << std::setprecision(3);

		shogi::shogi_ai ai;

		boost::filesystem::path filename = "shogi.dat";

		if (boost::filesystem::exists(filename))
			ai.load(filename);

		shogi::random_state_generator gen;

		while (true)
		{
			gen.randomize();

			shogi::game g = gen.get();

			std::vector<shogi::game> states;

			for (std::size_t i = 0; i < 1000; ++i)
			{
				float value = ai.evaluate(g);

				if (g.turn())
					std::cout << value;

				std::cout << std::endl;

				std::cout << g.state();

				if (!g.turn())
					std::cout << value;

				std::cout << std::endl;

				states.push_back(g);

				if (g.state().is_finished())
					break;

				auto next_g = g;
				next_g.apply(ai.step(g));

				g = next_g;
			}

			float sum = 0.0f;

			ai.zero_grads();

			for (std::size_t i = 0; i < states.size() - 1; ++i)
				sum += ai.train(states[i], states[i + 1]);

			ai.update();

			std::cout << "loss: " << sum << std::endl;

			ai.save(filename);
		}
	}
	catch (std::exception& e)
	{
		std::cerr << e.what() << std::endl;
		return -1;
	}

	return 0;
}


// Copyright 2015-2020 Tier IV, Inc. All rights reserved.
//
// 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.

#include <algorithm>
#include <behavior_tree_plugin/vehicle/behavior_tree.hpp>
#include <behavior_tree_plugin/vehicle/lane_change_action.hpp>
#include <boost/algorithm/clamp.hpp>
#include <memory>
#include <scenario_simulator_exception/exception.hpp>
#include <string>
#include <traffic_simulator/math/catmull_rom_spline.hpp>
#include <traffic_simulator/math/transform.hpp>
#include <vector>

namespace entity_behavior
{
namespace vehicle
{
LaneChangeAction::LaneChangeAction(const std::string & name, const BT::NodeConfiguration & config)
: entity_behavior::VehicleActionNode(name, config), current_s_(0)
{
}

const boost::optional<traffic_simulator_msgs::msg::Obstacle> LaneChangeAction::calculateObstacle(
  const traffic_simulator_msgs::msg::WaypointsArray &)
{
  return boost::none;
}

const traffic_simulator_msgs::msg::WaypointsArray LaneChangeAction::calculateWaypoints()
{
  if (!curve_) {
    THROW_SIMULATION_ERROR("curve is null");
  }
  if (!lane_change_parameters_) {
    THROW_SIMULATION_ERROR("lane change parameter is null");
  }
  if (entity_status.action_status.twist.linear.x >= 0) {
    traffic_simulator_msgs::msg::WaypointsArray waypoints;
    double horizon =
      boost::algorithm::clamp(entity_status.action_status.twist.linear.x * 5, 20, 50);
    auto following_lanelets =
      hdmap_utils->getFollowingLanelets(lane_change_parameters_->target.lanelet_id, 0);
    double l = curve_->getLength();
    double rest_s = current_s_ + horizon - l;
    if (rest_s < 0) {
      const auto curve_waypoints =
        curve_->getTrajectory(current_s_, current_s_ + horizon, 1.0, true);
      waypoints.waypoints = curve_waypoints;
    } else {
      std::vector<geometry_msgs::msg::Point> center_points =
        hdmap_utils->getCenterPoints(following_lanelets);
      // DIFFERENT SPLINE - recalculation needed
      traffic_simulator::math::CatmullRomSpline spline(center_points);
      const auto straight_waypoints = spline.getTrajectory(target_s_, target_s_ + rest_s, 1.0);
      waypoints.waypoints = straight_waypoints;
      const auto curve_waypoints = curve_->getTrajectory(current_s_, l, 1.0, true);
      std::copy(
        straight_waypoints.begin(), straight_waypoints.end(),
        std::back_inserter(waypoints.waypoints));
    }
    return waypoints;
  } else {
    return traffic_simulator_msgs::msg::WaypointsArray();
  }
}

void LaneChangeAction::getBlackBoardValues()
{
  VehicleActionNode::getBlackBoardValues();
  traffic_simulator::lane_change::Parameter lane_change_parameters;
  if (!getInput<traffic_simulator::lane_change::Parameter>(
        "lane_change_parameters", lane_change_parameters)) {
    lane_change_parameters_ = boost::none;
  } else {
    lane_change_parameters_ = lane_change_parameters;
  }
}

BT::NodeStatus LaneChangeAction::tick()
{
  getBlackBoardValues();
  if (request != "lane_change") {
    curve_ = boost::none;
    current_s_ = 0;
    return BT::NodeStatus::FAILURE;
  }
  if (!lane_change_parameters_) {
    curve_ = boost::none;
    current_s_ = 0;
    return BT::NodeStatus::FAILURE;
  }
  if (!curve_) {
    if (request == "lane_change") {
      if (!hdmap_utils->canChangeLane(
            entity_status.lanelet_pose.lanelet_id, lane_change_parameters_->target.lanelet_id)) {
        return BT::NodeStatus::FAILURE;
      }
      boost::optional<std::pair<traffic_simulator::math::HermiteCurve, double>> traj_with_goal;
      traffic_simulator_msgs::msg::LaneletPose along_pose, goal_pose;
      switch (lane_change_parameters_->constraint.type) {
        case traffic_simulator::lane_change::Constraint::Type::NONE:
          traj_with_goal = hdmap_utils->getLaneChangeTrajectory(
            hdmap_utils->toMapPose(entity_status.lanelet_pose).pose, lane_change_parameters_.get(),
            10.0, 20.0, 1.0);
          along_pose = hdmap_utils->getAlongLaneletPose(
            entity_status.lanelet_pose,
            traffic_simulator::lane_change::Parameter::default_lanechange_distance);
          break;
        case traffic_simulator::lane_change::Constraint::Type::LATERAL_VELOCITY:
          traj_with_goal = hdmap_utils->getLaneChangeTrajectory(
            entity_status.lanelet_pose, lane_change_parameters_.get());
          along_pose = hdmap_utils->getAlongLaneletPose(
            entity_status.lanelet_pose,
            traffic_simulator::lane_change::Parameter::default_lanechange_distance);
          break;
        case traffic_simulator::lane_change::Constraint::Type::LONGITUDINAL_DISTANCE:
          traj_with_goal = hdmap_utils->getLaneChangeTrajectory(
            entity_status.lanelet_pose, lane_change_parameters_.get());
          along_pose = hdmap_utils->getAlongLaneletPose(
            entity_status.lanelet_pose, lane_change_parameters_->constraint.value);
          break;
        case traffic_simulator::lane_change::Constraint::Type::TIME:
          traj_with_goal = hdmap_utils->getLaneChangeTrajectory(
            entity_status.lanelet_pose, lane_change_parameters_.get());
          along_pose = hdmap_utils->getAlongLaneletPose(
            entity_status.lanelet_pose, lane_change_parameters_->constraint.value);
          break;
      }
      if (traj_with_goal) {
        curve_ = traj_with_goal->first;
        target_s_ = traj_with_goal->second;
        goal_pose.lanelet_id = lane_change_parameters_->target.lanelet_id;
        goal_pose.s = traj_with_goal->second;
        double offset = std::fabs(
          traffic_simulator::math::getRelativePose(
            hdmap_utils->toMapPose(along_pose).pose, hdmap_utils->toMapPose(goal_pose).pose)
            .position.y);
        switch (lane_change_parameters_->constraint.type) {
          case traffic_simulator::lane_change::Constraint::Type::NONE:
            lane_change_velocity_ = entity_status.action_status.twist.linear.x;
            break;
          case traffic_simulator::lane_change::Constraint::Type::LATERAL_VELOCITY:
            lane_change_velocity_ =
              curve_->getLength() / (offset / lane_change_parameters_->constraint.value);
            break;
          case traffic_simulator::lane_change::Constraint::Type::LONGITUDINAL_DISTANCE:
            lane_change_velocity_ = entity_status.action_status.twist.linear.x;
            break;
          case traffic_simulator::lane_change::Constraint::Type::TIME:
            lane_change_velocity_ = curve_->getLength() / lane_change_parameters_->constraint.value;
            break;
        }
      } else {
        return BT::NodeStatus::FAILURE;
      }
    }
  }
  if (curve_) {
    double target_accel = 0;
    switch (lane_change_parameters_->constraint.policy) {
      /**
       * @brief Force changing speed in order to fullfill constraint.
       */
      case traffic_simulator::lane_change::Constraint::Policy::FORCE:
        entity_status.action_status.twist = geometry_msgs::msg::Twist();
        entity_status.action_status.accel = geometry_msgs::msg::Accel();
        entity_status.action_status.twist.linear.x = lane_change_velocity_;
        current_s_ = current_s_ + entity_status.action_status.twist.linear.x * step_time;
        break;
      /**
       * @brief Changing linear speed and try to fullfill constraint.
       */
      case traffic_simulator::lane_change::Constraint::Policy::BEST_EFFORT:
        target_accel =
          (lane_change_velocity_ - entity_status.action_status.twist.linear.x) / step_time;
        if (entity_status.action_status.twist.linear.x > target_speed) {
          target_accel = boost::algorithm::clamp(target_accel, driver_model.deceleration * -1, 0);
        } else {
          target_accel = boost::algorithm::clamp(target_accel, 0, driver_model.acceleration);
        }
        geometry_msgs::msg::Accel accel_new;
        accel_new.linear.x = target_accel;
        geometry_msgs::msg::Twist twist_new;
        twist_new.linear.x = boost::algorithm::clamp(
          entity_status.action_status.twist.linear.x + accel_new.linear.x * step_time, -10,
          vehicle_parameters.performance.max_speed);
        twist_new.linear.y = 0.0;
        twist_new.linear.z = 0.0;
        twist_new.angular.x = 0.0;
        twist_new.angular.y = 0.0;
        twist_new.angular.z = 0.0;
        entity_status.action_status.twist = twist_new;
        entity_status.action_status.accel = accel_new;
        current_s_ = current_s_ + entity_status.action_status.twist.linear.x * step_time;
        break;
    }
    if (current_s_ < curve_->getLength()) {
      geometry_msgs::msg::Pose pose = curve_->getPose(current_s_, true);
      traffic_simulator_msgs::msg::EntityStatus entity_status_updated;
      entity_status_updated.pose = pose;
      auto lanelet_pose = hdmap_utils->toLaneletPose(pose, entity_status.bounding_box, false);
      if (lanelet_pose) {
        entity_status_updated.lanelet_pose = lanelet_pose.get();
      } else {
        entity_status_updated.lanelet_pose_valid = false;
      }
      entity_status_updated.action_status = entity_status.action_status;
      setOutput("updated_status", entity_status_updated);
      const auto waypoints = calculateWaypoints();
      if (waypoints.waypoints.empty()) {
        return BT::NodeStatus::FAILURE;
      }
      const auto obstacle = calculateObstacle(waypoints);
      setOutput("waypoints", waypoints);
      setOutput("obstacle", obstacle);
      return BT::NodeStatus::RUNNING;
    } else {
      const auto waypoints = calculateWaypoints();
      if (waypoints.waypoints.empty()) {
        return BT::NodeStatus::FAILURE;
      }
      const auto obstacle = calculateObstacle(waypoints);
      setOutput("waypoints", waypoints);
      setOutput("obstacle", obstacle);
      double s = (current_s_ - curve_->getLength()) + target_s_;
      curve_ = boost::none;
      current_s_ = 0;
      lane_change_velocity_ = 0;
      traffic_simulator_msgs::msg::EntityStatus entity_status_updated;
      traffic_simulator_msgs::msg::LaneletPose lanelet_pose;
      lanelet_pose.lanelet_id = lane_change_parameters_->target.lanelet_id;
      lanelet_pose.s = s;
      lanelet_pose.offset = 0;
      entity_status_updated.pose = hdmap_utils->toMapPose(lanelet_pose).pose;
      entity_status_updated.lanelet_pose = lanelet_pose;
      entity_status_updated.action_status = entity_status.action_status;
      setOutput("updated_status", entity_status_updated);
      return BT::NodeStatus::SUCCESS;
    }
  }
  return BT::NodeStatus::FAILURE;
}
}  // namespace vehicle
}  // namespace entity_behavior


#include "ConnectFour.h"
#include <iostream>
#include <Eigen/Dense>
#include <Eigen/Core>
#include <stdexcept>
#include <sstream>
#include <memory>

using namespace Eigen;

ConnectFour::ConnectFour(){
	this->width = Position::WIDTH;
	this->height = Position::HEIGHT;
	this->board = MatrixXf::Zero(this->height, this->width);
	this->player = 1;
	this->inRow = 4;
	this->played = "";
	this->allowedToPlay = ArrayXf::Zero(this->width);
}

std::unique_ptr<Game> ConnectFour::copy(){
	return std::make_unique<ConnectFour>(*this);
}

int ConnectFour::getInputPlanes(){
	return 1;
}

int ConnectFour::getOutputPlanes(){
	return 1;
}

int ConnectFour::getActionSize(){
	return this->width;
}

std::vector<int> ConnectFour::getBoardSize(){
	return {this->width, this->height};
}

std::string ConnectFour::getPlayedMoves(){
	return this->played;
}

void ConnectFour::printBoard(){
	std::cout<< this->board << std::endl;
}

void ConnectFour::play(int action){
	int indexToPlay = this->getEmptyColIndex(action);

	if (indexToPlay == -1){
		std::ostringstream  error;
		error << "Invalid action: " << action << "\n" << this->board;
		throw std::invalid_argument(error.str());
	}

	this->board(indexToPlay, action) = this->player;
	this->allowedToPlay[action]++;
	this->board_rep.playCol(action);
	this->played += std::to_string(action + 1);
	this->player *= -1 ;
}

int ConnectFour::getEmptyColIndex(int action){
	if (this->allowedToPlay[action] == this->height){
		return -1;
	}
	return this->height - this->allowedToPlay[action] - 1;
}

bool ConnectFour::ended(){
	if(Position::alignment(this->board_rep.getCurrentPosition())){
		this->winner = this->player*-1;
		return true;
	} else if ((this->allowedToPlay == this->height).all()){
		this->winner = 0;
		return true;
	}

	return false;
}

int ConnectFour::getWinner(){
	return this->winner;
}

MatrixXf ConnectFour::getBoard(){
	return this->board;
}

float ConnectFour::getPlayer(){
	return this->player;
}

ArrayXf ConnectFour::getPossibleActions(){
	ArrayXf poss = ArrayXf::Zero(this->getActionSize());

	for(int i = 0; i < this->width; i++){
		poss[i] = this->board_rep.canPlay(i);
	}

	return poss;
}


#ifndef AGATE_PRIS_APLAS_TRANSFORM_HPP
#define AGATE_PRIS_APLAS_TRANSFORM_HPP

#include <boost/qvm/map_mat_vec.hpp>
#include <boost/qvm/map_vec_mat.hpp>
#include <boost/qvm/mat.hpp>
#include <boost/qvm/mat_operations4.hpp>
#include <boost/qvm/quat.hpp>
#include <boost/qvm/quat_operations.hpp>
#include <boost/qvm/swizzle3.hpp>
#include <boost/qvm/swizzle4.hpp>
#include <boost/qvm/vec.hpp>
#include <boost/qvm/vec_mat_operations4.hpp>
#include <boost/qvm/vec_operations.hpp>

#include <utility>

namespace agate_pris {
namespace aplas {
    template <template <typename> class Parent, typename Scalar>
    class transform {
    public:
        typedef Parent<transform> parent_type;
        typedef Parent<transform const> const_parent_type;
        typedef Scalar scalar_type;
        typedef boost::qvm::vec<scalar_type, 3> vector_3_type;
        typedef boost::qvm::quat<scalar_type> quaternion_type;
        typedef boost::qvm::mat<scalar_type, 3, 3> matrix_3_x_3_type;
        typedef boost::qvm::mat<scalar_type, 4, 4> matrix_4_x_4_type;

    private:
        parent_type m_parent;
        vector_3_type m_local_position;
        quaternion_type m_local_rotation;
        vector_3_type m_local_scale;

    public:
        inline transform()
            : m_local_position(boost::qvm::zero_vec<scalar_type, 3>())
            , m_local_rotation(boost::qvm::identity_quat<scalar_type>())
            , m_local_scale(boost::qvm::scalar_cast<scalar_type>(boost::qvm::_111()))
        {
        }
        inline transform(parent_type parent)
            : m_parent(std::move(parent))
            , m_local_position(boost::qvm::zero_vec<scalar_type, 3>())
            , m_local_rotation(boost::qvm::identity_quat<scalar_type>())
            , m_local_scale(boost::qvm::scalar_cast<scalar_type>(boost::qvm::_111()))
        {
        }
        inline const_parent_type get_const_parent() const
        {
            return m_parent;
        }
        inline const_parent_type get_parent() const
        {
            return get_const_parent();
        }
        inline parent_type get_mutable_parent()
        {
            return m_parent;
        }
        inline parent_type get_parent()
        {
            return get_mutable_parent();
        }
        inline vector_3_type const& get_const_local_position() const
        {
            return m_local_position;
        }
        inline vector_3_type const& get_local_position() const
        {
            return get_const_local_position();
        }
        inline vector_3_type& get_mutable_local_position()
        {
            return const_cast<vector_3&>(get_const_local_position());
        }
        inline vector_3_type& get_local_position()
        {
            return get_mutable_local_position();
        }
        inline void set_local_position(vector_3_type const& v)
        {
            get_mutable_local_position() = v;
        }
        inline quaternion_type const& get_const_local_rotation() const
        {
            return m_local_rotation;
        }
        inline quaternion_type const& get_local_rotation() const
        {
            return get_const_local_rotation();
        }
        inline quaternion_type& get_mutable_local_rotation()
        {
            return const_cast<quaternion_type&>(get_const_local_rotation());
        }
        inline quaternion_type& get_local_rotation()
        {
            return get_mutable_local_rotation();
        }
        inline void set_local_rotation(quaternion_type const& q)
        {
            get_mutable_local_rotation() = q;
        }
        inline vector_3_type const& get_const_local_scale() const
        {
            return m_local_scale;
        }
        inline vector_3_type const& get_local_scale() const
        {
            return get_const_local_scale();
        }
        inline vector_3_type& get_mutable_local_scale()
        {
            return const_cast<vector_3_type&>(get_const_local_scale());
        }
        inline vector_3_type& get_local_scale()
        {
            return get_mutable_local_scale();
        }
        inline vector_3_type get_inversed_local_scale() const
        {
            typedef boost::qvm::vec_traits<vector_3_type> vec_traits;
            auto const& v = get_const_local_scale();
            return vector_3_type{
                1 / vec_traits::read_element<0>(v),
                1 / vec_traits::read_element<1>(v),
                1 / vec_traits::read_element<2>(v)
            };
        }
        inline void set_local_scale(vector_3 const& v)
        {
            get_mutable_local_scale() = v;
        }
        inline matrix_4_x_4_type get_local_to_world_matrix()
        {
            auto m
                = boost::qvm::translation_mat(get_const_local_position())
                * boost::qvm::convert_to<matrix_4_x_4_type>(get_const_local_rotation())
                * boost::qvm::diag_mat(boost::qvm::XYZ1(get_const_local_scale()));

            for (auto t = get_const_parent(); t; t = t->get_const_parent()) {
                m
                    = boost::qvm::translation_mat(t->get_const_local_position())
                    * boost::qvm::convert_to<matrix_4_x_4_type>(t->get_const_local_rotation())
                    * boost::qvm::diag_mat(boost::qvm::XYZ1(t->get_const_local_scale()))
                    * m;
            }

            return m;
        }
        inline matrix_4_x_4_type get_world_to_local_matrix()
        {
            auto m
                = boost::qvm::diag_mat(boost::qvm::XYZ1(get_inversed_local_scale()))
                * boost::qvm::convert_to<matrix_4_x_4_type>(boost::qvm::inverse(get_const_local_rotation()))
                * boost::qvm::translation_mat(-get_const_local_position());

            for (auto t = get_const_parent(); t; t = t->get_const_parent()) {
                m
                    = m
                    * boost::qvm::diag_mat(boost::qvm::XYZ1(t->get_inversed_local_scale()))
                    * boost::qvm::convert_to<matrix_4_x_4_type>(boost::qvm::inverse(t->get_const_local_rotation()))
                    * boost::qvm::translation_mat(-t->get_const_local_position());
            }

            return m;
        }
        inline vector_3_type get_position() const
        {
            auto p = get_const_parent();
            if (!p)
                return get_const_local_position();
            return boost::qvm::XYZ(
                p->get_local_to_world_matrix()
                * boost::qvm::XYZ1(get_const_local_position()));
        }
        inline quaternion_type get_rotation() const
        {
            auto r = get_const_local_rotation();
            for (auto t = get_const_parent(); t; t = t->get_const_parent())
                r *= t->get_const_local_rotation();
            return r;
        }
        inline matrix_3_x_3_type get_scale() const
        {
            auto r
                = get_const_local_rotation();
            auto m
                = boost::qvm::convert_to<matrix_3_x_3_type>(get_const_local_rotation())
                * boost::qvm::diag_mat(get_const_local_scale());

            for (auto t = get_const_parent(); t; t = t->get_const_parent()) {
                r *= t->get_const_local_rotation();
                m
                    = boost::qvm::convert_to<matrix_3_x_3_type>(get_const_local_rotation())
                    * boost::qvm::diag_mat(get_const_local_scale())
                    * m;
            }

            m
                = boost::qvm::convert_to<matrix_3_x_3_type>(boost::qvm::inverse(r))
                * m;

            return m;
        }
        inline matrix_3_x_3_type get_inversed_scale() const
        {
            auto r
                = get_const_local_rotation();
            auto m
                = boost::qvm::diag_mat(get_inversed_scale())
                * boost::qvm::convert_to<matrix_3_x_3_type>(boost::qvm::inverse(get_const_local_rotation()));

            for (auto t = get_const_parent(); t; t = t->get_const_parent()) {
                r *= t->get_const_local_rotation();
                m
                    = m
                    * boost::qvm::diag_mat(get_inversed_scale())
                    * boost::qvm::convert_to<matrix_3_x_3_type>(boost::qvm::inverse(get_const_local_rotation()));
            }

            m
                = m
                * boost::qvm::convert_to<matrix_3_x_3_type>(r);

            return m;
        }
        inline vector_3_type get_lossy_scale() const
        {
            return boost::qvm::diag(get_scale());
        }
        inline void set_position(vector_3_type const& v)
        {
            if (auto parent = get_const_parent()) {
                set_local_position(boost::qvm::XYZ(
                    parent->get_world_to_local_matrix()
                    * boost::qvm::XYZ1(get_const_local_position()));
                return;
            }

            set_local_position(v);
        }
        inline void set_rotation(quaternion_type const& q)
        {
            if (auto parent = get_const_parent()) {
                set_local_rotation(
                    q
                    * boost::qvm::inverse(parent->get_rotation()));
                return;
            }

            set_local_rotation(q);
        }
        inline void translate(vector_3_type const& translation)
        {
            if (auto const parent = get_const_parent()) {
                get_mutable_local_position()
                    += boost::qvm::XYZ(
                        parent->get_inversed_scale()
                        * boost::qvm::convert_to<matrix_4_x_4_type>(get_const_local_rotation())
                        * boost::qvm::XYZ1(translation));
                return;
            }

            get_mutable_local_position()
                += get_const_local_rotation()
                * translation;
        }
        inline void rotate(vector3_type const& axis, scalar_type const angle)
        {
            set_local_rotation(
                boost::qvm::rot_quat(axis, angle)
                * get_const_local_rotation());
        }
    };
}
}

#endif


#include <catch2/catch.hpp>

#include "constitutive/internal_variables.hpp"
#include "constitutive/constitutive_model_factory.hpp"
#include "constitutive/mechanics/solid/gaussian_ageing_affine_microsphere.hpp"

#include "exceptions.hpp"
#include "numeric/dense_matrix.hpp"
#include "io/json.hpp"

#include <Eigen/Eigenvalues>

#include <iostream>

std::string json_input_file()
{
    return "{\"name\": \"rubber\", \"elastic_modulus\": 2.0, \"poissons_ratio\": 0.45}";
}

constexpr auto ZERO_MARGIN = 1E-5;

using neon::json;
using namespace neon;

TEST_CASE("Gaussian affine microsphere model with ageing")
{
    using namespace neon::mechanics::solid;

    auto variables = std::make_shared<internal_variables_t>(1);

    // Add the required variables for an updated Lagrangian formulation
    variables->add(variable::second::deformation_gradient, variable::second::cauchy_stress);

    variables->add(variable::scalar::DetF);

    auto const material_data{"{\"name\" : \"rubber\","
                             "\"shear_modulus\" : 2.0e6,"
                             "\"bulk_modulus\" : 100e6,"
                             "\"segments_per_chain\" : 50,"
                             "\"cure_time\" : 100,"
                             "\"scission_probability\" : 1.0e-5,"
                             "\"recombination_probability\" : 1.0e-5}"};

    auto const constitutive_data{"{\"constitutive\" : {\"name\": \"microsphere\","
                                 "\"type\":\"affine\","
                                 "\"statistics\":\"gaussian\","
                                 "\"quadrature\":\"BO21\","
                                 "\"ageing\":\"BAND\"}}"};

    auto affine = make_constitutive_model(variables,
                                          json::parse(material_data),
                                          json::parse(constitutive_data));

    auto [deformation_gradients,
          cauchy_stresses] = variables->get(variable::second::deformation_gradient,
                                            variable::second::cauchy_stress);

    auto& jacobians = variables->get(variable::scalar::DetF);
    std::fill(begin(jacobians), end(jacobians), 1.0);

    auto& material_tangents = variables->get(variable::fourth::tangent_operator);

    Eigen::EigenSolver<Eigen::MatrixXd> eigen_solver;

    SECTION("Setup checks")
    {
        REQUIRE(affine->is_symmetric());
        REQUIRE(affine->is_finite_deformation());
        REQUIRE(affine->intrinsic_material().name() == "rubber");

        REQUIRE(variables->has(variable::scalar::active_shear_modulus));
        REQUIRE(variables->has(variable::scalar::inactive_shear_modulus));
        REQUIRE(variables->has(variable::scalar::active_segments));
        REQUIRE(variables->has(variable::scalar::inactive_segments));
        REQUIRE(variables->has(variable::scalar::reduction_factor));

        for (auto segment : variables->get(variable::scalar::active_segments))
        {
            REQUIRE(segment == Approx(50.0));
        }
        for (auto shear_modulus : variables->get(variable::scalar::active_shear_modulus))
        {
            REQUIRE(shear_modulus == Approx(2.0e6));
        }
    }
    SECTION("no load")
    {
        std::fill(begin(deformation_gradients), end(deformation_gradients), neon::matrix3::Identity());

        affine->update_internal_variables(1.0);

        // Check the network parameters
        auto [active_segments,
              inactive_segments,
              active_shear_moduli,
              inactive_shear_moduli,
              reductions] = variables->get(variable::scalar::active_segments,
                                           variable::scalar::inactive_segments,
                                           variable::scalar::active_shear_modulus,
                                           variable::scalar::inactive_shear_modulus,
                                           variable::scalar::reduction_factor);

        for (auto active_segment : active_segments)
        {
            REQUIRE(active_segment > 49.0);
            REQUIRE(active_segment < 50.0);
        }
        for (auto inactive_segment : inactive_segments)
        {
            REQUIRE(inactive_segment > 20.0);
            REQUIRE(inactive_segment < 25.0);
        }
        for (auto shear_modulus : active_shear_moduli)
        {
            REQUIRE(shear_modulus < 2.01e6);
            REQUIRE(shear_modulus > 2.0e6);
        }
        for (auto shear_modulus : inactive_shear_moduli)
        {
            REQUIRE(shear_modulus < 2000.0);
            REQUIRE(shear_modulus > 1000.0);
        }
        for (auto reduction : reductions)
        {
            REQUIRE(reduction > 0.99);
            REQUIRE(reduction < 1.0);
        }

        // Check for zero Cauchy stresses under no load
        for (auto const& cauchy_stress : variables->get(variable::second::cauchy_stress))
        {
            REQUIRE(cauchy_stress.norm() == Approx(0.0).margin(ZERO_MARGIN));
        }

        for (auto const& material_tangent : material_tangents)
        {
            REQUIRE((material_tangent - material_tangent.transpose()).norm()
                    == Approx(0.0).margin(ZERO_MARGIN));

            eigen_solver.compute(material_tangent);
            REQUIRE((eigen_solver.eigenvalues().real().array() > 0.0).all());
        }
    }
    SECTION("uniaxial load")
    {
        std::fill(begin(deformation_gradients), end(deformation_gradients), matrix3::Identity());

        affine->update_internal_variables(1.0);

        for (auto& cauchy_stress : cauchy_stresses)
        {
            REQUIRE(cauchy_stress.norm() == Approx(0.0).margin(ZERO_MARGIN));
        }

        std::cout << "setting to stretch 1.1\n";

        for (auto& F : deformation_gradients)
        {
            F(0, 0) = 1.1;
            F(1, 1) = 1.0 / std::sqrt(1.1);
            F(2, 2) = 1.0 / std::sqrt(1.1);
        }

        affine->update_internal_variables(1.0);

        for (auto& cauchy_stress : cauchy_stresses)
        {
            std::cout << cauchy_stress << "\n\n";
            REQUIRE(cauchy_stress.norm() > 0.0);
        }

        affine->update_internal_variables(1.0);

        for (auto& cauchy_stress : cauchy_stresses)
        {
            std::cout << cauchy_stress << "\n\n";
            REQUIRE(cauchy_stress.norm() > 0.0);
        }

        std::cout << "setting to unity\n";
        std::fill(begin(deformation_gradients), end(deformation_gradients), matrix3::Identity());

        affine->update_internal_variables(1.0);

        for (auto& cauchy_stress : cauchy_stresses)
        {
            std::cout << cauchy_stress << "\n\n";
            REQUIRE(cauchy_stress.norm() >= 0.0);
        }

        std::cout << "finished the analysis\n";

        for (auto const& material_tangent : material_tangents)
        {
            REQUIRE((material_tangent - material_tangent.transpose()).norm()
                    == Approx(0.0).margin(ZERO_MARGIN));

            eigen_solver.compute(material_tangent);
            REQUIRE((eigen_solver.eigenvalues().real().array() > 0.0).all());
        }
    }
}
TEST_CASE("Gaussian affine microsphere model with crosslinking only")
{
    using namespace neon::mechanics::solid;

    auto variables = std::make_shared<internal_variables_t>(1);

    // Add the required variables for an updated Lagrangian formulation
    variables->add(variable::second::deformation_gradient, variable::second::cauchy_stress);

    variables->add(variable::scalar::DetF);

    auto const material_data{"{\"name\" : \"rubber\","
                             "\"shear_modulus\" : 2.0e6,"
                             "\"bulk_modulus\" : 100e6,"
                             "\"segments_per_chain\" : 50,"
                             "\"cure_time\" : 100,"
                             "\"scission_probability\" : 0.0,"
                             "\"recombination_probability\" : 1.0e-5}"};

    auto const constitutive_data{"{\"constitutive\" : {\"name\": \"microsphere\","
                                 "\"type\":\"affine\","
                                 "\"statistics\":\"gaussian\","
                                 "\"quadrature\":\"BO21\","
                                 "\"ageing\":\"BAND\"}}"};

    auto affine = make_constitutive_model(variables,
                                          json::parse(material_data),
                                          json::parse(constitutive_data));

    auto [deformation_gradients,
          cauchy_stresses] = variables->get(variable::second::deformation_gradient,
                                            variable::second::cauchy_stress);

    auto& jacobians = variables->get(variable::scalar::DetF);
    std::fill(begin(jacobians), end(jacobians), 1.0);

    auto& material_tangents = variables->get(variable::fourth::tangent_operator);

    Eigen::EigenSolver<Eigen::MatrixXd> eigen_solver;

    SECTION("no load")
    {
        std::fill(begin(deformation_gradients), end(deformation_gradients), neon::matrix3::Identity());

        affine->update_internal_variables(1.0);

        // Check the network parameters
        auto [active_segments,
              inactive_segments,
              active_shear_moduli,
              inactive_shear_moduli,
              reductions] = variables->get(variable::scalar::active_segments,
                                           variable::scalar::inactive_segments,
                                           variable::scalar::active_shear_modulus,
                                           variable::scalar::inactive_shear_modulus,
                                           variable::scalar::reduction_factor);

        for (auto active_segment : active_segments)
        {
            REQUIRE(active_segment > 49.0);
            REQUIRE(active_segment < 50.0);
        }
        for (auto inactive_segment : inactive_segments)
        {
            REQUIRE(inactive_segment == Approx(0.0).margin(0.0));
        }
        for (auto shear_modulus : active_shear_moduli)
        {
            REQUIRE(shear_modulus < 2.01e6);
            REQUIRE(shear_modulus > 2.0e6);
        }
        for (auto shear_modulus : inactive_shear_moduli)
        {
            REQUIRE(shear_modulus == Approx(0.0).margin(0.0));
        }
        for (auto reduction : reductions)
        {
            REQUIRE(reduction <= 1.0);
        }

        for (auto const& material_tangent : material_tangents)
        {
            REQUIRE((material_tangent - material_tangent.transpose()).norm()
                    == Approx(0.0).margin(ZERO_MARGIN));

            eigen_solver.compute(material_tangent);
            REQUIRE((eigen_solver.eigenvalues().real().array() > 0.0).all());
        }
    }
    SECTION("constant load")
    {
        std::fill(begin(deformation_gradients), end(deformation_gradients), neon::matrix3::Identity());

        affine->update_internal_variables(1.0);

        for (auto& F : deformation_gradients)
        {
            F(0, 0) = 1.1;
            F(1, 1) = 1.0 / std::sqrt(1.1);
            F(2, 2) = 1.0 / std::sqrt(1.1);
        }

        affine->update_internal_variables(1.0);

        for (auto const& cauchy_stress : cauchy_stresses)
        {
            std::cout << "Step 1: cauchy_stress\n" << cauchy_stress << "\n\n";
        }

        affine->update_internal_variables(1.0);

        for (auto const& cauchy_stress : cauchy_stresses)
        {
            std::cout << "Step 2: cauchy_stress\n" << cauchy_stress << "\n\n";
        }

        affine->update_internal_variables(1.0);

        for (auto const& cauchy_stress : cauchy_stresses)
        {
            std::cout << "Step 3: cauchy_stress\n" << cauchy_stress << "\n\n";
        }

        affine->update_internal_variables(1.0);

        for (auto const& cauchy_stress : cauchy_stresses)
        {
            std::cout << "Step 3: cauchy_stress\n" << cauchy_stress << "\n\n";
        }

        affine->update_internal_variables(1.0);

        for (auto const& cauchy_stress : cauchy_stresses)
        {
            std::cout << "Step 3: cauchy_stress\n" << cauchy_stress << "\n\n";
        }

        for (auto const& material_tangent : material_tangents)
        {
            std::cout << "tangent_matrix\n" << material_tangent << "\n";

            REQUIRE((material_tangent - material_tangent.transpose()).norm()
                    == Approx(0.0).margin(ZERO_MARGIN));

            eigen_solver.compute(material_tangent);
            REQUIRE((eigen_solver.eigenvalues().real().array() > 0.0).all());
        }
    }
}
// TEST_CASE("Finite J2 plasticity")
// {
//     using namespace neon::mechanics::solid;
//
//     // Create a json reader object from a string
//     std::string
//         input_data = "{\"name\": \"steel\", \"elastic_modulus\": 200.0e9, \"poissons_ratio\": 0.3,
//         "
//                      "\"yield_stress\": 200.0e6, \"isotropic_hardening_modulus\": 400.0e6}";
//
//     std::string simulation_input = "{\"constitutive\" : {\"name\" : \"small_strain_J2_plasticity\", "
//                                    "\"finite_strain\":true}}";
//
//     Json::Value material_data, simulation_data;
//
//     Json::Reader reader;
//
//     REQUIRE(reader.parse(input_data.c_str(), material_data));
//     REQUIRE(reader.parse(simulation_input.c_str(), simulation_data));
//
//     internal_variables variables(1);
//
//     // Add the required variables for an updated Lagrangian formulation
//     variables->add(variable::second::deformation_gradient,
//     variable::second::cauchy_stress); variables->add(variable::scalar::DetF);
//
//     auto small_strain_J2_plasticity = make_constitutive_model(variables, material_data,
//     simulation_data);
//
//     // Get the tensor variables
//     auto[deformation_gradients, cauchy_stresses] =
//     variables->get(variable::second::deformation_gradient,
//                                               variable::second::cauchy_stress);
//
//     auto& jacobians = variables->get(variable::scalar::DetF);
//
//     auto& material_tangents = variables->get(variable::fourth::tangent_operator);
//
//     for (auto& F : deformation_gradients) F = neon::matrix3::Identity();
//     for (auto& J : jacobians) J = 1.0;
//
//     variables.commit();
//
//     SECTION("Sanity checks")
//     {
//         REQUIRE(small_strain_J2_plasticity->is_finite_deformation() == true);
//         REQUIRE(small_strain_J2_plasticity->intrinsic_material().name() == "steel");
//
//         REQUIRE(variables->has(variable::scalar::von_mises_stress));
//         REQUIRE(variables->has(variable::scalar::effective_plastic_strain));
//         REQUIRE(variables->has(variable::second::hencky_strain_elastic));
//         REQUIRE(variables->has(variable::fourth::tangent_operator));
//     }
//     SECTION("Initial material tangent symmetry")
//     {
//         for (auto const& material_tangent : material_tangents)
//         {
//             REQUIRE((material_tangent - material_tangent.transpose()).norm()
//                     == Approx(0.0).margin(ZERO_MARGIN));
//         }
//     }
//     SECTION("No load")
//     {
//         small_strain_J2_plasticity->update_internal_variables(1.0);
//
//         for (auto const& material_tangent : material_tangents)
//         {
//             REQUIRE((material_tangent - material_tangent.transpose()).norm()
//                     == Approx(0.0).margin(ZERO_MARGIN));
//         }
//         for (auto& cauchy_stress : cauchy_stresses)
//         {
//             REQUIRE(cauchy_stress.norm() == Approx(0.0).margin(ZERO_MARGIN));
//         }
//     }
//     // SECTION("Uniaxial elastic load")
//     // {
//     //     for (auto& F : deformation_gradients) F(2, 2) = 1.001;
//     //
//     //     small_strain_J2_plasticity->update_internal_variables(1.0);
//     //
//     //     for (auto const& material_tangent : material_tangents)
//     //     {
//     //         REQUIRE((material_tangent - material_tangent.transpose()).norm() ==
//     //         Approx(0.0).margin(ZERO_MARGIN));
//     //     }
//     //     for (auto& cauchy_stress : cauchy_stresses)
//     //     {
//     //         REQUIRE((cauchy_stress - cauchy_stress.transpose()).norm() ==
//     Approx(0.0).margin(ZERO_MARGIN));
//     //
//     //         // Shear components should be close to zero
//     //         REQUIRE(cauchy_stress(0, 1) == Approx(0.0).margin(ZERO_MARGIN));
//     //         REQUIRE(cauchy_stress(0, 2) == Approx(0.0).margin(ZERO_MARGIN));
//     //         REQUIRE(cauchy_stress(1, 2) == Approx(0.0).margin(ZERO_MARGIN));
//     //         REQUIRE(cauchy_stress(0, 0) > 0.0);
//     //         REQUIRE(cauchy_stress(1, 1) > 0.0);
//     //         REQUIRE(cauchy_stress(2, 2) > 0.0);
//     //     }
//     //     for (auto& von_mises_stress :
//     variables->get(variable::scalar::von_mises_stress))
//     //     {
//     //         REQUIRE(von_mises_stress < 200.0e6);
//     //     }
//     // }
//     // SECTION("Plastic uniaxial elastic load")
//     // {
//     //     for (auto& F : deformation_gradients) F(2, 2) = 0.003;
//     //
//     //     small_strain_J2_plasticity->update_internal_variables(1.0);
//     //
//     //     auto[von_mises_stresses, accumulated_plastic_strains] =
//     //     variables->get(variable::scalar::von_mises_stress,
//     // variable::scalar::effective_plastic_strain);
//     //
//     //     // Ensure symmetry is correct
//     //     for (auto const& material_tangent : material_tangents)
//     //     {
//     //         REQUIRE((material_tangent - material_tangent.transpose()).norm() ==
//     //         Approx(0.0).margin(ZERO_MARGIN));
//     //     }
//     //     for (auto& cauchy_stress : cauchy_stresses)
//     //     {
//     //         REQUIRE(cauchy_stress.norm() != Approx(0.0).margin(ZERO_MARGIN));
//     //
//     //         // Shear components should be close to zero
//     //         REQUIRE(cauchy_stress(0, 1) == Approx(0.0).margin(ZERO_MARGIN));
//     //         REQUIRE(cauchy_stress(0, 2) == Approx(0.0).margin(ZERO_MARGIN));
//     //         REQUIRE(cauchy_stress(1, 2) == Approx(0.0).margin(ZERO_MARGIN));
//     //         REQUIRE(cauchy_stress(0, 0) > 0.0);
//     //         REQUIRE(cauchy_stress(1, 1) > 0.0);
//     //         REQUIRE(cauchy_stress(2, 2) > 0.0);
//     //     }
//     //     for (auto& accumulated_plastic_strain : accumulated_plastic_strains)
//     //     {
//     //         REQUIRE(accumulated_plastic_strain > 0.0);
//     //     }
//     //     for (auto& von_mises_stress : von_mises_stresses)
//     //     {
//     //         // Should experience hardening
//     //         REQUIRE(von_mises_stress <= 201.0e6);
//     //         REQUIRE(von_mises_stress > 200.0e6);
//     //     }
//     // }
// }


//==============================================================================
//         Copyright 2003 - 2011 LASMEA UMR 6602 CNRS/Univ. Clermont II
//         Copyright 2009 - 2011 LRI    UMR 8623 CNRS/Univ Paris Sud XI
//
//          Distributed under the Boost Software License, Version 1.0.
//                 See accompanying file LICENSE.txt or copy at
//                     http://www.boost.org/LICENSE_1_0.txt
//==============================================================================
#ifndef BOOST_SIMD_TOOLBOX_ARITHMETIC_FUNCTIONS_SCALAR_RANDOM_HPP_INCLUDED
#define BOOST_SIMD_TOOLBOX_ARITHMETIC_FUNCTIONS_SCALAR_RANDOM_HPP_INCLUDED
#include <boost/simd/toolbox/arithmetic/functions/random.hpp>
#include <boost/dispatch/meta/strip.hpp>
#include <boost/simd/include/functions/scalar/iround.hpp>
#include <boost/simd/include/functions/scalar/min.hpp>
#include <boost/simd/include/functions/scalar/dist.hpp>

// TO DO proper generator
/////////////////////////////////////////////////////////////////////////////
// Implementation when type A0 and A1 are not equal types
// This oddity is due to the fact that many tests up to now
// don't pass typed integer parameters to the macro NT2_CREATE_BUF
/////////////////////////////////////////////////////////////////////////////
namespace boost { namespace simd { namespace ext
{
  BOOST_SIMD_FUNCTOR_IMPLEMENTATION( boost::simd::tag::random_, tag::cpu_
                             , (A0)(A1)
                             , (scalar_< arithmetic_<A0> >)
                                 (scalar_< arithmetic_<A1> >)
                             )
  {
    typedef A0 result_type;
    BOOST_SIMD_FUNCTOR_CALL(2)
    {
      return boost::simd::random(static_cast<result_type>(a0),
                                 static_cast<result_type>(a1));
    }
  };

  BOOST_SIMD_FUNCTOR_IMPLEMENTATION( boost::simd::tag::random_, tag::cpu_
                            , (A0)
                            , (scalar_< arithmetic_<A0> >)
                        (scalar_< arithmetic_<A0> >)
                            )
  {
    typedef A0 result_type;
    BOOST_SIMD_FUNCTOR_CALL_REPEAT(2)
    {
      return iround(static_cast<double>(a1-a0)*(static_cast<double>(::rand())/RAND_MAX)+
                static_cast<double>(a0));
    }
  };

  BOOST_SIMD_FUNCTOR_IMPLEMENTATION( boost::simd::tag::random_, tag::cpu_
                            , (A0)
                            , (scalar_< unsigned_<A0> >)
                              (scalar_< unsigned_<A0> >)
                            )
  {
    typedef A0 result_type;
    BOOST_SIMD_FUNCTOR_CALL_REPEAT(2)
    {
      return iround(static_cast<double>(dist(a1, a0))*(static_cast<double>(::rand())/RAND_MAX)+
                static_cast<double>(min(a0, a1)));
    }
  };

  BOOST_SIMD_FUNCTOR_IMPLEMENTATION( boost::simd::tag::random_, tag::cpu_
                            , (A0)
                            , (scalar_< floating_<A0> >)
                              (scalar_< floating_<A0> >)
                            )
  {
    typedef A0 result_type;
    BOOST_SIMD_FUNCTOR_CALL_REPEAT(2)
    {
      return static_cast<A0>(
        static_cast<double>(a1-a0)*(static_cast<double>(::rand())/RAND_MAX)
        + static_cast<double>(a0)
      );
    }
  };
} } }

#endif


///////////////////////////////////////////////////////////////////////////////
// statistics::survival::data::meta::failure_distribution.hpp                //
//                                                                           //
//  Copyright 2009 Erwann Rogard. 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_STATISTICS_SURVIVAL_DATA_META_FAILURE_DISTRIBUTION_HPP_ER_2009
#define BOOST_STATISTICS_SURVIVAL_DATA_META_FAILURE_DISTRIBUTION_HPP_ER_2009
#include <boost/statistics/model/wrap/aggregate/model_covariate_parameter.hpp>

namespace boost{
namespace statistics{
namespace survival{
namespace data{

    // Model + Covariate + Parameter ---> failure time distribution (such as
    // defined in boost/math/distributions)
    //
    // These are only declarations. The implementation must specialize on M.
    template<typename M>
    struct meta_failure_distribution{
        // typedef ... type
    };

    template<typename M,typename X,typename P>
    typename meta_failure_distribution<M>::type
    make_failure_distribution(
        boost::statistics::model::model_covariate_parameter_<M,X,P> mcp
    ){
        typedef meta_failure_distribution<M> meta_;
        return meta_::make(
            mcp
        );
    }

}// data
}// survival
}// statistics
}// boost

#endif


#include "sv/dsol/hessian.h"

#include <Eigen/Cholesky>  // LLT and LDLT
#include <Eigen/LU>        // inverse

#include "sv/util/logging.h"
#include "sv/util/tbb.h"

namespace sv::dsol {

using Vector2d = Eigen::Vector2d;
using Vector3i = Eigen::Vector3i;
using LltLowerInplace = Eigen::LLT<MatrixXdRef, Eigen::Lower>;
using LdltLowerInplace = Eigen::LDLT<MatrixXdRef, Eigen::Lower>;

namespace {

constexpr int dp = Dim::kPose;
constexpr int df = Dim::kFrame;
constexpr int dd = Dim::kPoint;
constexpr int da = Dim::kAffine;

}  // namespace

/// ============================================================================
void PatchHessian::AddI(const Vector2d& It, double r, double w) noexcept {
  const auto wr = w * r;
  ItI.noalias() += (It * w) * It.transpose();
  Itr.noalias() += It * wr;
  r2 += r * r;
  wr2 += r * wr;
}

void PatchHessian::SetI(const Matrix2Kd& It,
                        const ArrayKd& r,
                        const ArrayKd& w) noexcept {
  const ArrayKd wr = w * r;
  ItI.noalias() = It * w.matrix().asDiagonal() * It.transpose();
  Itr.noalias() = It * wr.matrix();
  r2 = (r * r).sum();
  wr2 = (wr * r).sum();
}

/// ============================================================================
void PatchHessian1::AddA(const Vector2d& It,
                         const Vector2d& At,
                         double r,
                         double w) noexcept {
  ItA.noalias() += (It * w) * At.transpose();
  AtA.noalias() += (At * w) * At.transpose();
  Atr.noalias() += At * (w * r);
}

void PatchHessian1::SetA(const Matrix2Kd& It,
                         const Matrix2Kd& At,
                         const ArrayKd& r,
                         const ArrayKd& w) noexcept {
  const MatrixK2d wA = w.matrix().asDiagonal() * At.transpose();
  ItA.noalias() = It * wA;
  AtA.noalias() = At * wA;
  Atr.noalias() = At * (w * r).matrix();
}

/// ============================================================================
void PatchHessian2::AddA(const Vector2d& It,
                         const Vector2d& A0t,
                         const Vector2d& A1t,
                         double r,
                         double w) noexcept {
  const Vector2d Itw = It * w;
  ItA0.noalias() += Itw * A0t.transpose();
  ItA1.noalias() += Itw * A1t.transpose();

  A0tA0.noalias() += (A0t * w) * A0t.transpose();
  A0tA1.noalias() += (A0t * w) * A1t.transpose();
  A1tA1.noalias() += (A1t * w) * A1t.transpose();

  const auto wr = w * r;
  A0tr.noalias() += A0t * wr;
  A1tr.noalias() += A1t * wr;
}

void PatchHessian2::SetA(const Matrix2Kd& It,
                         const Matrix2Kd& A0t,
                         const Matrix2Kd& A1t,
                         const ArrayKd& r,
                         const ArrayKd& w) noexcept {
  const MatrixK2d wA0 = w.matrix().asDiagonal() * A0t.transpose();
  const MatrixK2d wA1 = w.matrix().asDiagonal() * A1t.transpose();
  ItA0.noalias() = It * wA0;
  ItA1.noalias() = It * wA1;

  A0tA0.noalias() = A0t * wA0;
  A0tA1.noalias() = A0t * wA1;
  A1tA1.noalias() = A1t * wA1;

  const ArrayKd wr = w * r;
  A0tr.noalias() = A0t * wr.matrix();
  A1tr.noalias() = A1t * wr.matrix();
}

/// ============================================================================
FrameHessian1& FrameHessian1::operator+=(const FrameHessian1& rhs) noexcept {
  if (rhs.n > 0) {
    H += rhs.H;
    b += rhs.b;
    n += rhs.n;
    c += rhs.c;
  }
  return *this;
}

void FrameHessian1::AddPatchHess(const PatchHessian1& ph,
                                 const Matrix26d& G,
                                 int affine_offset) noexcept {
  // [Gt*(It*I)*G, Gt*(It*A)] | [Gt*(It*r)]
  // [   (At*I)*G,    (At*A)] | [   (At*r)]

  ++n;
  // FrameHessian1 is 10x10, where the variable is ordered as follows
  // [pose, aff0, aff1]
  // Add pose block, which is always top-left 6x6
  // Gt * ItI * G
  H.topLeftCorner<dp, dp>().noalias() += G.transpose() * ph.ItI * G;
  // Gt * It * r
  b.head<dp>().noalias() -= G.transpose() * ph.Itr;
  // cost
  c += ph.r2;

  if (affine_offset < 0) return;

  // Depends on which camera, the affine parameters would be different
  // Therefore we use affine_offset to indicate where should we put the
  // corresponding Hessian block. affine_index is pose_dim + affine_offset, left
  // is pose_dim, right is pose_dim + affine_dim
  const int ia = dp + affine_offset;  // 6(L) or 8(R)

  // https://eigen.tuxfamily.org/dox/classEigen_1_1LLT.html
  // Performance: for best performance, it is recommended to use a column-major
  // storage format with the Lower triangular part (the default), or,
  // equivalently, a row-major storage format with the Upper triangular part.
  // Otherwise, you might get a 20% slowdown for the full factorization step,
  // and rank-updates can be up to 3 times slower.

  // AtI * G
  H.block<da, dp>(ia, 0).noalias() += ph.ItA.transpose() * G;
  // AtA
  H.block<da, da>(ia, ia) += ph.AtA;
  // Atr
  b.segment<da>(ia) -= ph.Atr;
}

/// ============================================================================
FrameHessian2::FrameHessian2(int fi, int fj) : i_{fi}, j_{fj} {
  CHECK_GE(fi, 0);
  CHECK_GE(fj, 0);
  CHECK_NE(fi, fj);
}

std::string FrameHessian2::Repr() const {
  return fmt::format("FrameHessian2(i={}, j={}, n={}, c={:.4e})", i_, j_, n, c);
}

FrameHessian2& FrameHessian2::operator+=(const FrameHessian2& rhs) {
  CHECK_GE(i_, 0);
  CHECK_GE(j_, 0);
  CHECK_EQ(i_, rhs.i_);
  CHECK_EQ(j_, rhs.j_);

  if (rhs.n > 0) {
    Hii += rhs.Hii;
    Hij += rhs.Hij;
    Hjj += rhs.Hjj;

    bi += rhs.bi;
    bj += rhs.bj;

    c += rhs.c;
    n += rhs.n;
  }

  return *this;
}

void FrameHessian2::AddPatchHess(const PatchHessian2& ph,
                                 const Matrix26d& G0,
                                 const Matrix26d& G1,
                                 int affine_offset) noexcept {
  ++n;

  // ---------------------------------------------------------------------------
  // JtJ                                              | -Jtr
  // ---------------------------------------------------------------------------
  // [Hii,                  , Hij                   ] | [bi]
  // [                      , Hjj                   ] | [bj]
  // ---------------------------------------------------------------------------
  // [[G0t*ItI*G0, G0t*ItA0], [G0t*ItI*G1, G0t*ItA1]] | [G0t*Itr]
  // [[   A0tI*G0,   A0tA0)], [   A0tI*G1,    A0tA1]] | [  A0t*r]
  //                          [G1t*ItI*G1, G1t*ItA1]] | [G1t*Itr]
  //                          [   A1tI*G1,   A1t*A1]] | [  A1t*r]
  // ---------------------------------------------------------------------------
  // Note that we don't strictly enforce i < j, so it is possible when i > j, we
  // have a lower triangular form. This is taken care of in
  // FullBlockHessian.Add()

  // Pose block
  // G0t * ItI * G0
  Hii.topLeftCorner<dp, dp>().noalias() += G0.transpose() * ph.ItI * G0;
  // G0t * ItI * G1
  Hij.topLeftCorner<dp, dp>().noalias() += G0.transpose() * ph.ItI * G1;
  // G1t * ItI * G1
  Hjj.topLeftCorner<dp, dp>().noalias() += G1.transpose() * ph.ItI * G1;

  // G0t * Itr
  bi.head<dp>().noalias() -= G0.transpose() * ph.Itr;
  // G1t * Itr
  bj.head<dp>().noalias() -= G1.transpose() * ph.Itr;

  // cost
  c += ph.r2;

  if (affine_offset < 0) return;
  // since we only project points from left image, ia0 should always be 0
  // ia1 depends on which image we project into, 0 for left, 2 for right
  const int ia0 = dp;
  const int ia1 = dp + affine_offset;

  // We only fill the upper-triangular part for now and make it symmetric later
  // before solving

  // G0t * ItA0
  Hii.block<dp, da>(0, ia0).noalias() += G0.transpose() * ph.ItA0;
  // A0t * A0
  Hii.block<da, da>(ia0, ia0) += ph.A0tA0;

  // This is already an off-diagonal block so we need it to be full
  // G0t * It * A1
  Hij.block<dp, da>(0, ia1).noalias() += G0.transpose() * ph.ItA1;
  // A0t * I * G1
  Hij.block<da, dp>(ia0, 0).noalias() += ph.ItA0.transpose() * G1;
  // A0t * A1
  Hij.block<da, da>(ia0, ia1) += ph.A0tA1;

  // G1t * ItA1
  Hjj.block<dp, da>(ia1, 0).noalias() += G1.transpose() * ph.ItA1;
  // A1t * A1
  Hjj.block<da, da>(ia1, ia1) += ph.A1tA1;

  // A0t * r
  bi.segment<da>(ia0) -= ph.A0tr;
  // A1t * r
  bj.segment<da>(ia1) -= ph.A1tr;
}

/// ============================================================================
void FrameHessianX::ResetData() noexcept {
  Reset();

  if (!data_.empty()) {
    storage().setZero();
  }
}

int FrameHessianX::MapFrames(int nframes) {
  const int dim = nframes * Dim::kFrame;
  const int size = dim * dim + dim;
  ResizeData(size);

  auto* ptr = data_ptr();
  new (&Hpp) MatrixXdMap(ptr, dim, dim);
  ptr += Hpp.size();
  new (&bp) VectorXdMap(ptr, dim);
  ptr += bp.size();

  CHECK_EQ(ptr - data_ptr(), size);
  return size;
}

void FrameHessianX::Scale(double s) noexcept {
  if (s == 1.0) return;
  Hpp.array() *= s;
  bp.array() *= s;
}

void FrameHessianX::AddDiag(int ind, int size, double val) {
  CHECK_GE(ind, 0);
  CHECK_LE(ind + size, Hpp.rows())
      << fmt::format("ind: {}, size: {}, H: {}", ind, size, Hpp.rows());
  Hpp.diagonal().segment(ind, size).array() += val;
}

void FrameHessianX::FixGauge(double lambda, bool fix_scale) {
  CHECK_GE(num_frames(), 1);
  AddDiag(0, df, lambda);

  if (fix_scale) {
    // Fix the translation parameter of the second frame
    // Ideally one only need to fix one direction of translation
    CHECK_GE(num_frames(), 2);
    AddDiag(df + 3, 3, lambda);
  }
}

/// ============================================================================
void SchurFrameHessian::AddPriorHess(const PriorFrameHessian& prior) {
  CHECK(prior.Ok());
  CHECK_GE(num_frames(), prior.num_frames());

  const auto m = prior.dim_frames();
  Hpp.topLeftCorner(m, m) += prior.Hpp;
  bp.head(m) += prior.bp;
  n += prior.n;
}

void SchurFrameHessian::Solve(VectorXdRef xp, VectorXdRef yp) {
  CHECK(Ok());
  CHECK(!empty());
  CHECK_EQ(xp.size(), dim_frames());

  // s is a copy
  const auto s =
      (Hpp.diagonal().array().abs() + 10).sqrt().inverse().matrix().eval();
  const auto S = s.asDiagonal();

  // Use inplace decomposition
  // https://eigen.tuxfamily.org/dox/group__InplaceDecomposition.html
  // For 4 frames, inplace solve is roughly 10x faster than normal solve (1us vs
  // 10us).
  // Because we add both a large and small value to the diagonal, we use ldlt to
  // ensure stability of the decomposition
  LdltLowerInplace solver(Hpp);
  xp = solver.solve(bp);
  yp = S.inverse() * xp;

  // Inplace decomp modifies storage, so we reset n
  n = 0;
}

void SchurFrameHessian::MargFrame(PriorFrameHessian& prior, int fid) {
  CHECK(Ok());
  CHECK_GE(fid, 0);
  CHECK_LT(fid, num_frames());
  CHECK_EQ(num_frames(), prior.num_frames() + 1);

  // First rotate the block that needs to be marginalized to top left. It is
  // possible to also rotate it to bottom right, but it is highly likely that
  // the oldest frame will be marginalized, so move it to top left will take
  // fewer operations (noop if its already there). Therefore, the subsequent
  // frame marginalization operation will be different from the point
  // marginalization one.
  StableRotateBlockTopLeft(Hpp, bp, fid, df);

  // Now we can marginalize the top left block of Hsc on to the rest
  MargTopLeftBlock(Hpp, bp, prior.Hpp, prior.bp, df);

  // Record how many points are marginalized
  prior.n = n;
}

/// ============================================================================
int FramePointHessian::CalcDataSize(int frame_dim, int point_dim) noexcept {
  // Hpp + Hpm + Hmm + bp + bm + xp + xm
  return frame_dim * frame_dim +  // Hpp
         frame_dim * point_dim +  // Hpm
         point_dim +              // Hmm
         frame_dim +              // bp
         point_dim +              // bm
         frame_dim +              // xp
         point_dim;               // xm
}

std::string FramePointHessian::Repr() const {
  return fmt::format(
      "FullBlockHessian(nframes={}, npoints={}, dframes={}, dpoints={}, "
      "num={}, costs={:.4e}, size={}, capacity={}, usage={:.2f}%)",
      num_frames(),
      num_points(),
      dim_frames(),
      dim_points(),
      num_costs(),
      cost(),
      size(),
      capacity(),
      (100.0 * size()) / capacity());
}

auto FramePointHessian::XpAt(int i) const noexcept
    -> Eigen::Map<const Vector10d> {
  CHECK_GE(i, 0);
  CHECK_LT(i, num_frames());
  return Eigen::Map<const Vector10d>(xp.data() + i * df);
}

auto FramePointHessian::XpAt(int i) noexcept -> Eigen::Map<Vector10d> {
  CHECK_GE(i, 0);
  CHECK_LT(i, num_frames());
  return Eigen::Map<Vector10d>(xp.data() + i * df);
}

int FramePointHessian::MapFull(int nframes, int npoints) {
  const int frame_dim = nframes * Dim::kFrame;
  const int point_dim = npoints * Dim::kPoint;
  const int size = CalcDataSize(frame_dim, point_dim);
  ResizeData(size);

  // https://eigen.tuxfamily.org/dox
  // group__TutorialMapClass.html#TutorialMapPlacementNew
  auto* ptr = data_ptr();

  new (&Hpp) MatrixXdMap(ptr, frame_dim, frame_dim);
  ptr += Hpp.size();
  new (&Hpm) MatrixXdMap(ptr, frame_dim, point_dim);
  ptr += Hpm.size();
  new (&Hmm) VectorXdMap(ptr, point_dim);
  new (&Hmm_inv) VectorXdMap(ptr, point_dim);
  ptr += Hmm.size();

  new (&bp) VectorXdMap(ptr, frame_dim);
  ptr += bp.size();
  new (&bm) VectorXdMap(ptr, point_dim);
  ptr += bm.size();

  new (&xp) VectorXdMap(ptr, frame_dim);
  ptr += xp.size();
  new (&xm) VectorXdMap(ptr, point_dim);
  ptr += xm.size();

  CHECK_EQ(ptr - data_ptr(), size);
  return size;
}

void FramePointHessian::ReserveFull(int nframes, int npoints) {
  CHECK_GT(nframes, 0);
  CHECK_GT(npoints, 0);

  const int frame_dim = nframes * Dim::kFrame;
  const int point_dim = npoints * Dim::kPoint;
  const auto size = CalcDataSize(frame_dim, point_dim);
  ReserveData(size);
}

void FramePointHessian::ResetFull(double Hpp_diag) noexcept {
  ready_ = false;
  ResetData();
  if (Hpp_diag > 0) {
    Hpp.diagonal().setConstant(Hpp_diag);
  }
}

void FramePointHessian::AddFrameHess(const FrameHessian2& fh) {
  if (!fh.Ok()) return;
  const int ii = fh.ii();
  const int jj = fh.jj();

  CHECK_GE(ii, 0);
  CHECK_GE(jj, 0);
  CHECK_LT(ii, dim_frames());
  CHECK_LT(jj, dim_frames());
  CHECK_NE(ii, jj);

  // Add each sub block in pose hessian (only upper triangular block)
  Hpp.block<df, df>(ii, ii) += fh.Hii;
  Hpp.block<df, df>(jj, jj) += fh.Hjj;

  if (ii < jj) {
    Hpp.block<df, df>(ii, jj) += fh.Hij;
  } else {
    Hpp.block<df, df>(jj, ii).noalias() += fh.Hij.transpose();
  }

  bp.segment<df>(ii) += fh.bi;
  bp.segment<df>(jj) += fh.bj;

  c += fh.c;
  n += fh.n;
}

void FramePointHessian::AddPatchHess(const PatchHessian2& ph,
                                     const Matrix26d& G0,
                                     const Matrix26d& G1,
                                     const Vector2d& Gd,
                                     const Vector3i& ijh,
                                     int affine_offset) noexcept {
  const int if0 = ijh[0];  // start index of frame 0
  const int if1 = ijh[1];  // start index of frame 1
  const int hid = ijh[2];  // index of depth point

  // G0t * ItI * Gd
  Hpm.block<dp, dd>(if0, hid).noalias() += G0.transpose() * ph.ItI * Gd;
  // G1t * ItI * Gd
  Hpm.block<dp, dd>(if1, hid).noalias() += G1.transpose() * ph.ItI * Gd;

  // Gdt * ItI * Gd
  Hmm[hid] += Gd.transpose() * ph.ItI * Gd;
  // Gdt * Itr
  bm[hid] -= Gd.transpose() * ph.Itr;

  if (affine_offset < 0) return;
  const int ia0 = if0 + dp;  // aff_offset = 0, since it's always left
  const int ia1 = if1 + dp + affine_offset;
  // A0t * I * Gd
  Hpm.block<da, dd>(ia0, hid).noalias() += ph.ItA0.transpose() * Gd;
  // A1t * I * Gd
  Hpm.block<da, dd>(ia1, hid).noalias() += ph.ItA1.transpose() * Gd;
}

void FramePointHessian::Solve() {
  CHECK(ready_);

  // Back-substitute xp to get xm
  xm = Hmm_inv.asDiagonal() * (bm - Hpm.transpose() * xp);
}

void FramePointHessian::Prepare() {
  CHECK(!ready_);
  SafeCwiseInverse(Hmm);  // Hmm_inv
  FillLowerTriangular(Hpp);
  ready_ = true;
}

int FramePointHessian::MargPointsAll(SchurFrameHessian& schur,
                                     int gsize) const {
  CHECK(Ok());
  CHECK(ready_);
  // Hsc = Hpp - Hpm * Hmm^-1 * Hmp
  // bsc = bp  - Hpm * Hmm^-1 * bm
  // Also update n for schur to store number of points marginalized
  const auto n_marged =
      MargPointsToFrames(Hpp, Hpm, Hmm_inv, bp, bm, schur.Hpp, schur.bp, gsize);
  schur.n = n_marged;
  return n_marged;
}

int FramePointHessian::MargPointsRange(SchurFrameHessian& schur,
                                       const cv::Range& range,
                                       int gsize) {
  CHECK(Ok()) << Repr();
  CHECK_GE(range.start, 0);
  CHECK_LT(range.end, dim_points());
  CHECK_GT(range.size(), 0);

  const int i = range.start;
  const int m = range.end - range.start;
  // m stores the number of points being marginalized, which might be different
  // from n, as some points might be considered outlier and have no
  // contribution. Also m is allowed to be 0 since the MargPointsToFramesLower
  // function can handle this case.

  // Manually prepare for marginalization
  FillLowerTriangular(Hpp);
  SafeCwiseInverse(Hmm.segment(i, m));
  const auto n_marged = MargPointsToFrames(Hpp,
                                           Hpm.middleCols(i, m),
                                           Hmm_inv.segment(i, m),
                                           bp,
                                           bm.segment(i, m),
                                           schur.Hpp,
                                           schur.bp,
                                           gsize);
  schur.n = n_marged;
  return n_marged;
}

int MargPointsToFrames(const MatrixXdCRef& Hpp,
                       const MatrixXdCRef& Hpm,
                       const VectorXdCRef& Hmm_inv,
                       const VectorXdCRef& bp,
                       const VectorXdCRef& bm,
                       MatrixXdRef Hsc,
                       VectorXdRef bsc,
                       int gsize) {
  // https://www.robots.ox.ac.uk/~gsibley/Personal/Papers/gsibley-springer2007.pdf
  // [ Hpp Hpm ] [ xp ] = [ bp ]
  // [ Hpm Hmm ] [ xm ] = [ bm ]
  // Hsc = Hpp - Hpm * Hmm^-1 * Hmp
  // bsc = bp  - Hpm * Hmm^-1 * bm

  // Pre-condition
  // 1. Hpp is square/symmetric and matches Hpm and bp
  // 2. Hsc is square and matches Hpp and bsc
  const auto dim_frames = bp.size();
  const auto dim_points = bm.size();
  CHECK_GT(dim_frames, 0);
  CHECK_EQ(Hpp.rows(), dim_frames);
  CHECK_EQ(Hpp.cols(), dim_frames);
  CHECK_EQ(Hpm.rows(), dim_frames);
  CHECK_EQ(Hpm.cols(), dim_points);
  CHECK_EQ(Hmm_inv.size(), dim_points);
  CHECK_EQ(Hpp, Hpp.transpose()) << "\n" << Hpp;

  CHECK_EQ(Hsc.rows(), dim_frames);
  CHECK_EQ(Hsc.cols(), dim_frames);
  CHECK_EQ(bsc.size(), dim_frames);

  // Writing efficient matrix product expressions
  // https://eigen.tuxfamily.org/dox/TopicWritingEfficientProductExpression.html
  // m1.noalias() = m4 + m2 * m3;
  // Should be
  // m1 = m4;
  // m1.noalias() += m2 * m3;
  // Simple benchmark shows no huge difference though, probably because Hpp is
  // much smaller compared to Hpm. However, doing this allows us to nicely
  // handle the case where no points are to be marginalized
  Hsc = Hpp;
  bsc = bp;

  // If no points to marg, then just set Hsc to Hpp and return
  if (dim_points == 0) return 0;

  if (gsize <= 0) {
    // Single thread
    Hsc.noalias() -= Hpm * Hmm_inv.asDiagonal() * Hpm.transpose();
  } else {
    // Parallel
    // First compute Hpm * Hmm^-1 * Hpm^T, only fill lower triangular part
    // Hsc is column major, thus we fill each column per thread
    // for each column in lower triangular part we compute
    // Hsc[i:,i] -= Hpm[i:,:] * (Hmm_inv * Hpm[i,:])
    ParallelFor({0, static_cast<int>(dim_frames), gsize}, [&](int i) {
      const auto r = dim_frames - i;
      Hsc.col(i).tail(r).noalias() -=
          Hpm.bottomRows(r) * (Hmm_inv.asDiagonal() * Hpm.row(i).transpose());
    });
  }

  // Make sure Hsc is symmetric
  FillUpperTriangular(Hsc);

  bsc.noalias() -= Hpm * (Hmm_inv.asDiagonal() * bm);

  // Post-condition
  // 1. Hsc shape doesn't change
  CHECK_EQ(Hsc.rows(), dim_frames);
  CHECK_EQ(Hsc.cols(), dim_frames);
  CHECK_EQ(bsc.size(), dim_frames);
  return static_cast<int>((bm.array() != 0).count());
}

int PointHessian::MapPoints(int npoints) {
  const int dim = npoints * Dim::kPoint;
  const int size = dim * 3;
  data_.resize(size);

  auto* ptr = data_.data();
  new (&H) VectorXdMap(ptr, dim);
  ptr += H.size();
  new (&b) VectorXdMap(ptr, dim);
  ptr += b.size();
  new (&x) VectorXdMap(ptr, dim);
  ptr += x.size();

  CHECK_EQ(ptr - data_.data(), size);
  return size;
}

}  // namespace sv::dsol


///////////////////////////////////////////////////////////////////////////////
// statistics::survival::model::concept::random_failure.hpp                  //
//                                                                           //
//  Copyright 2009 Erwann Rogard. 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_STATISTICS_SURVIVAL_MODEL_CONCEPT_RANDOM_FAILURE_HPP_ER_2009
#define BOOST_STATISTICS_SURVIVAL_MODEL_CONCEPT_RANDOM_FAILURE_HPP_ER_2009
#include <boost/concept_check.hpp>
#include <boost/mpl/empty_base.hpp>
#include <boost/statistics/model/wrap/aggregate/model_covariate_parameter.hpp>

namespace boost{

namespace random{ // fwd
    template<typename T,typename M,typename U>
    T failure_time(M m ,U& u);
}

namespace statistics{
namespace survival{
namespace model{

    template<
        typename T,
        typename M,
        typename X,
        typename P,
        typename U,
        typename B = mpl::empty_base
    >
    class RandomFailure : B
    {
        public:
        typedef statistics::model::model_covariate_parameter_<M,X,P> mcp_;

        BOOST_CONCEPT_USAGE(RandomFailure){

            ft = boost::random::failure_time<T>(mcp,urng);

        }

        private:
        mcp_ mcp;
        U urng;
        T ft;

    };

}// model
}// survival
}// statistics
}// boost

#endif