#include <stdlib.h>
#include <stdio.h>
#include <vector>
#include <string>
#include <sstream>
#include <iostream>
#include <stdexcept>
#include <Eigen/Geometry>
#include <visualization_msgs/Marker.h>
#include <arc_utilities/eigen_helpers.hpp>
#include <arc_utilities/voxel_grid.hpp>
#include <sdf_tools/SDF.h>

#ifndef SDF_HPP
#define SDF_HPP

inline std::vector<uint8_t> FloatToBinary(float value)
{
    uint32_t binary_value = 0;
    memcpy(&binary_value, &value, sizeof(uint32_t));
    std::vector<uint8_t> binary(4);
    // Copy byte 1, least-significant byte
    binary[3] = binary_value & 0x000000ff;
    // Copy byte 2
    binary_value = binary_value >> 8;
    binary[2] = binary_value & 0x000000ff;
    // Copy byte 3
    binary_value = binary_value >> 8;
    binary[1] = binary_value & 0x000000ff;
    // Copy byte 4, most-significant byte
    binary_value = binary_value >> 8;
    binary[0] = binary_value & 0x000000ff;
    return binary;
}

inline float FloatFromBinary(std::vector<uint8_t>& binary)
{
    if (binary.size() != 4)
    {
        std::cerr << "Binary value is not 4 bytes" << std::endl;
        return NAN;
    }
    else
    {
        uint32_t binary_value = 0;
        // Copy in byte 4, most-significant byte
        binary_value = binary_value | binary[0];
        binary_value = binary_value << 8;
        // Copy in byte 3
        binary_value = binary_value | binary[1];
        binary_value = binary_value << 8;
        // Copy in byte 2
        binary_value = binary_value | binary[2];
        binary_value = binary_value << 8;
        // Copy in byte 1, least-significant byte
        binary_value = binary_value | binary[3];
        // Convert binary to float and store
        float field_value = 0.0;
        memcpy(&field_value, &binary_value, sizeof(float));
        return field_value;
    }
}

namespace sdf_tools
{
    class SignedDistanceField
    {
    protected:

        VoxelGrid::VoxelGrid<float> distance_field_;
        std::string frame_;
        bool initialized_;
        bool locked_;

        std::vector<uint8_t> GetInternalBinaryRepresentation(const std::vector<float> &field_data);

        std::vector<float> UnpackFieldFromBinaryRepresentation(std::vector<uint8_t>& binary);

        /*
         * You *MUST* provide valid indices to this function, hence why it is protected (there are safe wrappers available - use them!)
         */
        void FollowGradientsToLocalMaximaUnsafe(VoxelGrid::VoxelGrid<Eigen::Vector3d>& watershed_map, const int64_t x_index, const int64_t y_index, const int64_t z_index) const;

    public:

        EIGEN_MAKE_ALIGNED_OPERATOR_NEW

        inline SignedDistanceField(std::string frame, double resolution, double x_size, double y_size, double z_size, float OOB_value) : initialized_(true), locked_(false)
        {
            frame_ = frame;
            VoxelGrid::VoxelGrid<float> new_field(resolution, x_size, y_size, z_size, OOB_value);
            distance_field_ = new_field;
        }

        inline SignedDistanceField(Eigen::Affine3d origin_transform, std::string frame, double resolution, double x_size, double y_size, double z_size, float OOB_value) : initialized_(true), locked_(false)
        {
            frame_ = frame;
            VoxelGrid::VoxelGrid<float> new_field(origin_transform, resolution, x_size, y_size, z_size, OOB_value);
            distance_field_ = new_field;
        }

        inline SignedDistanceField() : initialized_(false), locked_(false) {}

        inline bool IsInitialized() const
        {
            return initialized_;
        }

        inline bool IsLocked() const
        {
            return locked_;
        }

        inline void Lock()
        {
            locked_ = true;
        }

        inline void Unlock()
        {
            locked_ = false;
        }

        inline float Get(const double x, const double y, const double z) const
        {
            return distance_field_.GetImmutable(x, y, z).first;
        }

        inline float Get3d(const Eigen::Vector3d& location) const
        {
            return distance_field_.GetImmutable3d(location).first;
        }

        inline float Get4d(const Eigen::Vector4d& location) const
        {
            return distance_field_.GetImmutable4d(location).first;
        }

        inline float Get(const int64_t x_index, const int64_t y_index, const int64_t z_index) const
        {
            return distance_field_.GetImmutable(x_index, y_index, z_index).first;
        }

        inline std::pair<float, bool> GetSafe(const double x, const double y, const double z) const
        {
            return distance_field_.GetImmutable(x, y, z);
        }

        inline std::pair<float, bool> GetSafe3d(const Eigen::Vector3d& location) const
        {
            return distance_field_.GetImmutable3d(location);
        }

        inline std::pair<float, bool> GetSafe4d(const Eigen::Vector4d& location) const
        {
            return distance_field_.GetImmutable4d(location);
        }

        inline std::pair<float, bool> GetSafe(const int64_t x_index, const int64_t y_index, const int64_t z_index) const
        {
            return distance_field_.GetImmutable(x_index, y_index, z_index);
        }

        /*
         * Setter functions MUST be used carefully - If you arbitrarily change SDF values, it is not a proper SDF any more!
         *
         * Use of these functions can be prevented by calling SignedDistanceField::Lock() on the SDF, at which point these functions
         * will fail with a warning printed to std_err.
         */
        inline bool Set(const double x, const double y, const double z, float value)
        {
            if (!locked_)
            {
                return distance_field_.SetValue(x, y, z, value);
            }
            else
            {
                std::cerr << "Attempt to set value in locked SDF" << std::endl;
                return false;
            }
        }

        inline bool Set3d(const Eigen::Vector3d& location, float value)
        {
            if (!locked_)
            {
                return distance_field_.SetValue3d(location, value);
            }
            else
            {
                std::cerr << "Attempt to set value in locked SDF" << std::endl;
                return false;
            }
        }

        inline bool Set4d(const Eigen::Vector4d& location, float value)
        {
            if (!locked_)
            {
                return distance_field_.SetValue4d(location, value);
            }
            else
            {
                std::cerr << "Attempt to set value in locked SDF" << std::endl;
                return false;
            }
        }

        inline bool Set(const int64_t x_index, const int64_t y_index, const int64_t z_index, const float value)
        {
            if (!locked_)
            {
                return distance_field_.SetValue(x_index, y_index, z_index, value);
            }
            else
            {
                std::cerr << "Attempt to set value in locked SDF" << std::endl;
                return false;
            }
        }

        inline bool Set(const VoxelGrid::GRID_INDEX& index, const float value)
        {
            if (!locked_)
            {
                return distance_field_.SetValue(index, value);
            }
            else
            {
                std::cerr << "Attempt to set value in locked SDF" << std::endl;
                return false;
            }
        }

        inline bool CheckInBounds3d(const Eigen::Vector3d& location) const
        {
            return distance_field_.GetImmutable3d(location).second;
        }

        inline bool CheckInBounds4d(const Eigen::Vector4d& location) const
        {
            return distance_field_.GetImmutable4d(location).second;
        }

        inline bool CheckInBounds(const double x, const double y, const double z) const
        {
            return distance_field_.GetImmutable(x, y, z).second;
        }

        inline bool CheckInBounds(const VoxelGrid::GRID_INDEX& index) const
        {
            return distance_field_.GetImmutable(index.x, index.y, index.z).second;
        }

        inline bool CheckInBounds(const int64_t x_index, const int64_t y_index, const int64_t z_index) const
        {
            return distance_field_.GetImmutable(x_index, y_index, z_index).second;
        }

        inline double GetXSize() const
        {
            return distance_field_.GetXSize();
        }

        inline double GetYSize() const
        {
            return distance_field_.GetYSize();
        }

        inline double GetZSize() const
        {
            return distance_field_.GetZSize();
        }

        inline double GetResolution() const
        {
            return distance_field_.GetCellSizes()[0];
        }

        inline float GetOOBValue() const
        {
            return distance_field_.GetDefaultValue();
        }

        inline int64_t GetNumXCells() const
        {
            return distance_field_.GetNumXCells();
        }

        inline int64_t GetNumYCells() const
        {
            return distance_field_.GetNumYCells();
        }

        inline int64_t GetNumZCells() const
        {
            return distance_field_.GetNumZCells();
        }

    protected:

        inline std::pair<Eigen::Vector3d, double> GetPrimaryComponentsVector(const Eigen::Vector3d& raw_vector) const
        {
            if (std::abs(raw_vector.x()) > std::abs(raw_vector.y()) && std::abs(raw_vector.x()) > std::abs(raw_vector.z()))
            {
                if (raw_vector.x() >= 0.0)
                {
                    return std::make_pair(Eigen::Vector3d(GetResolution() * 0.5, 0.0, 0.0), GetResolution() * 0.5);
                }
                else
                {
                    return std::make_pair(Eigen::Vector3d(GetResolution() * -0.5, 0.0, 0.0), GetResolution() * 0.5);
                }
            }
            else if (std::abs(raw_vector.y()) > std::abs(raw_vector.x()) && std::abs(raw_vector.y()) > std::abs(raw_vector.z()))
            {
                if (raw_vector.y() >= 0.0)
                {
                    return std::make_pair(Eigen::Vector3d(0.0, GetResolution() * 0.5, 0.0), GetResolution() * 0.5);
                }
                else
                {
                    return std::make_pair(Eigen::Vector3d(0.0, GetResolution() * -0.5, 0.0), GetResolution() * 0.5);
                }
            }
            else if (std::abs(raw_vector.z()) > std::abs(raw_vector.x()) && std::abs(raw_vector.z()) > std::abs(raw_vector.y()))
            {
                if (raw_vector.z() >= 0.0)
                {
                    return std::make_pair(Eigen::Vector3d(0.0, 0.0, GetResolution() * 0.5), GetResolution() * 0.5);
                }
                else
                {
                    return std::make_pair(Eigen::Vector3d(0.0, 0.0, GetResolution() * -0.5), GetResolution() * 0.5);
                }
            }
            else if (std::abs(raw_vector.x()) == std::abs(raw_vector.y()))
            {
                const Eigen::Vector3d temp_vector(raw_vector.x(), raw_vector.y(), 0.0);
                return std::make_pair((temp_vector / (temp_vector.norm())) * std::sqrt((GetResolution() * GetResolution() * 0.25) * 2.0), std::sqrt((GetResolution() * GetResolution() * 0.25) * 2.0));
            }
            else if (std::abs(raw_vector.y()) == std::abs(raw_vector.z()))
            {
                const Eigen::Vector3d temp_vector(0.0, raw_vector.y(), raw_vector.x());
                return std::make_pair((temp_vector / (temp_vector.norm())) * std::sqrt((GetResolution() * GetResolution() * 0.25) * 2.0), std::sqrt((GetResolution() * GetResolution() * 0.25) * 2.0));
            }
            else if (std::abs(raw_vector.x()) == std::abs(raw_vector.z()))
            {
                const Eigen::Vector3d temp_vector(raw_vector.x(), 0.0, raw_vector.z());
                return std::make_pair((temp_vector / (temp_vector.norm())) * std::sqrt((GetResolution() * GetResolution() * 0.25) * 2.0), std::sqrt((GetResolution() * GetResolution() * 0.25) * 2.0));
            }
            else
            {
                return std::make_pair((raw_vector / (raw_vector.norm())) * std::sqrt((GetResolution() * GetResolution() * 0.25) * 3.0), std::sqrt((GetResolution() * GetResolution() * 0.25) * 3.0));
            }
        }

        inline double ComputeAxisMatch(const double axis_value, const double check_value) const
        {
            if ((axis_value >= 0.0) == (check_value >= 0.0))
            {
                return std::abs(check_value - axis_value);
            }
            else
            {
                return -std::abs(check_value - axis_value);
            }
        }

        inline Eigen::Vector3d GetBestMatchSurfaceVector(const Eigen::Vector3d& possible_surfaces_vector, const Eigen::Vector3d& center_to_location_vector) const
        {
            const Eigen::Vector3d location_rejected_on_possible = EigenHelpers::VectorRejection(possible_surfaces_vector, center_to_location_vector);
            // Find the axis with the best-match components
            const double x_axis_match = ComputeAxisMatch(possible_surfaces_vector.x(), location_rejected_on_possible.x());
            const double y_axis_match = ComputeAxisMatch(possible_surfaces_vector.y(), location_rejected_on_possible.y());
            const double z_axis_match = ComputeAxisMatch(possible_surfaces_vector.z(), location_rejected_on_possible.z());
            if ((x_axis_match > y_axis_match) && (x_axis_match > z_axis_match))
            {
                return Eigen::Vector3d(possible_surfaces_vector.x(), 0.0, 0.0);
            }
            else if ((y_axis_match > x_axis_match) && (y_axis_match > z_axis_match))
            {
                return Eigen::Vector3d(0.0, possible_surfaces_vector.y(), 0.0);
            }
            else if ((z_axis_match > x_axis_match) && (z_axis_match > y_axis_match))
            {
                return Eigen::Vector3d(0.0, 0.0, possible_surfaces_vector.z());
            }
            else
            {
                assert(false);
                return possible_surfaces_vector;
            }
        }

        /**
         * @brief GetPrimaryEntrySurfaceVector Estimates the real distance of the provided point, comparing it with the cell center location and gradient vector
         * @param boundary_direction_vector
         * @param center_to_location_vector
         * @return vector from center of voxel to primary entry surface, and magnitude of that vector
         */
        inline std::pair<Eigen::Vector3d, double> GetPrimaryEntrySurfaceVector(const Eigen::Vector3d& boundary_direction_vector, const Eigen::Vector3d& center_to_location_vector) const
        {
            if (boundary_direction_vector.squaredNorm() > std::numeric_limits<double>::epsilon())
            {
                const std::pair<Eigen::Vector3d, double> primary_components_vector_query = GetPrimaryComponentsVector(boundary_direction_vector);
                // If the cell is on a surface
                if (primary_components_vector_query.second == (GetResolution() * 0.5))
                {
                    return primary_components_vector_query;
                }
                // If the cell is on an edge or surface
                else
                {
                    // Pick the best-match of the two/three exposed surfaces
                    return std::make_pair(GetBestMatchSurfaceVector(primary_components_vector_query.first, center_to_location_vector), GetResolution() * 0.5);
                }
            }
            else
            {
                return GetPrimaryComponentsVector(center_to_location_vector);
            }
        }

        inline double EstimateDistanceInternal(const double x, const double y, const double z, const int64_t x_idx, const int64_t y_idx, const int64_t z_idx) const
        {
            const std::vector<double> cell_center = GridIndexToLocation(x_idx, y_idx, z_idx);
            const Eigen::Vector3d cell_center_to_location_vector(x - cell_center[0], y - cell_center[1], z - cell_center[2]);
            const double nominal_sdf_distance = (double)distance_field_.GetImmutable(x_idx, y_idx, z_idx).first;

            // Determine vector from "entry surface" to center of voxel
            // TODO: Needs special handling if there's no gradient to work with
            const std::vector<double> raw_gradient = GetGradient(x_idx, y_idx, z_idx, true);
            const Eigen::Vector3d gradient = EigenHelpers::StdVectorDoubleToEigenVector3d(raw_gradient);
            const Eigen::Vector3d direction_to_boundary = (nominal_sdf_distance >= 0.0) ? -gradient : gradient;
            const std::pair<Eigen::Vector3d, double> entry_surface_information = GetPrimaryEntrySurfaceVector(direction_to_boundary, cell_center_to_location_vector);
            const Eigen::Vector3d& entry_surface_vector = entry_surface_information.first;
            const double minimum_distance_magnitude = entry_surface_information.second;

            // Adjust for calculating distance to boundary of voxels instead of center of voxels
            const double center_adjusted_nominal_distance = (nominal_sdf_distance >= 0.0) ? nominal_sdf_distance - (GetResolution() * 0.5) : nominal_sdf_distance + (GetResolution() * 0.5);
            const double minimum_adjusted_distance = arc_helpers::SpreadValue(center_adjusted_nominal_distance, -minimum_distance_magnitude, 0.0, minimum_distance_magnitude);

            // Account for target location being not at the exact center of the voxel
            const double raw_distance_adjustment = EigenHelpers::VectorProjection(entry_surface_vector, cell_center_to_location_vector).norm();
            const double real_distance_adjustment = (minimum_adjusted_distance >= 0.0) ? -raw_distance_adjustment: raw_distance_adjustment;
            const double final_adjusted_distance = minimum_adjusted_distance + real_distance_adjustment;

            // Perform minimum distance thresholding and error checking
            // TODO: do we need to address this magic number somehow?
            if (std::abs(final_adjusted_distance) < GetResolution() * 0.001)
            {
                return 0.0;
            }
            if ((minimum_adjusted_distance >= 0.0) == (final_adjusted_distance >= 0.0))
            {
                return final_adjusted_distance;
            }
            else
            {
                std::cerr << "Center adjusted nominal distance " << minimum_adjusted_distance << " final adjusted_distance " << final_adjusted_distance << std::endl;
                assert(false && "Mismatched minimum and final adjusted distance signs");
            }
        }

    public:

        inline std::pair<double, bool> EstimateDistance(const double x, const double y, const double z) const
        {
            return EstimateDistance4d(Eigen::Vector4d(x, y, z, 1.0));
        }

        inline std::pair<double, bool> EstimateDistance3d(const Eigen::Vector3d& location) const
        {
            const std::vector<int64_t> indices = LocationToGridIndex3d(location);
            if (indices.size() == 3)
            {
                return std::make_pair(EstimateDistanceInternal(location.x(), location.y(), location.z(), indices[0], indices[1], indices[2]), true);
            }
            else
            {
                return std::make_pair((double)distance_field_.GetOOBValue(), false);
            }
        }

        inline std::pair<double, bool> EstimateDistance4d(const Eigen::Vector4d& location) const
        {
            const std::vector<int64_t> indices = LocationToGridIndex4d(location);
            if (indices.size() == 3)
            {
                return std::make_pair(EstimateDistanceInternal(location(0), location(1), location(2), indices[0], indices[1], indices[2]), true);
            }
            else
            {
                return std::make_pair((double)distance_field_.GetOOBValue(), false);
            }
        }

        inline std::vector<double> GetGradient(const double x, const double y, const double z, const bool enable_edge_gradients=false) const
        {
            return GetGradient4d(Eigen::Vector4d(x, y, z, 1.0), enable_edge_gradients);
        }

        inline std::vector<double> GetGradient3d(const Eigen::Vector3d& location, const bool enable_edge_gradients=false) const
        {
            const std::vector<int64_t> indices = LocationToGridIndex3d(location);
            if (indices.size() == 3)
            {
                return GetGradient(indices[0], indices[1], indices[2], enable_edge_gradients);
            }
            else
            {
                return std::vector<double>();
            }
        }

        inline std::vector<double> GetGradient4d(const Eigen::Vector4d& location, const bool enable_edge_gradients=false) const
        {
            const std::vector<int64_t> indices = LocationToGridIndex4d(location);
            if (indices.size() == 3)
            {
                return GetGradient(indices[0], indices[1], indices[2], enable_edge_gradients);
            }
            else
            {
                return std::vector<double>();
            }
        }

        inline std::vector<double> GetGradient(const VoxelGrid::GRID_INDEX& index, const bool enable_edge_gradients=false) const
        {
            return GetGradient(index.x, index.y, index.z, enable_edge_gradients);
        }

        inline std::vector<double> GetGradient(const int64_t x_index, const int64_t y_index, const int64_t z_index, const bool enable_edge_gradients=false) const
        {
            // Make sure the index is inside bounds
            if ((x_index >= 0) && (y_index >= 0) && (z_index >= 0) && (x_index < GetNumXCells()) && (y_index < GetNumYCells()) && (z_index < GetNumZCells()))
            {
                // Make sure the index we're trying to query is one cell in from the edge
                if ((x_index > 0) && (y_index > 0) && (z_index > 0) && (x_index < (GetNumXCells() - 1)) && (y_index < (GetNumYCells() - 1)) && (z_index < (GetNumZCells() - 1)))
                {
                    double inv_twice_resolution = 1.0 / (2.0 * GetResolution());
                    double gx = (Get(x_index + 1, y_index, z_index) - Get(x_index - 1, y_index, z_index)) * inv_twice_resolution;
                    double gy = (Get(x_index, y_index + 1, z_index) - Get(x_index, y_index - 1, z_index)) * inv_twice_resolution;
                    double gz = (Get(x_index, y_index, z_index + 1) - Get(x_index, y_index, z_index - 1)) * inv_twice_resolution;
                    return std::vector<double>{gx, gy, gz};
                }
                // If we're on the edge, handle it specially
                else if (enable_edge_gradients)
                {
                    // Get the "best" indices we can use
                    int64_t low_x_index = std::max((int64_t)0, x_index - 1);
                    int64_t high_x_index = std::min(GetNumXCells() - 1, x_index + 1);
                    int64_t low_y_index = std::max((int64_t)0, y_index - 1);
                    int64_t high_y_index = std::min(GetNumYCells() - 1, y_index + 1);
                    int64_t low_z_index = std::max((int64_t)0, z_index - 1);
                    int64_t high_z_index = std::min(GetNumZCells() - 1, z_index + 1);
                    // Compute the axis increments
                    double x_increment = (high_x_index - low_x_index) * GetResolution();
                    double y_increment = (high_y_index - low_y_index) * GetResolution();
                    double z_increment = (high_z_index - low_z_index) * GetResolution();
                    // Compute the gradients for each axis - by default these are zero
                    double gx = 0.0;
                    double gy = 0.0;
                    double gz = 0.0;
                    // Only if the increments are non-zero do we compute the gradient of an axis
                    if (x_increment > 0.0)
                    {
                        double inv_x_increment = 1.0 / x_increment;
                        double high_x_value = Get(high_x_index, y_index, z_index);
                        double low_x_value = Get(low_x_index, y_index, z_index);
                        // Compute the gradient
                        gx = (high_x_value - low_x_value) * inv_x_increment;
                    }
                    if (y_increment > 0.0)
                    {
                        double inv_y_increment = 1.0 / y_increment;
                        double high_y_value = Get(x_index, high_y_index, z_index);
                        double low_y_value = Get(x_index, low_y_index, z_index);
                        // Compute the gradient
                        gy = (high_y_value - low_y_value) * inv_y_increment;
                    }
                    if (z_increment > 0.0)
                    {
                        double inv_z_increment = 1.0 / z_increment;
                        double high_z_value = Get(x_index, y_index, high_z_index);
                        double low_z_value = Get(x_index, y_index, low_z_index);
                        // Compute the gradient
                        gz = (high_z_value - low_z_value) * inv_z_increment;
                    }
                    // Assemble and return the computed gradient
                    return std::vector<double>{gx, gy, gz};
                }
                // Edge gradients disabled, return no gradient
                else
                {
                    return std::vector<double>();
                }
            }
            // If we're out of bounds, return no gradient
            else
            {
                return std::vector<double>();
            }
        }

        inline Eigen::Vector3d ProjectOutOfCollision(const double x, const double y, const double z, const double stepsize_multiplier = 1.0 / 10.0) const
        {
            const Eigen::Vector4d result = ProjectOutOfCollision4d(Eigen::Vector4d(x, y, z, 1.0), stepsize_multiplier);
            return result.head<3>();
        }

        inline Eigen::Vector3d ProjectOutOfCollisionToMinimumDistance(const double x, const double y, const double z, const double minimum_distance, const double stepsize_multiplier = 1.0 / 10.0) const
        {
            const Eigen::Vector4d result = ProjectOutOfCollisionToMinimumDistance4d(Eigen::Vector4d(x, y, z, 1.0), minimum_distance, stepsize_multiplier);
            return result.head<3>();
        }

        inline Eigen::Vector3d ProjectOutOfCollision3d(const Eigen::Vector3d& location, const double stepsize_multiplier = 1.0 / 10.0) const
        {
            return ProjectOutOfCollision(location.x(), location.y(), location.z(), stepsize_multiplier);
        }

        inline Eigen::Vector3d ProjectOutOfCollisionToMinimumDistance3d(const Eigen::Vector3d& location, const double minimum_distance, const double stepsize_multiplier = 1.0 / 10.0) const
        {
            return ProjectOutOfCollisionToMinimumDistance(location.x(), location.y(), location.z(), minimum_distance, stepsize_multiplier);
        }

        inline Eigen::Vector4d ProjectOutOfCollision4d(const Eigen::Vector4d& location, const double stepsize_multiplier = 1.0 / 10.0) const
        {
            return ProjectOutOfCollisionToMinimumDistance4d(location, 0.0, stepsize_multiplier);
        }

        inline Eigen::Vector4d ProjectOutOfCollisionToMinimumDistance4d(const Eigen::Vector4d& location, const double minimum_distance, const double stepsize_multiplier = 1.0 / 10.0) const
        {
            Eigen::Vector4d mutable_location = location;
            const bool enable_edge_gradients = true;

            double sdf_dist = EstimateDistance4d(mutable_location).first;
            if (sdf_dist < minimum_distance && CheckInBounds4d(location))
            {
                while (sdf_dist < minimum_distance)
                {
                    const std::vector<double> gradient = GetGradient4d(mutable_location, enable_edge_gradients);
                    const Eigen::Vector3d grad_eigen = EigenHelpers::StdVectorDoubleToEigenVector3d(gradient);

                    assert(grad_eigen.norm() > GetResolution() / 4.0); // Sanity check
                    mutable_location.head<3>() += grad_eigen.normalized() * GetResolution() * stepsize_multiplier;

                    sdf_dist = EstimateDistance4d(mutable_location).first;
                }
            }

            return mutable_location;
        }

        inline const Eigen::Affine3d& GetOriginTransform() const
        {
            return distance_field_.GetOriginTransform();
        }

        inline const Eigen::Affine3d& GetInverseOriginTransform() const
        {
            return distance_field_.GetInverseOriginTransform();
        }

        inline std::string GetFrame() const
        {
            return frame_;
        }

        inline std::vector<int64_t> LocationToGridIndex3d(const Eigen::Vector3d& location) const
        {
            return distance_field_.LocationToGridIndex3d(location);
        }

        inline std::vector<int64_t> LocationToGridIndex4d(const Eigen::Vector4d& location) const
        {
            return distance_field_.LocationToGridIndex4d(location);
        }

        inline std::vector<int64_t> LocationToGridIndex(const double x, const double y, const double z) const
        {
            return distance_field_.LocationToGridIndex(x, y, z);
        }

        inline std::vector<double> GridIndexToLocation(const VoxelGrid::GRID_INDEX& index) const
        {
            return distance_field_.GridIndexToLocation(index);
        }

        inline std::vector<double> GridIndexToLocation(const int64_t x_index, const int64_t y_index, const int64_t z_index) const
        {
            return distance_field_.GridIndexToLocation(x_index, y_index, z_index);
        }

        inline std::vector<double> GridIndexToLocation(std::vector<int64_t> index) const
        {
            return distance_field_.GridIndexToLocation(index[0], index[1], index[2]);
        }

        bool SaveToFile(const std::string& filepath);

        bool LoadFromFile(const std::string& filepath);

        sdf_tools::SDF GetMessageRepresentation();

        bool LoadFromMessageRepresentation(sdf_tools::SDF& message);

        visualization_msgs::Marker ExportForDisplay(float alpha = 0.01f) const;

        visualization_msgs::Marker ExportForDisplayCollisionOnly(float alpha = 0.01f) const;

        visualization_msgs::Marker ExportForDebug(float alpha = 0.5f) const;

        /*
         * The following function can be *VERY EXPENSIVE* to compute, since it performs gradient ascent across the SDF
         */
        VoxelGrid::VoxelGrid<Eigen::Vector3d> ComputeLocalMaximaMap() const;

        inline bool GradientIsEffectiveFlat(const Eigen::Vector3d& gradient) const
        {
            // A gradient is at a local maxima if the absolute value of all components (x,y,z) are less than 1/2 SDF resolution
            double half_resolution = GetResolution() * 0.5;
            if (fabs(gradient.x()) <= half_resolution && fabs(gradient.y()) <= half_resolution && fabs(gradient.z()) <= half_resolution)
            {
                return true;
            }
            else
            {
                return false;
            }
        }

        inline VoxelGrid::GRID_INDEX GetNextFromGradient(const VoxelGrid::GRID_INDEX& index, const Eigen::Vector3d& gradient) const
        {
            // Given the gradient, pick the "best fit" of the 26 neighboring points
            VoxelGrid::GRID_INDEX next_index = index;
            double half_resolution = GetResolution() * 0.5;
            if (gradient.x() > half_resolution)
            {
                next_index.x++;
            }
            else if (gradient.x() < -half_resolution)
            {
                next_index.x--;
            }
            if (gradient.y() > half_resolution)
            {
                next_index.y++;
            }
            else if (gradient.y() < -half_resolution)
            {
                next_index.y--;
            }
            if (gradient.z() > half_resolution)
            {
                next_index.z++;
            }
            else if (gradient.z() < -half_resolution)
            {
                next_index.z--;
            }
            return next_index;
        }
    };
}

#endif // SDF_HPP


/*
 * Copyright (C) 2012-2014, Computing Systems Laboratory (CSLab), NTUA.
 * Copyright (C) 2012-2014, Athena Elafrou
 * All rights reserved.
 *
 * This file is distributed under the BSD License. See LICENSE.txt for details.
 */

/**
 * \file Rcm.hpp
 * \brief Reverse Cuthill-Mckee Ordering Interface
 *
 * \author Computing Systems Laboratory (CSLab), NTUA
 * \date 2011–2014
 * \copyright This file is distributed under the BSD License. See LICENSE.txt
 * for details.
 */

#ifndef SPARSEX_INTERNALS_RCM_HPP
#define SPARSEX_INTERNALS_RCM_HPP

#include <sparsex/internals/Csr.hpp>
#include <sparsex/internals/Mmf.hpp>
#include <sparsex/internals/logger/Logger.hpp>
#include <boost/config.hpp>
#include <boost/graph/adjacency_list.hpp>
#include <boost/graph/cuthill_mckee_ordering.hpp>
#include <boost/graph/properties.hpp>
#include <boost/graph/bandwidth.hpp>

using namespace sparsex::io;

namespace sparsex {
  namespace utilities {

    using namespace boost;

#define bad_reorder 0

    typedef pair<size_t, size_t> Pair;
    typedef adjacency_list<vecS, vecS, undirectedS,
			   property<vertex_color_t, default_color_type,
				    property<vertex_degree_t,int> > > Graph;
    typedef graph_traits<Graph>::vertex_descriptor Vertex;
    typedef graph_traits<Graph>::vertices_size_type size_type;

    template<typename IndexType, typename ValueType>
    struct ColumnCompare {
    public:
      bool operator() (const pair<IndexType, ValueType> &lhs,
		       const pair<IndexType, ValueType> &rhs) const
      {
	if (lhs.first < rhs.first) return true;
	if (lhs.first > rhs.first) return false;
	return false;
      }
    };

    /**
     *  @return perm      permutation from the old ordering to the new one.
     *  @return inv_perm  permutation from the new ordering to the old one.
     *  @param  graph     the incidence graph of the matrix.
     */
    void FindPerm(vector<size_t>& perm, vector<size_t>& inv_perm, Graph& graph);

    /**
     *  Loads matrix from an MMF file and applies the reverse Cuthill-Mckee
     *  reordering algorithm
     *
     *  @param file_name   name of the file where the matrix is kept.
     *  @param rt_config   runtime configuration.
     *  @param mat         handler of MMF class.
     *  @param graph       the incidence graph of the matrix.
     *  @param perm        permutation from the old ordering to the new one.
     *  @return            spm class object with the characteristics of the
     *                     matrix.
     */
    template<typename IndexType, typename ValueType>
    void ReorderMat_MMF(MMF<IndexType, ValueType>& mat,
			const vector<size_t>& perm);
    template<typename IndexType, typename ValueType>
    Graph& ConstructGraph_MMF(Graph& graph, MMF<IndexType, ValueType>& mat);
    template<typename IndexType, typename ValueType>
    void DoReorder_RCM(MMF<IndexType, ValueType>& mat, vector<size_t>& perm);

    /**
     *  Loads matrix from CSR format and applies the reverse Cuthill-Mckee
     *  reordering algorithm
     *
     *  @param rowptr      array "rowptr" of CSR format.
     *  @param colind      array "colind" of CSR format.
     *  @param values      array "values" of CSR format.
     *  @param nr_rows     number of rows.
     *  @param nr_cols     number of columns.
     *  @param zero_based  indexing.
     *  @param rt_config   runtime configuration.
     *  @param perm        permutation from the old ordering to the new one.
     *  @param inv_perm    permutation from the new ordering to the old one.
     *  @param graph       the incidence graph of the matrix.
     *  @return            spm class object with the characteristics of the
     *                     matrix.
     */
    template<typename IndexType, typename ValueType>
    void ReorderMat_CSR(CSR<IndexType, ValueType>& mat,
			const vector<size_t>& perm,
			vector<size_t>& inv_perm);
    template<typename IterT>
    Graph& ConstructGraph_CSR(Graph& graph, IterT& iter, const IterT& iter_end,
			      size_t nr_nzeros, bool symmetric);
    template<typename IndexType, typename ValueType>
    void DoReorder_RCM(CSR<IndexType, ValueType>& mat, vector<size_t>& perm);

    /*
     * Implementation of RCM interface
     */
    void FindPerm(vector<size_t>& perm, vector<size_t>& invperm, Graph& graph)
    {
      graph_traits<Graph>::vertex_iterator ui, ui_end;
      property_map<Graph,vertex_degree_t>::type deg = get(vertex_degree, graph);
      for (boost::tie(ui, ui_end) = vertices(graph); ui != ui_end; ++ui) {
        deg[*ui] = degree(*ui, graph);
      }

      // Original ordering
      property_map<Graph, vertex_index_t>::type
        index_map = get(vertex_index, graph);

      size_t ob = bandwidth(graph);
      LOG_INFO << "Original Bandwidth: " << ob << "\n";

      vector<Vertex> inv_perm(num_vertices(graph));
      perm.resize(num_vertices(graph));
      invperm.reserve(num_vertices(graph));

      // Reverse Cuthill Mckee Ordering
      cuthill_mckee_ordering(graph, inv_perm.rbegin(), get(vertex_color, graph),
			     make_degree_map(graph));

      for (size_t i = 0; i != inv_perm.size(); ++i) {
        perm[index_map[inv_perm[i]]] = i;
        invperm.push_back(inv_perm[i]);
      }

      // cout << "Permutation: ";
      // for (size_t i = 0; i < inv_perm.size(); i++) {
      //     cout << inv_perm[i] << " ";
      // }
      // cout << endl;

      size_t nb = bandwidth(graph, make_iterator_property_map
			    (&perm[0], index_map, perm[0]));
      LOG_INFO << "Final Bandwidth: " << nb << "\n";
    }

    template<typename IndexType, typename ValueType>
    Graph& ConstructGraph_MMF(Graph& graph, MMF<IndexType, ValueType>& mat)
    {
      // The flag must be set before using the MMF iterator
      mat.SetReordered(true);

      // Make a pessimistic guess for nr_edges
      size_t nr_edges;
      if (mat.IsSymmetric()) {
        // If main diagonal is full (nr_nzeros - nr_rows) / 2
        nr_edges = mat.GetNrNonzeros() / 2;
      } else {
        nr_edges = mat.GetNrNonzeros();
      }

      Pair *edges = new Pair[nr_edges];

      typename MMF<IndexType, ValueType>::iterator iter = mat.begin();
      typename MMF<IndexType, ValueType>::iterator iter_end = mat.end();
      size_t index = 0;
      if (!mat.IsSymmetric() && !mat.IsColWise()) {
        mat.InitMatrix(mat.GetNrNonzeros());
        for (; iter != iter_end; ++iter) {
	  mat.InsertElement(*iter);
	  if ((*iter).GetRow() != (*iter).GetCol())
	    edges[index++] = make_pair((*iter).GetRow() - 1,
				       (*iter).GetCol() - 1);
        }
      } else {
        for (;iter != iter_end; ++iter)
	  if ((*iter).GetRow() < (*iter).GetCol())
	    edges[index++] = make_pair((*iter).GetRow() - 1,
				       (*iter).GetCol() - 1);
      }

      // index -> actual nr_edges
      assert(index <= nr_edges);
      if (index == 0) {
        LOG_WARNING << "no reordering available for this matrix\n";
        delete[] edges;
        throw bad_reorder;
      }

      for (size_t i = 0; i < index; ++i) {
        add_edge(edges[i].first, edges[i].second, graph);
      }

      delete[] edges;
      return graph;
    }

    template<typename IndexType, typename ValueType>
    void ReorderMat_MMF(MMF<IndexType, ValueType>& mat,
			const vector<size_t>& perm)
    {
      for (size_t i = 0; i < mat.GetNrNonzeros(); i++) {
        pair<IndexType, ValueType> coord = mat.GetCoordinates(i);
        mat.SetCoordinates(i, perm[coord.first-1] + 1,
                           perm[coord.second-1] + 1);
      }

      mat.Sort();
    }

    template<typename IndexType, typename ValueType>
    void DoReorder_RCM(MMF<IndexType, ValueType>& mat, vector<size_t> &perm)
    {
      vector<size_t> inv_perm;

      LOG_INFO << "Reordering input matrix...\n";
      // Construct graph
      Graph graph(mat.GetNrRows());
      try {
        graph = ConstructGraph_MMF(graph, mat);
      } catch (int e) {
        mat.ResetStream();
        mat.SetReordered(false);
        LOG_WARNING << "reordering failed\n";
        return;
      }
      // Find permutation
      FindPerm(perm, inv_perm, graph);
      // Reorder original matrix
      ReorderMat_MMF(mat, perm);
      LOG_INFO << "Reordering complete\n";
    }

    template<typename IterT>
    Graph& ConstructGraph_CSR(Graph& graph, IterT& iter, const IterT& iter_end,
			      size_t nr_nzeros, bool symmetric)
    {
      // make a pessimistic guess for nr_edges
      size_t nr_edges;
      if (symmetric) {    //if main diagonal is full (nr_nzeros - nr_rows) / 2
        nr_edges = nr_nzeros / 2;
      } else {
        nr_edges = nr_nzeros;
      }

      Pair *edges = new Pair[nr_edges];
      size_t index = 0;
      if (symmetric) {
        for (;iter != iter_end; ++iter) {
	  if ((*iter).GetRow() < (*iter).GetCol()) {
	    edges[index++] = make_pair(((*iter).GetRow()) - 1,
				       ((*iter).GetCol()) - 1);
	  }
        }
      } else {
        for (;iter != iter_end; ++iter) {
	  if ((*iter).GetRow() != (*iter).GetCol()) {
	    edges[index++] = make_pair(((*iter).GetRow()) - 1,
				       ((*iter).GetCol()) - 1);
	  }
        }
      }

      // index -> actual nr_edges
      assert(index <= nr_edges);

      if (index == 0) {
        LOG_WARNING << "no reordering available for this matrix\n";
        delete[] edges;
        throw bad_reorder;
      }

      for (size_t i = 0; i < index; i++) {
        add_edge(edges[i].first, edges[i].second, graph);
      }

      delete[] edges;
      return graph;
    }

    template<typename IndexType, typename ValueType>
    void ReorderMat_CSR(CSR<IndexType, ValueType>& mat,
			const vector<size_t>& perm,
			vector<size_t>& inv_perm)
    {
      // Have to work on a copy of colind and values if CSR structures are not
      // be shared
      IndexType *new_colind = new IndexType[mat.GetNrNonzeros()];
      ValueType *new_values = new ValueType[mat.GetNrNonzeros()];
      memcpy(new_values, mat.values_, sizeof(ValueType) * mat.GetNrNonzeros());

      // Apply permutation only to colind
      for (size_t i = 0; i < mat.GetNrNonzeros(); i++) {
        new_colind[i] = perm[mat.colind_[i] - !mat.IsZeroBased()] +
	  !mat.IsZeroBased();
      }

      mat.colind_ = new_colind;
      mat.values_ = new_values;

      // Simultaneously sort colind and values per row
      for (size_t i = 0; i < mat.GetNrRows(); i++) {
        sort(mat.row_begin(i), mat.row_end(i),
             ColumnCompare<IndexType, ValueType>());
      }
      mat.SetReordered(inv_perm);
      //assert(inv_perm.capacity() == 0);
    }

    template<typename IndexType, typename ValueType>
    void DoReorder_RCM(CSR<IndexType, ValueType>& mat, vector<size_t>& perm)
    {
      typename CSR<IndexType, ValueType>::iterator iter = mat.begin();
      typename CSR<IndexType, ValueType>::iterator iter_end = mat.end();
      vector<size_t> inv_perm;

      LOG_INFO << "Reordering input matrix...\n";
      // Construct graph
      Graph graph(mat.GetNrRows());
      try {
        graph = ConstructGraph_CSR(graph, iter, iter_end, mat.GetNrNonzeros(),
                                   mat.IsSymmetric());
      } catch (int e) {
        LOG_INFO << "reordering failed\n";
        return;
      }
      // Find permutation
      FindPerm(perm, inv_perm, graph);
      // Reorder original matrix
      ReorderMat_CSR<IndexType, ValueType>(mat, perm, inv_perm);
      LOG_INFO << "Reordering complete\n";
    }

  } // end of namespace utilities
} // end of namespace sparsex

#endif // SPARSEX_INTERNALS_RCM_HPP


/* * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * */
/*                                                                           */
/*  This file is part of the library KASKADE 7                               */
/*    see http://www.zib.de/projects/kaskade7-finite-element-toolbox         */
/*                                                                           */
/*  Copyright (C) 2002-2012 Zuse Institute Berlin                            */
/*                                                                           */
/*  KASKADE 7 is distributed under the terms of the ZIB Academic License.    */
/*    see $KASKADE/academic.txt                                              */
/*                                                                           */
/* * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * */

#ifndef ALGORITHM_HIERARCHIC_ERROR_ESTIMATOR_HH
#define ALGORITHM_HIERARCHIC_ERROR_ESTIMATOR_HH

#include <algorithm>
#include <iostream>
#include <utility>
#include <vector>

#include <boost/timer/timer.hpp>
#include <boost/fusion/include/at_c.hpp>

#include "algorithm/newton_bridge.hh"
#include "linalg/jacobiPreconditioner.hh"
#include "fem/hierarchicErrorEstimator.hh"
#include "fem/variables.hh"
#include "linalg/direct.hh"
#include "linalg/tcg.hh"
#include "utilities/enums.hh"

// forward declarations
namespace Dune
{
  struct InverseOperatorResult;
  template <class,int> class FieldVector;
}

namespace Kaskade
{
  namespace HierarchicErrorEstimator_Detail{
    bool biggerThanAbs(const std::pair<double,int>& x1, const std::pair<double,int>&  x2)
    {
      return fabs(x1.first)>fabs(x2.first);
    }

    double add(const double& x1, const std::pair<double,int>&  x2)
    {
      return x1+x2.first;
    }
  }

  template <class CorrectionVector>
  struct LeaveRHS
  {
    explicit LeaveRHS(CorrectionVector const&) {}

    template <class RhsVector>
    void operator()(RhsVector& rhs)
    {}
  };

  template <class CorrectionVector>
  struct SubstractCorrection
  {
    explicit SubstractCorrection(CorrectionVector const& cor_)
    : cor(cor_)
    {}

    template <class RhsVector>
    void operator()(RhsVector& rhs)
    {
      rhs -= cor;
    }

  private:
    CorrectionVector const& cor;
  };

  template <class Functional, /*class LF,*/ class ExtensionVariableSetDescription, class ExtensionSpace, class NormFunctional=Functional, template <class> class AdjustRHS=LeaveRHS>
  class HierarchicalBasisErrorEstimator : public AbstractHierarchicalErrorEstimator
  {
    static constexpr int yId = 1, uId = 0, pId = 2;
    static constexpr int noOfVariables = ExtensionVariableSetDescription::noOfVariables;
    //    typedef typename LF::AnsatzVars LVars;
  public:
    typedef typename Functional::Scalar Scalar;
    typedef typename Functional::AnsatzVars::VariableSet VariableSet;
    static constexpr int dim = VariableSet::Descriptions::Grid::dimension;
    typedef HierarchicErrorEstimator<LinearizationAt<Functional>,ExtensionVariableSetDescription,ExtensionVariableSetDescription,HierarchicErrorEstimatorDetail::TakeAllD2<LinearizationAt<Functional> > > ErrorEstimator;
    typedef VariationalFunctionalAssembler<ErrorEstimator> Assembler;
    typedef typename ExtensionVariableSetDescription::template CoefficientVectorRepresentation<0,2>::type CoefficientVector02;
    typedef typename ExtensionVariableSetDescription::template CoefficientVectorRepresentation<0,noOfVariables>::type CoefficientVector;
    typedef typename ExtensionVariableSetDescription::GridView::template Codim<0>::Iterator CellIterator;

    //    HierarchicalBasisErrorEstimator(Functional& f_, LF& lf_, LVars& lvars, ExtensionVariableSetDescription& extensionVariableSetDescription_, ExtensionSpace& extensionSpace_, Scalar fraction=0.7, bool verbose_=false)
    //    : f(f_), lf(lf_), normFunctional(f_), extensionVariableSetDescription(extensionVariableSetDescription_), extensionSpace(extensionSpace_), assembler(extensionVariableSetDescription.spaces),
    //    squaredFraction(fraction*fraction), verbose(verbose_)
    //    {}

    HierarchicalBasisErrorEstimator(Functional& f_, /*LF& lf_,*/ /*LVars& lvars_,*/ NormFunctional& normFunctional_, ExtensionVariableSetDescription& extensionVariableSetDescription_,
        ExtensionSpace& extensionSpace_, Scalar fraction=0.7, bool verbose_=false)
    : f(f_), /*lf(lf_), lvars(lvars_),*/ normFunctional(normFunctional_), extensionVariableSetDescription(extensionVariableSetDescription_), extensionSpace(extensionSpace_), assembler(extensionVariableSetDescription.spaces),
      squaredFraction(fraction*fraction), verbose(verbose_)
    {}

    virtual ~HierarchicalBasisErrorEstimator(){}

    void operator()(AbstractVector const& x_, AbstractVector const& dx_, int step, AbstractVector const& lowerOrderRhs)
    {
      if(verbose) std::cout << "ERROR ESTIMATOR: Start." << std::endl;
      Bridge::Vector<VariableSet> const& x = dynamic_cast<const Bridge::Vector<VariableSet>&>(x_);
      Bridge::Vector<VariableSet> const& dx = dynamic_cast<const Bridge::Vector<VariableSet>&>(dx_);
      //    Bridge::Vector<VariableSet> const& loRhs = dynamic_cast<const Bridge::Vector<VariableSet>&>(lowerOrderRhs);

      boost::timer::cpu_timer atimer;
      assembler.assemble(ErrorEstimator(LinearizationAt<Functional>(f,x.get()),dx.get()));
      if(verbose) std::cout << "ERROR ESTIMATOR: assembly time: " << boost::timer::format(atimer.elapsed(), 2, "%t") << "s" << std::endl;

      typedef AssembledGalerkinOperator<Assembler> Operator;
      typedef AssembledGalerkinOperator<Assembler,0,2,0,2> CorrectionOperator;
      typedef typename Functional::AnsatzVars::template CoefficientVectorRepresentation<0,2>::type CorrectionVector;

      Operator op(assembler);
      //      CorrectionOperator(assembler);

      CoefficientVector estRhs(assembler.rhs());
      std::vector<Scalar> tmpRepVec;
      IstlInterfaceDetail::toVector(estRhs,tmpRepVec);
      //estRhs = Bridge::getImpl<VariableSet>(lowerOrderRhs);
      std::cout << "rhs entries:" << std::endl;
      Scalar maxRhs=std::numeric_limits<Scalar>::min(), minRhs=std::numeric_limits<Scalar>::max(), rhsSum=0;
      for(size_t i=0; i<tmpRepVec.size(); ++i){
        if(std::isnan(tmpRepVec[i])) std::cout << "nan at " << i << std::endl;

        Scalar val = tmpRepVec[i];
        if(val < minRhs) minRhs = val;
        if(val > maxRhs) maxRhs = val;
        rhsSum += val*val;
      }
      std::cout << "max rhs: " << maxRhs << std::endl;
      std::cout << "min rhs: " << minRhs << std::endl;
      std::cout << "sum: " << rhsSum << std::endl;
      // possible adjustment of rhs
      //adjustRHS(estRhs);

      boost::timer::cpu_timer timer;
      CoefficientVector estSol(ExtensionVariableSetDescription::template CoefficientVectorRepresentation<0,noOfVariables>::init(extensionVariableSetDescription));

      directInverseOperator(op,DirectType::UMFPACK,MatrixProperties::GENERAL).apply(estRhs,estSol);

      if(verbose) std::cout << "ERROR ESTIMATOR: computation time: " << boost::timer::format(timer.elapsed(), 2, "%t") << "s" << std::endl;
      typename ExtensionVariableSetDescription::VariableSet mySol(extensionVariableSetDescription);
      mySol = estSol;
      using namespace boost::fusion;

      Scalar maxVal = std::numeric_limits<Scalar>::min(),
          minVal = std::numeric_limits<Scalar>::max(),
          sum = 0;

      if(verbose)
      {
        for(int i=0; i<at_c<yId>(mySol.data).size();  ++i)
        {
          for(int j=0; j<at_c<yId>(mySol.data).coefficients()[i].size; ++j)
          {
            Scalar val = at_c<yId>(mySol.data).coefficients()[i][j];
            sum += val*val;
            if(val > maxVal) maxVal = val;
            if(val < minVal) minVal = val;
          }
        }
        std::cout << "y: minimal value: " << minVal << std::endl;
        std::cout << "y: maximal value: " << maxVal << std::endl;
        std::cout << "y: l2-sum: " << sum << std::endl;
        Scalar tmpMax = maxVal, tmpMin = minVal, tmpSum = sum;
        maxVal = std::numeric_limits<Scalar>::min();
        minVal = std::numeric_limits<Scalar>::max();
        sum = 0;

        for(int i=0; i<at_c<uId>(mySol.data).size();  ++i)
        {
          for(int j=0; j<at_c<uId>(mySol.data).coefficients()[i].size; ++j)
          {
            Scalar val = at_c<uId>(mySol.data).coefficients()[i][0];
            sum += val*val;
            if(val > maxVal) maxVal = val;
            if(val < minVal) minVal = val;
          }
        }
        std::cout << "u: minimal value: " << minVal << std::endl;
        std::cout << "u: maximal value: " << maxVal << std::endl;
        std::cout << "u: l2-sum: " << sum << std::endl;
        tmpMax = std::max(tmpMax,maxVal), tmpMin = std::min(tmpMin,minVal), tmpSum += sum;
        maxVal = std::numeric_limits<Scalar>::min();
        minVal = std::numeric_limits<Scalar>::max();
        sum = 0;

        for(int i=0; i<at_c<pId>(mySol.data).size();  ++i)
        {
          for(int j=0; j<at_c<pId>(mySol.data).coefficients()[i].size; ++j)
          {
            Scalar val = at_c<pId>(mySol.data).coefficients()[i][j];
            sum += val*val;
            if(val > maxVal) maxVal = val;
            if(val < minVal) minVal = val;
          }
        }
        std::cout << "p: minimal value: " << minVal << std::endl;
        std::cout << "p: maximal value: " << maxVal << std::endl;
        std::cout << "p: l2-sum: " << sum << std::endl;

        tmpSum += sum;
        std::cout << "minimal value: " << std::min(tmpMin,minVal) << std::endl;
        std::cout << "maximal value: " << std::max(tmpMax,maxVal) << std::endl;
        std::cout << "l2-sum: " << tmpSum << std::endl;
      }
      // Transfer error indicators to cells.
      auto const& is = extensionSpace.gridManager().grid().leafIndexSet();
      errorDistribution.clear();
      errorDistribution.resize(is.size(0),std::make_pair(0.0,0));
      Scalar errLevel(0.0), minRefine(0.2);

      typedef ErrorDistribution<NormFunctional,ExtensionVariableSetDescription> EnergyError;
      EnergyError energyError(normFunctional,x.get(),mySol);
      typedef VariationalFunctionalAssembler<LinearizationAt<EnergyError> > EnergyErrorAssembler;
      EnergyErrorAssembler eeAssembler(extensionSpace.gridManager(), energyError.getSpaces());

      eeAssembler.assemble(linearization(energyError,x.get()), EnergyErrorAssembler::RHS);
      typename EnergyError::ErrorVector distError( eeAssembler.rhs() );
      typename EnergyError::AnsatzVars::VariableSet mde( energyError.getVariableSetDescription() );
      mde = distError;

      if(verbose)
      {
        std::string name = "errorDistribution_";
        name += boost::lexical_cast<std::string>(step);
        writeVTKFile(mde.descriptions.gridView,mde.descriptions,mde,name);
      }
      // Fehler bzgl L_xx : ( sum_i estSol_i * (L_xx * estSol)_i )^{1/2}  (nur die x - Komponente, nicht p)
      // Lokalisierung: einfachste Idee: v_i = estSol_i * (L_xx * estSol)_i
      // Aufteilen von v_i auf die einzelnen Zellen
      CellIterator cend = extensionVariableSetDescription.gridView.template end<0>();
      for (CellIterator ci=extensionVariableSetDescription.gridView.template begin<0>(); ci!=cend; ++ci)
        errorDistribution[is.index(*ci)] = std::make_pair( fabs(boost::fusion::at_c<0>(mde.data).value(*ci,Dune::FieldVector<Scalar,dim>(0.3))) , is.index(*ci));

      totalErrorSquared = std::accumulate(errorDistribution.begin(), errorDistribution.end(), 0.0, HierarchicErrorEstimator_Detail::add);
    }

    void refineGrid()
    {
      //      int lastIndexForDoubleRefinement = -1;
      std::sort(errorDistribution.begin(), errorDistribution.end(), HierarchicErrorEstimator_Detail::biggerThanAbs);

      Scalar bulkCriterionTolerance = squaredFraction*totalErrorSquared;
      if(verbose)
      {
        std::cout << "ERROR ESTIMATOR: totalErrorSquared: " << totalErrorSquared << std::endl;
        std::cout << "ERROR ESTIMATOR: bulkCriterionTolerance: " << bulkCriterionTolerance << std::endl;
      }
      Scalar bulkErrorSquared = 0;

      std::vector<std::pair<Scalar,size_t> > bulkErrorDistribution;
      auto const& is = extensionSpace.gridManager().grid().leafIndexSet();
      CellIterator cend = extensionVariableSetDescription.gridView.template end<0>();

      while(bulkErrorSquared < bulkCriterionTolerance)
      {
        bulkErrorDistribution.push_back(errorDistribution[bulkErrorDistribution.size()]);
        bulkErrorSquared += bulkErrorDistribution.back().first;
      }

      if(verbose) std::cout << "ERROR ESTIMATOR: number of candidates for refinement: " << bulkErrorDistribution.size() << std::endl;

      //      size_t lastIndex = bulkErrorDistribution.size()-1;
      //      size_t firstIndex = 0;

      /*while(firstIndex+3 < lastIndex)
      {
        if(bulkErrorDistribution.size() > 4)
        {
          Scalar firstContribution = bulkErrorDistribution[firstIndex].first;
          firstContribution *= 15.0/256.0;

          Scalar lastContributions = bulkErrorDistribution[lastIndex--].first;
          lastContributions += bulkErrorDistribution[lastIndex--].first;
          lastContributions += bulkErrorDistribution[lastIndex].first;
          lastContributions *= 15.0/16.0;

          if(lastContributions < firstContribution)
          {
            ++lastIndexForDoubleRefinement;
            ++firstIndex;
            --lastIndex;
            bulkErrorDistribution.erase(bulkErrorDistribution.end()-3, bulkErrorDistribution.end());
          }
          else break;
        }
      }*/

      // Refine mesh.
      for (CellIterator ci=extensionVariableSetDescription.gridView.template begin<0>(); ci!=cend; ++ci) // iterate over cells
      {
        for(int i=0; i<bulkErrorDistribution.size(); ++i) // iterate over chosen part of the error distribution
          if(is.index(*ci) == bulkErrorDistribution[i].second)
          {
            //            if(i <= lastIndexForDoubleRefinement) {
            //              extensionSpace.gridManager().mark(2,*ci);
            //            }
            //            else{
            extensionSpace.gridManager().mark(1,*ci);
            //            }
          }
      }

      extensionSpace.gridManager().adaptAtOnce();
    }

    double estimatedAbsoluteError() const final
        {
      return sqrt(fabs(totalErrorSquared));
        }

    size_t gridSize() const final
        {
      return extensionSpace.gridManager().grid().size(0);
        }

  private:
    Functional& f;
    //    LF& lf;
    //    LVars& lvars;
    NormFunctional& normFunctional;
    ExtensionVariableSetDescription& extensionVariableSetDescription;
    ExtensionSpace& extensionSpace;
    Assembler assembler;
    Scalar squaredFraction;
    Scalar totalErrorSquared;
    std::vector<std::pair<double,int> > errorDistribution;
    bool verbose;
    // AdjustRHS<CoefficientVector> adjustRHS;
  };
}

#endif


// Software License for MTL
//
// Copyright (c) 2007 The Trustees of Indiana University.
//               2008 Dresden University of Technology and the Trustees of Indiana University.
//               2010 SimuNova UG (haftungsbeschränkt), www.simunova.com.
// All rights reserved.
// Authors: Peter Gottschling and Andrew Lumsdaine
//
// This file is part of the Matrix Template Library
//
// See also license.mtl.txt in the distribution.

#ifndef MTL_MAT_VEC_MULT_INCLUDE
#define MTL_MAT_VEC_MULT_INCLUDE

#include <cassert>
// #include <iostream>
#include <boost/mpl/bool.hpp>

#include <boost/numeric/mtl/config.hpp>
#include <boost/numeric/mtl/mtl_fwd.hpp>
#include <boost/numeric/mtl/utility/category.hpp>
#include <boost/numeric/mtl/utility/property_map.hpp>
#include <boost/numeric/mtl/utility/range_generator.hpp>
#include <boost/numeric/mtl/utility/tag.hpp>
#include <boost/numeric/mtl/utility/is_static.hpp>
#include <boost/numeric/mtl/utility/tag.hpp>
#include <boost/numeric/mtl/utility/enable_if.hpp>
#include <boost/numeric/mtl/utility/multi_tmp.hpp>
#include <boost/numeric/mtl/utility/static_assert.hpp>
#include <boost/numeric/mtl/vector/dense_vector.hpp>
#include <boost/numeric/mtl/utility/omp_size_type.hpp>
#include <boost/numeric/mtl/operation/set_to_zero.hpp>
#include <boost/numeric/mtl/operation/update.hpp>
#include <boost/numeric/linear_algebra/identity.hpp>
#include <boost/numeric/meta_math/loop.hpp>
#include <boost/numeric/mtl/interface/vpt.hpp>

namespace mtl { namespace matrix {

namespace impl {

    template <std::size_t Index0, std::size_t Max0, std::size_t Index1, std::size_t Max1, typename Assign>
    struct fully_unroll_mat_cvec_mult
      : public meta_math::loop2<Index0, Max0, Index1, Max1>
    {
	typedef meta_math::loop2<Index0, Max0, Index1, Max1>                              base;
	typedef fully_unroll_mat_cvec_mult<base::next_index0, Max0, base::next_index1, Max1, Assign>  next_t;

	template <typename Matrix, typename VectorIn, typename VectorOut>
	static inline void apply(const Matrix& A, const VectorIn& v, VectorOut& w)
	{
	    Assign::update(w[base::index0], A[base::index0][base::index1] * v[base::index1]);
	    next_t::apply(A, v, w);
	}
    };

    // need specialization here for not going back to column 0 but column 1
    template <std::size_t Index0, std::size_t Max0, std::size_t Max1, typename Assign>
    struct fully_unroll_mat_cvec_mult<Index0, Max0, Max1, Max1, Assign>
      : public meta_math::loop2<Index0, Max0, Max1, Max1>
    {
	typedef meta_math::loop2<Index0, Max0, Max1, Max1>                              base;
	typedef fully_unroll_mat_cvec_mult<base::next_index0, Max0, 2, Max1, Assign>  next_t;

	template <typename Matrix, typename VectorIn, typename VectorOut>
	static inline void apply(const Matrix& A, const VectorIn& v, VectorOut& w)
	{
	    Assign::update(w[base::index0], A[base::index0][base::index1] * v[base::index1]);
	    next_t::apply(A, v, w);
	}
    };

    template <std::size_t Max0, std::size_t Max1, typename Assign>
    struct fully_unroll_mat_cvec_mult<Max0, Max0, Max1, Max1, Assign>
      : public meta_math::loop2<Max0, Max0, Max1, Max1>
    {
	typedef meta_math::loop2<Max0, Max0, Max1, Max1>                              base;

	template <typename Matrix, typename VectorIn, typename VectorOut>
	static inline void apply(const Matrix& A, const VectorIn& v, VectorOut& w)
	{
	    Assign::update(w[base::index0], A[base::index0][base::index1] * v[base::index1]);
	}
    };

    struct noop
    {
	template <typename Matrix, typename VectorIn, typename VectorOut>
	static inline void apply(const Matrix&, const VectorIn&, VectorOut&) {}
    };
} // impl

// Dense matrix vector multiplication with run-time matrix size
template <typename Matrix, typename VectorIn, typename VectorOut, typename Assign>
inline void dense_mat_cvec_mult(const Matrix& A, const VectorIn& v, VectorOut& w, Assign, boost::mpl::true_)
{
    vampir_trace<3017> tracer;
    typedef typename static_num_rows<Matrix>::type size_type;
    static const size_type rows_a= static_num_rows<Matrix>::value, cols_a= static_num_cols<Matrix>::value;

    assert(rows_a > 0 && cols_a > 0);
    // w= A[all][0] * v[0];  N.B.: 1D is unrolled by the compiler faster (at least on gcc)
    for (size_type i= 0; i < rows_a; i++)
	Assign::first_update(w[i], A[i][0] * v[0]);

    // corresponds to w+= A[all][1:] * v[1:]; if necessary
    typedef impl::fully_unroll_mat_cvec_mult<1, rows_a, 2, cols_a, Assign>  f2;
    typedef typename boost::mpl::if_c<(cols_a > 1), f2, impl::noop>::type   f3;
    f3::apply(A, v, w);
}

template <unsigned Index, unsigned Size>
struct init_ptrs
{
    template <typename Matrix, typename Ptrs>
    inline static void apply(const Matrix& A, Ptrs& ptrs)
    {
	ptrs.value= &A[Index][0];
	init_ptrs<Index+1, Size>::apply(A, ptrs.sub);
    }
};

template <unsigned Size>
struct init_ptrs<Size, Size>
{
    template <typename Matrix, typename Ptrs>
    inline static void apply(const Matrix&, Ptrs&) {}
};

template <unsigned Index, unsigned Size>
struct square_cvec_mult_rows
{
    template <typename Tmps, typename Ptrs, typename ValueIn>
    inline static void compute(Tmps& tmps, Ptrs& ptrs, ValueIn vi)
    {
	tmps.value+= *ptrs.value++ * vi;
	square_cvec_mult_rows<Index+1, Size>::compute(tmps.sub, ptrs.sub, vi);
    }

    template <typename As, typename VectorOut, typename Tmps>
    inline static void update(As, VectorOut& w, const Tmps& tmps)
    {
	As::first_update(w[Index], tmps.value);
	square_cvec_mult_rows<Index+1, Size>::update(As(), w, tmps.sub);
    }
};

template <unsigned Size>
struct square_cvec_mult_rows<Size, Size>
{
    template <typename Tmps, typename Ptrs, typename ValueIn>
    inline static void compute(Tmps&, Ptrs&, ValueIn) {}

    template <typename As, typename VectorOut, typename Tmps>
    inline static void update(As, VectorOut&, const Tmps&) {}
};

template <unsigned Index, unsigned Size>
struct square_cvec_mult_cols
{
    template <typename Tmps, typename Ptrs, typename VPtr>
    inline static void compute(Tmps& tmps, Ptrs& ptrs, VPtr vp)
    {
	square_cvec_mult_rows<0, Size>::compute(tmps, ptrs, *vp);
	square_cvec_mult_cols<Index+1, Size>::compute(tmps, ptrs, ++vp);
    }
};

template <unsigned Size>
struct square_cvec_mult_cols<Size, Size>
{
    template <typename Tmps, typename Ptrs, typename VPtr>
    inline static void compute(Tmps&, Ptrs&, VPtr) {}
};

// Dense matrix vector multiplication with run-time matrix size
template <unsigned Size, typename MValue, typename MPara, typename ValueIn, typename ParaIn,
	  typename VectorOut, typename Assign>
inline void square_cvec_mult(const dense2D<MValue, MPara>& A, const mtl::vector::dense_vector<ValueIn, ParaIn>& v, VectorOut& w, Assign)
{
    // vampir_trace<3067> tracer;
    MTL_STATIC_ASSERT((mtl::traits::is_row_major<MPara>::value), "Only row-major matrices supported in this function.");

    typedef typename Collection<VectorOut>::value_type value_type;
    multi_tmp<Size, value_type> tmps(math::zero(w[0]));

    multi_tmp<Size, const MValue*>    ptrs;
    init_ptrs<0, Size>::apply(A, ptrs);

    const ValueIn* vp= &v[0];

    square_cvec_mult_cols<0, Size>::compute(tmps, ptrs, vp); // outer loop over columns
    square_cvec_mult_rows<0, Size>::update(Assign(), w, tmps);         // update rows
}

// Dense matrix vector multiplication with run-time matrix size
template <typename MValue, typename MPara, typename ValueIn, typename ParaIn, typename VectorOut, typename Assign>
typename boost::enable_if<mtl::traits::is_row_major<MPara> >::type
inline dense_mat_cvec_mult(const dense2D<MValue, MPara>& A, const mtl::vector::dense_vector<ValueIn, ParaIn>& v, VectorOut& w, Assign, boost::mpl::false_)
{
    // vampir_trace<3066> tracer;

    using math::zero;
    if (mtl::vector::size(w) == 0) return;

    typedef typename Collection<VectorOut>::value_type value_type;
    // typedef ValueIn                                    value_in_type;
    typedef typename MPara::size_type                  size_type;

    const size_type  nr= num_rows(A), nc= num_cols(A);
    if (nr == nc && nr <= 8)
	switch (nr) {
	  case 1: Assign::first_update(w[0], A[0][0] * v[0]); return;
	  case 2: square_cvec_mult<2>(A, v, w, Assign()); return;
	  case 3: square_cvec_mult<3>(A, v, w, Assign()); return;
	  case 4: square_cvec_mult<4>(A, v, w, Assign()); return;
	  case 5: square_cvec_mult<5>(A, v, w, Assign()); return;
	  case 6: square_cvec_mult<6>(A, v, w, Assign()); return;
	  case 7: square_cvec_mult<7>(A, v, w, Assign()); return;
	  case 8: square_cvec_mult<8>(A, v, w, Assign()); return;
	}

    const size_type  nrb= nr / 4 * 4;
    const value_type z(math::zero(w[0]));

    for (size_type i= 0; i < nrb; i+= 4) {
	value_type      tmp0(z), tmp1(z), tmp2(z), tmp3(z);
	const MValue *p0= &A[i][0], *pe= p0 + nc, *p1= &A[i+1][0], *p2= &A[i+2][0], *p3= &A[i+3][0];
	const ValueIn* vp= &v[0];
	for (; p0 != pe; ) {
	    const ValueIn vj= *vp++;
	    tmp0+= *p0++ * vj;
	    tmp1+= *p1++ * vj;
	    tmp2+= *p2++ * vj;
	    tmp3+= *p3++ * vj;
	}
	Assign::first_update(w[i], tmp0);
	Assign::first_update(w[i+1], tmp1);
	Assign::first_update(w[i+2], tmp2);
	Assign::first_update(w[i+3], tmp3);
    }

    for (size_type i= nrb; i < nr; i++) {
	value_type tmp= z;
	const ValueIn* vp= &v[0];
	for (const MValue *p0= &A[i][0], *pe= p0 + nc; p0 != pe; )
	    tmp+= *p0++ * *vp++;
	Assign::first_update(w[i], tmp);
    }
}

// Dense matrix vector multiplication with run-time matrix size
template <typename Matrix, typename VectorIn, typename VectorOut, typename Assign>
inline void dense_mat_cvec_mult(const Matrix& A, const VectorIn& v, VectorOut& w, Assign, boost::mpl::false_)
{
    vampir_trace<3018> tracer;
    // Naive implementation, will be moved to a functor and complemented with more efficient ones

    using math::zero;
    if (mtl::vector::size(w) == 0) return;
    // std::cout << "Bin in richtiger Funktion\n";

    // if (Assign::init_to_zero) set_to_zero(w); // replace update with first_update insteda

    typedef typename Collection<VectorOut>::value_type value_type;
    typedef typename Collection<VectorIn>::value_type  value_in_type;
    typedef typename Collection<Matrix>::size_type     size_type;

    const value_type z(math::zero(w[0]));
    const size_type nr= num_rows(A), nrb= nr / 4 * 4, nc= num_cols(A);

    for (size_type i= 0; i < nrb; i+= 4) {
	value_type      tmp0(z), tmp1(z), tmp2(z), tmp3(z);
	for (size_type j= 0; j < nc; j++) {
	    const value_in_type vj= v[j];
	    tmp0+= A[i][j] * vj;
	    tmp1+= A[i+1][j] * vj;
	    tmp2+= A[i+2][j] * vj;
	    tmp3+= A[i+3][j] * vj;
	}
	Assign::first_update(w[i], tmp0);
	Assign::first_update(w[i+1], tmp1);
	Assign::first_update(w[i+2], tmp2);
	Assign::first_update(w[i+3], tmp3);
    }

    for (size_type i= nrb; i < nr; i++) {
	value_type tmp= z;
	for (size_type j= 0; j < nc; j++)
	    tmp+= A[i][j] * v[j];
	Assign::first_update(w[i], tmp);
    }
}

// Dense matrix vector multiplication
template <typename Matrix, typename VectorIn, typename VectorOut, typename Assign>
inline void mat_cvec_mult(const Matrix& A, const VectorIn& v, VectorOut& w, Assign, tag::flat<tag::dense>)
{
# ifdef MTL_NOT_UNROLL_FSIZE_MAT_VEC_MULT
    boost::mpl::false_        selector;
# else
	mtl::traits::is_static<Matrix> selector;
# endif
    dense_mat_cvec_mult(A, v, w, Assign(), selector);
}

// Element structure vector multiplication
template <typename Matrix, typename VectorIn, typename VectorOut, typename Assign>
inline void mat_cvec_mult(const Matrix& A, const VectorIn& v, VectorOut& w, Assign, tag::flat<tag::element_structure>)
{
    vampir_trace<3048> tracer;
    if (mtl::vector::size(w) == 0) return;

    typedef typename Collection<VectorOut>::value_type value_type;
    typedef typename Collection<VectorIn>::value_type  value_in_type;

    value_in_type varray[1024];
    value_type    warray[1024];

    if (Assign::init_to_zero) set_to_zero(w);
    for(int elmi= 0; elmi < A.m_total_elements; elmi++){
	const typename Matrix::element_type& elementi= A.m_elements[elmi];
	const typename Matrix::element_type::index_type& indices= elementi.get_indices();
	unsigned int n(size(indices));

	if (n <= 1024) {
	    VectorIn vtmp(n, varray);
	    for (unsigned int i= 0; i < n; i++)
		vtmp[i]= v[indices[i]];
	    VectorOut wtmp(n, warray);
	    wtmp= elementi.get_values() * vtmp;
	    for (unsigned int i= 0; i < n; i++)
		Assign::update(w[indices[i]], wtmp[i]);
	} else {
	    VectorIn vtmp(n);
	    for (unsigned int i= 0; i < n; i++)
		vtmp[i]= v[indices[i]];
	    VectorOut wtmp(elementi.get_values() * vtmp);
	    for (unsigned int i= 0; i < n; i++)
		Assign::update(w[indices[i]], wtmp[i]);
	}
    }
}

// Multi-vector vector multiplication (tag::multi_vector is derived from dense)
template <typename Matrix, typename VectorIn, typename VectorOut, typename Assign>
inline void mat_cvec_mult(const Matrix& A, const VectorIn& v, VectorOut& w, Assign, tag::flat<tag::multi_vector>)
{
    vampir_trace<3019> tracer;
    if (Assign::init_to_zero) set_to_zero(w);
    for (unsigned i= 0; i < num_cols(A); i++)
	Assign::update(w, A.vector(i) * v[i]);
}

// Transposed multi-vector vector multiplication (tag::transposed_multi_vector is derived from dense)
template <typename TransposedMatrix, typename VectorIn, typename VectorOut, typename Assign>
inline void mat_cvec_mult(const TransposedMatrix& A, const VectorIn& v, VectorOut& w, Assign, tag::flat<tag::transposed_multi_vector>)
{
    vampir_trace<3020> tracer;
    typename TransposedMatrix::const_ref_type B= A.ref; // Referred matrix

    if (Assign::init_to_zero) set_to_zero(w);
    for (unsigned i= 0; i < num_cols(B); i++)
	Assign::update(w[i], dot_real(B.vector(i), v));
}

// Hermitian multi-vector vector multiplication (tag::hermitian_multi_vector is derived from dense)
template <typename HermitianMatrix, typename VectorIn, typename VectorOut, typename Assign>
inline void mat_cvec_mult(const HermitianMatrix& A, const VectorIn& v, VectorOut& w, Assign, tag::flat<tag::hermitian_multi_vector>)
{
    vampir_trace<3021> tracer;
    typename HermitianMatrix::const_ref_type B= A.const_ref(); // Referred matrix

    if (Assign::init_to_zero) set_to_zero(w);
    for (unsigned i= 0; i < num_cols(B); i++)
	Assign::update(w[i], dot(B.vector(i), v));
}

// Sparse row-major matrix vector multiplication
template <typename Matrix, typename VectorIn, typename VectorOut, typename Assign>
inline void smat_cvec_mult(const Matrix& A, const VectorIn& v, VectorOut& w, Assign, tag::row_major)
{
    vampir_trace<3022> tracer;
    using namespace tag;
    using mtl::traits::range_generator;
    using math::zero;
    using mtl::vector::set_to_zero;

    typedef typename range_generator<row, Matrix>::type       a_cur_type;
    typedef typename range_generator<nz, a_cur_type>::type    a_icur_type;
    typename mtl::traits::col<Matrix>::type                   col_a(A);
    typename mtl::traits::const_value<Matrix>::type           value_a(A);

    if (Assign::init_to_zero) set_to_zero(w);

    typedef typename Collection<VectorOut>::value_type        value_type;
    a_cur_type ac= begin<row>(A), aend= end<row>(A);
    for (unsigned i= 0; ac != aend; ++ac, ++i) {
	value_type tmp= zero(w[i]);
	for (a_icur_type aic= begin<nz>(ac), aiend= end<nz>(ac); aic != aiend; ++aic)
	    tmp+= value_a(*aic) * v[col_a(*aic)];
	Assign::update(w[i], tmp);
    }
}

// Row-major compressed2D with very few entries (i.e. Very Sparse MATrix) times vector
template <typename MValue, typename MPara, typename VectorIn, typename VectorOut, typename Assign>
inline void vsmat_cvec_mult(const compressed2D<MValue, MPara>& A, const VectorIn& v, VectorOut& w, Assign, tag::row_major)
{
    vampir_trace<3064> tracer;
    using math::zero;

    typedef compressed2D<MValue, MPara>                       Matrix;
    typedef typename Collection<Matrix>::size_type            size_type;
    typedef typename Collection<VectorOut>::value_type        value_type;

    if (size(w) == 0) return;
    const value_type z(math::zero(w[0]));

    // std::cout << "very sparse: nnz = " << A.nnz() << ", num_rows = " << num_rows(A) << '\n';

    size_type nr= num_rows(A);
    for (size_type i1= 0, i2= std::min<size_type>(1024, nr); i1 < i2; i1= i2, i2= std::min<size_type>(i2 + 1024, nr)) {
	// std::cout << "range = " << i1 << " .. " << i2 << "\n";
	if (A.ref_major()[i1] < A.ref_major()[i2])
	    for (size_type i= i1; i < i2; ++i) {
		const size_type cj0= A.ref_major()[i], cj1= A.ref_major()[i+1];
		value_type      tmp0(z);
		for (size_type j0= cj0; j0 != cj1; ++j0)
		    tmp0+= A.data[j0] * v[A.ref_minor()[j0]];
		Assign::first_update(w[i], tmp0);
	    }
    }
}

#ifdef MTL_CRS_CVEC_MULT_TUNING
template <unsigned Index, unsigned BSize, typename SizeType>
struct crs_cvec_mult_block
{
    template <typename Matrix, typename VectorIn, typename CBlock, typename TBlock>
    void operator()(const Matrix& A, const VectorIn& v, const CBlock& cj, TBlock& tmp) const
    {
	for (SizeType j= cj.value; j != cj.sub.value; ++j) // cj is one index larger
	    tmp.value+= A.data[j] * v[A.ref_minor()[j]];
	sub(A, v, cj.sub, tmp.sub);
    }

    template <typename VectorOut, typename TBlock, typename Assign>
    void first_update(VectorOut& w, SizeType i, const TBlock& tmp, Assign as) const
    {
	Assign::first_update(w[i + Index], tmp.value);
	sub.first_update(w, i, tmp.sub, as);
    }

    crs_cvec_mult_block<Index+1, BSize, SizeType> sub;
};

template <unsigned BSize, typename SizeType>
struct crs_cvec_mult_block<BSize, BSize, SizeType>
{
    template <typename Matrix, typename VectorIn, typename CBlock, typename TBlock>
    void operator()(const Matrix& A, const VectorIn& v, const CBlock& cj, TBlock& tmp) const
    {
	for (SizeType j= cj.value; j != cj.sub.value; ++j)// cj is one index larger
	    tmp.value+= A.data[j] * v[A.ref_minor()[j]];
    }

    template <typename VectorOut, typename TBlock, typename Assign>
    void first_update(VectorOut& w, SizeType i, const TBlock& tmp, Assign) const
    {
	Assign::first_update(w[i + BSize], tmp.value);
    }
};

// Row-major compressed2D vector multiplication
template <unsigned BSize, typename MValue, typename MPara, typename VectorIn, typename VectorOut, typename Assign>
inline void smat_cvec_mult(const compressed2D<MValue, MPara>& A, const VectorIn& v, VectorOut& w, Assign as, tag::row_major)
{
    vampir_trace<3049> tracer;
    using math::zero;

    if (A.nnz() < num_rows(A)) {
	vsmat_cvec_mult(A, v, w, as, tag::row_major());
	return;
    }

    typedef compressed2D<MValue, MPara>                       Matrix;
    typedef typename Collection<VectorOut>::value_type        value_type;
    typedef typename mtl::traits::omp_size_type<typename Collection<Matrix>::size_type>::type size_type;

    if (size(w) == 0) return;
    const value_type z(math::zero(w[0]));

    size_type nr= num_rows(A), nrb= nr / BSize * BSize;

    #ifdef MTL_WITH_OPENMP
    #   pragma omp parallel
    #endif
    {
    	#ifdef MTL_WITH_OPENMP
	    vampir_trace<8004> tracer;
    	#   pragma omp for
    	#endif
	for (size_type i= 0; i < nrb; i+= BSize) {
	    multi_constant_from_array<0, BSize+1, size_type> cj(A.ref_major(), i);
	    multi_tmp<BSize, value_type>                     tmp(z);
	    crs_cvec_mult_block<0, BSize-1, size_type>       block;

	    block(A, v, cj, tmp);
	    block.first_update(w, i, tmp, as);
	}
    }

    for (size_type i= nrb; i < nr; ++i) {
	const size_type cj0= A.ref_major()[i], cj1= A.ref_major()[i+1];
	value_type      tmp0(z);
	for (size_type j0= cj0; j0 != cj1; ++j0)
	    tmp0+= A.data[j0] * v[A.ref_minor()[j0]];
	Assign::first_update(w[i], tmp0);
    }
}

template <typename MValue, typename MPara, typename VectorIn, typename VectorOut, typename Assign>
typename mtl::traits::enable_if_scalar<typename Collection<VectorOut>::value_type>::type
inline smat_cvec_mult(const compressed2D<MValue, MPara>& A, const VectorIn& v, VectorOut& w, Assign, tag::row_major)
{
    smat_cvec_mult<crs_cvec_mult_block_size>(A, v, w, Assign(), tag::row_major());
}
#endif

#if !defined(MTL_CRS_CVEC_MULT_NO_ACCEL) && !defined(MTL_CRS_CVEC_MULT_TUNING)

template <typename MValue, typename MPara, typename VectorIn, typename VectorOut, typename Assign>
typename mtl::traits::enable_if_scalar<typename Collection<VectorOut>::value_type>::type
inline adapt_crs_cvec_mult(const compressed2D<MValue, MPara>& A, const VectorIn& v, VectorOut& w, Assign)
{
    vampir_trace<3065> tracer;
    using math::zero;
    assert(!Assign::init_to_zero);

    typedef compressed2D<MValue, MPara>                       Matrix;
    typedef typename Collection<Matrix>::size_type            size_type;
    typedef typename Collection<VectorOut>::value_type        value_type;

    const value_type z(math::zero(w[0]));
    size_type nr= num_rows(A), nrb= nr / 4 * 4, nrb2= nr / 64 * 64;

    for (size_type i1= 0; i1 < nrb2; i1+= 64)
	if (A.ref_major()[i1] != A.ref_major()[i1 + 64])
	    for (size_type i2= i1, i2e= i1+64; i2 < i2e; i2+= 16)
		if (A.ref_major()[i2] != A.ref_major()[i2 + 16])
		    for (size_type i= i2, i3e= i2+16; i < i3e; i+= 4)
			if (A.ref_major()[i] != A.ref_major()[i + 4]) {
			    const size_type cj0= A.ref_major()[i], cj1= A.ref_major()[i+1], cj2= A.ref_major()[i+2],
				cj3= A.ref_major()[i+3], cj4= A.ref_major()[i+4];
			    value_type      tmp0(z), tmp1(z), tmp2(z), tmp3(z);
			    for (size_type j0= cj0; j0 != cj1; ++j0)
				tmp0+= A.data[j0] * v[A.ref_minor()[j0]];
			    for (size_type j1= cj1; j1 != cj2; ++j1)
				tmp1+= A.data[j1] * v[A.ref_minor()[j1]];
			    for (size_type j2= cj2; j2 != cj3; ++j2)
				tmp2+= A.data[j2] * v[A.ref_minor()[j2]];
			    for (size_type j3= cj3; j3 != cj4; ++j3)
				tmp3+= A.data[j3] * v[A.ref_minor()[j3]];

			    Assign::first_update(w[i], tmp0);
			    Assign::first_update(w[i+1], tmp1);
			    Assign::first_update(w[i+2], tmp2);
			    Assign::first_update(w[i+3], tmp3);
			}

    for (size_type i= nrb2; i < nrb; i+= 4)
	if (A.ref_major()[i] != A.ref_major()[i + 4])  {
	    const size_type cj0= A.ref_major()[i], cj1= A.ref_major()[i+1], cj2= A.ref_major()[i+2],
		cj3= A.ref_major()[i+3], cj4= A.ref_major()[i+4];
	    value_type      tmp0(z), tmp1(z), tmp2(z), tmp3(z);
	    for (size_type j0= cj0; j0 != cj1; ++j0)
		tmp0+= A.data[j0] * v[A.ref_minor()[j0]];
	    for (size_type j1= cj1; j1 != cj2; ++j1)
		tmp1+= A.data[j1] * v[A.ref_minor()[j1]];
	    for (size_type j2= cj2; j2 != cj3; ++j2)
		tmp2+= A.data[j2] * v[A.ref_minor()[j2]];
	    for (size_type j3= cj3; j3 != cj4; ++j3)
		tmp3+= A.data[j3] * v[A.ref_minor()[j3]];

	    Assign::first_update(w[i], tmp0);
	    Assign::first_update(w[i+1], tmp1);
	    Assign::first_update(w[i+2], tmp2);
	    Assign::first_update(w[i+3], tmp3);
	}

    for (size_type i= nrb; i < nr; ++i) {
	const size_type cj0= A.ref_major()[i], cj1= A.ref_major()[i+1];
	value_type      tmp0(z);
	for (size_type j0= cj0; j0 != cj1; ++j0)
	    tmp0+= A.data[j0] * v[A.ref_minor()[j0]];
	Assign::first_update(w[i], tmp0);
    }
}

// Row-major compressed2D vector multiplication
template <typename MValue, typename MPara, typename VectorIn, typename VectorOut, typename Assign>
typename mtl::traits::enable_if_scalar<typename Collection<VectorOut>::value_type>::type
inline smat_cvec_mult(const compressed2D<MValue, MPara>& A, const VectorIn& v, VectorOut& w, Assign as, tag::row_major)
{
    vampir_trace<3049> tracer;
    // vampir_trace<5056> tttracer;
    using math::zero;

    if (A.nnz() < num_rows(A) && !as.init_to_zero) {
	vsmat_cvec_mult(A, v, w, as, tag::row_major());
	return;
    }

    typedef compressed2D<MValue, MPara>                       Matrix;
    typedef typename Collection<VectorOut>::value_type        value_type;
    typedef typename mtl::traits::omp_size_type<typename Collection<Matrix>::size_type>::type size_type;

    if (size(w) == 0) return;
    const value_type z(math::zero(w[0]));

    size_type nr= num_rows(A), nrb= nr / 4 * 4;
    if (nr > 10) {
	size_type nh= nr / 2, nq= nr / 4, nt= nr - nq;
	if (!as.init_to_zero &&
	    (A.ref_major()[1] == A.ref_major()[0]
	     || A.ref_major()[nq] == A.ref_major()[nq+1]
	     || A.ref_major()[nh] == A.ref_major()[nh+1]
	     || A.ref_major()[nt] == A.ref_major()[nt+1]
	     || A.ref_major()[nr-1] == A.ref_major()[nr])) {
	    adapt_crs_cvec_mult(A, v, w, as);
	    return;
	}
    }

    #ifdef MTL_WITH_OPENMP
    #   pragma omp parallel
    #endif
    {
    	#ifdef MTL_WITH_OPENMP
	    vampir_trace<8004> tracer;
    	#   pragma omp for
    	#endif
	    for (size_type i= 0; i < nrb; i+= 4) {
		const size_type cj0= A.ref_major()[i], cj1= A.ref_major()[i+1], cj2= A.ref_major()[i+2],
		    cj3= A.ref_major()[i+3], cj4= A.ref_major()[i+4];
		value_type      tmp0(z), tmp1(z), tmp2(z), tmp3(z);
		for (size_type j0= cj0; j0 != cj1; ++j0)
		    tmp0+= A.data[j0] * v[A.ref_minor()[j0]];
		for (size_type j1= cj1; j1 != cj2; ++j1)
		    tmp1+= A.data[j1] * v[A.ref_minor()[j1]];
		for (size_type j2= cj2; j2 != cj3; ++j2)
		    tmp2+= A.data[j2] * v[A.ref_minor()[j2]];
		for (size_type j3= cj3; j3 != cj4; ++j3)
		    tmp3+= A.data[j3] * v[A.ref_minor()[j3]];

		Assign::first_update(w[i], tmp0);
		Assign::first_update(w[i+1], tmp1);
		Assign::first_update(w[i+2], tmp2);
		Assign::first_update(w[i+3], tmp3);
	    }
    }

    for (size_type i= nrb; i < nr; ++i) {
	const size_type cj0= A.ref_major()[i], cj1= A.ref_major()[i+1];
	value_type      tmp0(z);
	for (size_type j0= cj0; j0 != cj1; ++j0)
	    tmp0+= A.data[j0] * v[A.ref_minor()[j0]];
	Assign::first_update(w[i], tmp0);
    }
}
#endif

// Row-major ell_matrix vector multiplication
template <typename MValue, typename MPara, typename VectorIn, typename VectorOut, typename Assign>
typename mtl::traits::enable_if_scalar<typename Collection<VectorOut>::value_type>::type
inline smat_cvec_mult(const ell_matrix<MValue, MPara>& A, const VectorIn& v, VectorOut& w, Assign, tag::row_major)
{
    typedef typename MPara::size_type size_type;

    const size_type stride= A.stride(), slots= A.slots();
    for (size_type r= 0; r < A.dim1(); ++r) {
	MValue s(0);
	for (size_type k= r, i= 0; i < slots; ++i, k+= stride)
	    s+= A.ref_data()[k] * v[A.ref_minor()[k]];
	Assign::first_update(w[r], s);
    }
 }

// Row-major sparse_banded vector multiplication
template <typename MValue, typename MPara, typename VectorIn, typename VectorOut, typename Assign>
typename mtl::traits::enable_if_scalar<typename Collection<VectorOut>::value_type>::type
inline smat_cvec_mult(const sparse_banded<MValue, MPara>& A, const VectorIn& v, VectorOut& w, Assign, tag::row_major)
{
    vampir_trace<3069> tracer;
    typedef sparse_banded<MValue, MPara>                      Matrix;
    typedef typename Collection<VectorOut>::value_type        value_type;
    typedef typename Matrix::band_size_type                   band_size_type;
    typedef typename MPara::size_type                         size_type;
    typedef mtl::vector::dense_vector<band_size_type, vector::parameters<> > vector_type;

    if (size(w) == 0) return;
    const value_type z(math::zero(w[0]));

    size_type nr= num_rows(A), nc= num_cols(A), nb= A.ref_bands().size();
    if (nb == size_type(0) && Assign::init_to_zero) {
	set_to_zero(w);
	return;
    }

    vector_type bands(A.ref_bands()), begin_rows(max(0, -bands)), end_rows(min(nr, nc - bands));
    assert(end_rows[nb-1] > 0);

    // std::cout << "bands = " << bands << ", begin_rows = " << begin_rows << ", end_rows = " << end_rows << "\n";
    size_type begin_pos= 0, end_pos= nb - 1;

    // find lowest diagonal in row 0
    while (begin_pos < nb && begin_rows[begin_pos] > 0) begin_pos++;
    // if at the end, the first rows are empty
    if (begin_pos == nb && Assign::init_to_zero) {
	w[irange(begin_rows[--begin_pos])]= z;
	// std::cout << "w[0.." << begin_rows[begin_pos] << "] <- 0\n";
    }

    band_size_type from= begin_rows[begin_pos];
    // find first entry with same value
    while (begin_pos > 0 && begin_rows[begin_pos - 1] == from) {
	assert(from == 0); // should only happen when multiple bands start in row 0
	begin_pos--;
    }
    for (bool active= true; active; ) {
	// search backwards for the next-largest entry
	band_size_type to= begin_pos > 0 && begin_rows[begin_pos - 1] <= end_rows[end_pos] ? begin_rows[begin_pos - 1] : end_rows[end_pos];

	// std::cout << "rows " << from << ".." << to << ": with bands ";
	// for (size_type i= begin_pos; i <= end_pos; i++)
	//     std::cout << bands[i] << (i < end_pos ? ", " : "\n");

	const MValue* Aps= A.ref_data() + (from * nb + begin_pos);

	const band_size_type blocked_to= ((to - from) & -4) + from;
	assert((blocked_to - from) % 4 == 0 && blocked_to >= band_size_type(from) && blocked_to <= band_size_type(to));
	for (band_size_type r= from; r < blocked_to; r+= 4) {
	    value_type     tmp0(z), tmp1(z), tmp2(z), tmp3(z);
	    const MValue   *Ap0= Aps, *Ap1= Aps + nb, *Ap2= Ap1 + nb, *Ap3= Ap2 + nb;
	    for (size_type b= begin_pos; b <= end_pos; ++b, ++Ap0, ++Ap1, ++Ap2, ++Ap3) {
		tmp0+= *Ap0 * v[r + bands[b]];
		tmp1+= *Ap1 * v[r + bands[b] + 1];
		tmp2+= *Ap2 * v[r + bands[b] + 2];
		tmp3+= *Ap3 * v[r + bands[b] + 3];
	    }
	    Assign::first_update(w[r], tmp0);
	    Assign::first_update(w[r+1], tmp1);
	    Assign::first_update(w[r+2], tmp2);
	    Assign::first_update(w[r+3], tmp3);
	    Aps+= 4 * nb;
	}

	for (band_size_type r= blocked_to; r < band_size_type(to); r++) {
	    value_type     tmp(z);
	    const MValue*  Ap= Aps;
	    for (size_type b= begin_pos; b <= end_pos; ++b, ++Ap)
		tmp+= *Ap * v[r + bands[b]];
	    Assign::first_update(w[r], tmp);
	    Aps+= nb;
	}

	if (begin_pos > 0) {
	    if (begin_rows[begin_pos-1] == to)
		begin_pos--;
	    if (end_rows[end_pos] == to)
		end_pos--;
	} else { // begin == 0 -> decrement end_pos or finish
	    if (end_rows[0] == to) {
		active= false;
		assert(end_rows[0] = end_rows[end_pos]);
	    } else {
		assert(end_pos > 0);
		end_pos--;
	    }
	}
	assert(begin_pos <= end_pos);
	from= to;
    }

    if (size_type(end_rows[0]) < nr  && Assign::init_to_zero) {
	w[irange(end_rows[0], nr)]= z;
	// std::cout << "w[" << end_rows[0] << ".." << nr << "] <- 0\n";
    }
}

// Sparse column-major matrix vector multiplication
template <typename Matrix, typename VectorIn, typename VectorOut, typename Assign>
inline void smat_cvec_mult(const Matrix& A, const VectorIn& v, VectorOut& w, Assign, tag::col_major)
{
    vampir_trace<3023> tracer;
    using namespace tag; namespace traits = mtl::traits;
    using traits::range_generator;
    using mtl::vector::set_to_zero;
    typedef typename range_generator<col, Matrix>::type       a_cur_type;
    typedef typename range_generator<nz, a_cur_type>::type    a_icur_type;

    typename traits::row<Matrix>::type                        row_a(A);
    typename traits::const_value<Matrix>::type                value_a(A);

    if (Assign::init_to_zero) set_to_zero(w);

    unsigned rv= 0; // traverse all rows of v
    for (a_cur_type ac= begin<col>(A), aend= end<col>(A); ac != aend; ++ac, ++rv) {
	typename Collection<VectorIn>::value_type    vv= v[rv];
	for (a_icur_type aic= begin<nz>(ac), aiend= end<nz>(ac); aic != aiend; ++aic)
	    Assign::update(w[row_a(*aic)], value_a(*aic) * vv);
    }
}

// Sparse matrix vector multiplication
template <typename Matrix, typename VectorIn, typename VectorOut, typename Assign>
inline void mat_cvec_mult(const Matrix& A, const VectorIn& v, VectorOut& w, Assign, tag::flat<tag::sparse>)
{
    smat_cvec_mult(A, v, w, Assign(), typename OrientedCollection<Matrix>::orientation());
}

}} // namespace mtl::matrix

#endif // MTL_MAT_VEC_MULT_INCLUDE


// Copyright 2018 The Simons Foundation, 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.

#ifndef NETKET_METROPOLISFLIPT_HPP
#define NETKET_METROPOLISFLIPT_HPP

#include <mpi.h>
#include <Eigen/Dense>
#include <iostream>
#include "Utils/parallel_utils.hpp"
#include "Utils/random_utils.hpp"
#include "abstract_sampler.hpp"

namespace netket {

// Metropolis sampling generating local changes
// Parallel tempering is also used
class MetropolisLocalPt : public AbstractSampler {
  AbstractMachine& psi_;

  const AbstractHilbert& hilbert_;

  // number of visible units
  const int nv_;

  // states of visible units
  // for each sampled temperature
  std::vector<Eigen::VectorXd> v_;

  Eigen::VectorXd accept_;
  Eigen::VectorXd moves_;

  int mynode_;
  int totalnodes_;

  // clusters to do updates
  std::vector<std::vector<int>> clusters_;

  // Look-up tables
  std::vector<typename AbstractMachine::LookupType> lt_;

  int nrep_;

  std::vector<double> beta_;

  int nstates_;
  std::vector<double> localstates_;

 public:
  // Constructor with one replica by default
  explicit MetropolisLocalPt(AbstractMachine& psi, int nreplicas = 1)
      : psi_(psi),
        hilbert_(psi.GetHilbert()),
        nv_(hilbert_.Size()),
        nrep_(nreplicas) {
    Init();
  }

  void Init() {
    MPI_Comm_size(MPI_COMM_WORLD, &totalnodes_);
    MPI_Comm_rank(MPI_COMM_WORLD, &mynode_);

    nstates_ = hilbert_.LocalSize();
    localstates_ = hilbert_.LocalStates();

    SetNreplicas(nrep_);

    InfoMessage() << "Metropolis sampler with parallel tempering is ready "
                  << std::endl;
    InfoMessage() << "Nreplicas is equal to " << nrep_ << std::endl;
  }

  void SetNreplicas(int nrep) {
    nrep_ = nrep;
    v_.resize(nrep_);
    for (int i = 0; i < nrep_; i++) {
      v_[i].resize(nv_);
    }

    for (int i = 0; i < nrep_; i++) {
      beta_.push_back(1. - double(i) / double(nrep_));
    }

    lt_.resize(nrep_);

    accept_.resize(2 * nrep_);
    moves_.resize(2 * nrep_);

    Reset(true);
  }

  void Reset(bool initrandom = false) override {
    if (initrandom) {
      for (int i = 0; i < nrep_; i++) {
        hilbert_.RandomVals(v_[i], this->GetRandomEngine());
      }
    }

    for (int i = 0; i < nrep_; i++) {
      psi_.InitLookup(v_[i], lt_[i]);
    }

    accept_ = Eigen::VectorXd::Zero(2 * nrep_);
    moves_ = Eigen::VectorXd::Zero(2 * nrep_);
  }

  // Exchange sweep at given temperature
  void LocalSweep(int rep) {
    std::vector<int> tochange(1);
    std::vector<double> newconf(1);

    std::uniform_real_distribution<double> distu;
    std::uniform_int_distribution<int> distrs(0, nv_ - 1);
    std::uniform_int_distribution<int> diststate(0, nstates_ - 1);

    for (int i = 0; i < nv_; i++) {
      // picking a random site to be changed
      int si = distrs(this->GetRandomEngine());
      assert(si < nv_);
      tochange[0] = si;

      // picking a random state
      int newstate = diststate(this->GetRandomEngine());
      newconf[0] = localstates_[newstate];

      // make sure that the new state is not equal to the current one
      while (std::abs(newconf[0] - v_[rep](si)) <
             std::numeric_limits<double>::epsilon()) {
        newstate = diststate(this->GetRandomEngine());
        newconf[0] = localstates_[newstate];
      }

      const auto lvd = psi_.LogValDiff(v_[rep], tochange, newconf, lt_[rep]);
      double ratio = std::norm(std::exp(beta_[rep] * lvd));

#ifndef NDEBUG
      const auto psival1 = psi_.LogVal(v_[rep]);
      if (std::abs(
              std::exp(psi_.LogVal(v_[rep]) - psi_.LogVal(v_[rep], lt_[rep])) -
              1.) > 1.0e-8) {
        std::cerr << psi_.LogVal(v_[rep]) << "  and LogVal with Lt is "
                  << psi_.LogVal(v_[rep], lt_[rep]) << std::endl;
        std::abort();
      }
#endif
      // Metropolis acceptance test
      if (ratio > distu(this->GetRandomEngine())) {
        accept_(rep) += 1;

        psi_.UpdateLookup(v_[rep], tochange, newconf, lt_[rep]);
        hilbert_.UpdateConf(v_[rep], tochange, newconf);

#ifndef NDEBUG
        const auto psival2 = psi_.LogVal(v_[rep]);
        if (std::abs(std::exp(psival2 - psival1 - lvd) - 1.) > 1.0e-8) {
          std::cerr << psival2 - psival1 << " and logvaldiff is " << lvd
                    << std::endl;
          std::cerr << psival2 << " and LogVal with Lt is "
                    << psi_.LogVal(v_[rep], lt_[rep]) << std::endl;
          std::abort();
        }
#endif
      }
      moves_(rep) += 1;
    }
  }

  void Sweep() override {
    // First we do local sweeps
    for (int i = 0; i < nrep_; i++) {
      LocalSweep(i);
    }

    // Tempearture exchanges
    std::uniform_real_distribution<double> distribution(0, 1);

    for (int r = 1; r < nrep_; r += 2) {
      if (ExchangeProb(r, r - 1) > distribution(this->GetRandomEngine())) {
        Exchange(r, r - 1);
        accept_(nrep_ + r) += 1.;
        accept_(nrep_ + r - 1) += 1;
      }
      moves_(nrep_ + r) += 1.;
      moves_(nrep_ + r - 1) += 1;
    }

    for (int r = 2; r < nrep_; r += 2) {
      if (ExchangeProb(r, r - 1) > distribution(this->GetRandomEngine())) {
        Exchange(r, r - 1);
        accept_(nrep_ + r) += 1.;
        accept_(nrep_ + r - 1) += 1;
      }
      moves_(nrep_ + r) += 1.;
      moves_(nrep_ + r - 1) += 1;
    }
  }

  // computes the probability to exchange two replicas
  double ExchangeProb(int r1, int r2) {
    const double lf1 = 2 * std::real(psi_.LogVal(v_[r1], lt_[r1]));
    const double lf2 = 2 * std::real(psi_.LogVal(v_[r2], lt_[r2]));

    return std::exp((beta_[r1] - beta_[r2]) * (lf2 - lf1));
  }

  void Exchange(int r1, int r2) {
    std::swap(v_[r1], v_[r2]);
    std::swap(lt_[r1], lt_[r2]);
  }

  Eigen::VectorXd Visible() override { return v_[0]; }

  void SetVisible(const Eigen::VectorXd& v) override { v_[0] = v; }

  AbstractMachine& GetMachine() noexcept override { return psi_; }

  const AbstractHilbert& GetHilbert() const noexcept override {
    return hilbert_;
  }

  Eigen::VectorXd Acceptance() const override {
    Eigen::VectorXd acc = accept_;
    for (int i = 0; i < acc.size(); i++) {
      acc(i) /= moves_(i);
    }
    return acc;
  }
};

}  // namespace netket

#endif