/*!@file
 * @copyright This code is licensed under the 3-clause BSD license.
 *   Copyright ETH Zurich, Laboratory of Physical Chemistry, Reiher Group.
 *   See LICENSE.txt for details.
 */

#include "Molassembler/Stereopermutators/AtomStereopermutatorImpl.h"

#include "boost/range/join.hpp"
#include "Molassembler/Shapes/Properties.h"
#include "Molassembler/Shapes/PropertyCaching.h"
#include "Molassembler/Shapes/ContinuousMeasures.h"
#include <Eigen/Dense>
#include <Eigen/SparseCore>
#include "Molassembler/Stereopermutation/Manipulation.h"
#include "Molassembler/Temple/Adaptors/AllPairs.h"
#include "Molassembler/Temple/Adaptors/Iota.h"
#include "Molassembler/Temple/Adaptors/Transform.h"
#include "Molassembler/Temple/Adaptors/Zip.h"
#include "Molassembler/Temple/Functional.h"
#include "Molassembler/Temple/Functor.h"
#include "Molassembler/Temple/Optionals.h"
#include "Molassembler/Temple/OrderedPair.h"
#include "Molassembler/Temple/Random.h"
#include "Molassembler/Temple/TinySet.h"
#include "Molassembler/Temple/constexpr/Math.h"
#include "Molassembler/Temple/constexpr/Numeric.h"

#include "Molassembler/Cycles.h"
#include "Molassembler/Detail/BuildTypeSwitch.h"
#include "Molassembler/Detail/Cartesian.h"
#include "Molassembler/DistanceGeometry/SpatialModel.h"
#include "Molassembler/DistanceGeometry/ValueBounds.h"
#include "Molassembler/Graph/PrivateGraph.h"
#include "Molassembler/Log.h"
#include "Molassembler/Modeling/CommonTrig.h"
#include "Molassembler/Stereopermutators/ShapeVertexMaps.h"
#include "Molassembler/Stereopermutators/RankingMapping.h"

#include "Utils/Geometry/ElementInfo.h"

using namespace std::string_literals;
using namespace std::placeholders;

namespace Scine {
namespace Molassembler {

namespace {

Shapes::Shape pickTransition(
  const Shapes::Shape shape,
  const unsigned T,
  const boost::optional<Shapes::Vertex>& removedVertexOptional
) {
  boost::optional<Shapes::Properties::ShapeTransitionGroup> bestTransition;
  std::vector<Shapes::Shape> propositions;

  auto replaceOrAdd = [&](
    Shapes::Shape newName,
    const Shapes::Properties::ShapeTransitionGroup& transitionGroup
  ) {
    if(bestTransition){
      if(transitionGroup.angularDistortion < bestTransition->angularDistortion) {
        bestTransition = transitionGroup;
        propositions = {newName};
      } else if(transitionGroup.angularDistortion == bestTransition->angularDistortion) {
        propositions.push_back(newName);
      }
    }

    if(!bestTransition) {
      bestTransition = transitionGroup;
      propositions = {newName};
    }
  };

  // Populate the list of shape with candidates
  for(const Shapes::Shape propositionalShape : Shapes::allShapes) {
    if(Shapes::size(propositionalShape) != T) {
      continue;
    }

    if(auto transitionOptional = Shapes::getMapping(shape, propositionalShape, removedVertexOptional)) {
      if(transitionOptional->indexMappings.empty()) {
        continue;
      }

      if(
        Options::chiralStatePreservation == ChiralStatePreservation::EffortlessAndUnique
        && transitionOptional->indexMappings.size() == 1
        && transitionOptional->angularDistortion <= 0.2
      ) {
        replaceOrAdd(propositionalShape, *transitionOptional);
      }

      if(
        Options::chiralStatePreservation == ChiralStatePreservation::Unique
        && transitionOptional->indexMappings.size() == 1
      ) {
        replaceOrAdd(propositionalShape, *transitionOptional);
      }

      if(Options::chiralStatePreservation == ChiralStatePreservation::RandomFromMultipleBest) {
        replaceOrAdd(propositionalShape, *transitionOptional);
      }
    }
  }

  /* In case we have no propositions, add all of fitting size. We need to
   * propose something, even if state will not be propagated due to the
   * preservation criteria.
   */
  if(propositions.empty()) {
    return Shapes::Properties::mostSymmetric(T);
  }

  return Shapes::Properties::mostSymmetric(std::move(propositions));
}

std::pair<Shapes::Shape, std::vector<Shapes::Vertex>> classifyShape(const Eigen::Matrix<double, 3, Eigen::Dynamic>& sitePositions) {
  const unsigned S = sitePositions.cols() - 1;
  auto normalized = Shapes::Continuous::normalize(sitePositions);

  std::vector<Shapes::Shape> viableShapes;
  for(const Shapes::Shape shape : Shapes::allShapes) {
    if(Shapes::size(shape) == S) {
      viableShapes.push_back(shape);
    }
  }
  const unsigned shapesCount = viableShapes.size();
  std::vector<Shapes::Continuous::ShapeResult> shapeMeasureResults (shapesCount);
  std::vector<boost::optional<double>> randomCloudProbabilities (shapesCount);

#pragma omp parallel for
  for(unsigned i = 0; i < shapesCount; ++i) {
    const Shapes::Shape candidateShape = viableShapes[i];
    shapeMeasureResults[i] = Shapes::Continuous::shapeCentroidLast(normalized, candidateShape);
    // Shape classification for size 2 is better based on continuous shape measures themselves
    if(Shapes::size(candidateShape) > 2) {
      randomCloudProbabilities[i] = Shapes::Continuous::probabilityRandomCloud(shapeMeasureResults[i].measure, candidateShape);
    }
  }

  // Ensure centroids are mapped against one another
  assert(
    Temple::all_of(
      shapeMeasureResults,
      [](const auto& shapeResult) -> bool {
        assert(!shapeResult.mapping.empty());
        return shapeResult.mapping.back() == shapeResult.mapping.size() - 1;
      }
    )
  );

  // Prefer probabilities for comparison
  if(Temple::all_of(randomCloudProbabilities)) {
    const auto minElementIter = std::min_element(
      std::begin(randomCloudProbabilities),
      std::end(randomCloudProbabilities),
      [](const boost::optional<double>& a, const boost::optional<double>& b) -> bool {
        // We know all optionals are Somes from the previous all_of call
        return a.value() < b.value();
      }
    );
    const unsigned minimalShapeIndex = minElementIter - std::begin(randomCloudProbabilities);
    const Shapes::Shape minimalShape = viableShapes.at(minimalShapeIndex);
    return std::make_pair(
      minimalShape,
      std::move(shapeMeasureResults.at(minimalShapeIndex).mapping)
    );
  }

  // Fall back to minimal shape measure
  const auto minElementIter = std::min_element(
    std::begin(shapeMeasureResults),
    std::end(shapeMeasureResults),
    [](const auto& a, const auto& b) -> bool {
      return a.measure < b.measure;
    }
  );
  const unsigned minimalShapeIndex = minElementIter - std::begin(shapeMeasureResults);
  const Shapes::Shape minimalShape = viableShapes.at(minimalShapeIndex);
  return std::make_pair(
    minimalShape,
    std::move(minElementIter->mapping)
  );
}

} // namespace

Shapes::Shape AtomStereopermutator::Impl::up(const Shapes::Shape shape) {
  return pickTransition(shape, Shapes::size(shape) + 1, boost::none);
}

Shapes::Shape AtomStereopermutator::Impl::down(const Shapes::Shape shape, const Shapes::Vertex removedVertex) {
  return pickTransition(shape, Shapes::size(shape) - 1, removedVertex);
}

boost::optional<std::vector<Shapes::Vertex>> AtomStereopermutator::Impl::selectTransitionMapping(
  const Shapes::Properties::ShapeTransitionGroup& mappingsGroup,
  const ChiralStatePreservation& preservationOption
) {
  if(mappingsGroup.indexMappings.empty()) {
    return boost::none;
  }

  if(
    preservationOption == ChiralStatePreservation::EffortlessAndUnique
    && mappingsGroup.indexMappings.size() == 1
    && mappingsGroup.angularDistortion <= 0.2
  ) {
    return mappingsGroup.indexMappings.front();
  }

  if(
    preservationOption == ChiralStatePreservation::Unique
    && mappingsGroup.indexMappings.size() == 1
  ) {
    return mappingsGroup.indexMappings.front();
  }

  if(preservationOption == ChiralStatePreservation::RandomFromMultipleBest) {
    return Temple::Random::pick(mappingsGroup.indexMappings, randomnessEngine());
  }

  return boost::none;
}

bool AtomStereopermutator::Impl::thermalized(
  const Graph& graph,
  const AtomIndex centerAtom,
  const Shapes::Shape shape,
  const RankingInformation& ranking
) {
  if(Options::Thermalization::pyramidalInversion) {
    /* Nitrogen atom inversion */
    constexpr unsigned nitrogenZ = Utils::ElementInfo::Z(Utils::ElementType::N);
    const bool isNitrogenIsotope = Utils::ElementInfo::Z(graph.elementType(centerAtom)) == nitrogenZ;

    if(isNitrogenIsotope && shape == Shapes::Shape::VacantTetrahedron) {
      const bool inSmallCycle = Temple::any_of(
        ranking.links,
        [](const RankingInformation::Link& link) {
          return link.cycleSequence.size() <= 4;
        }
      );
      return !inSmallCycle;
    }
  }

  if(ranking.links.empty()) {
    if(Options::Thermalization::berryPseudorotation && shape == Shapes::Shape::TrigonalBipyramid) {
      return true;
    }

    if(Options::Thermalization::bartellMechanism && shape == Shapes::Shape::PentagonalBipyramid) {
      return true;
    }
  }

  return false;
}

/* Constructors */
AtomStereopermutator::Impl::Impl(
  const Graph& graph,
  // The symmetry of this Stereopermutator
  const Shapes::Shape shape,
  // The atom this Stereopermutator is centered on
  const AtomIndex centerAtom,
  // Ranking information of substituents
  RankingInformation ranking
) : centerAtom_ {centerAtom},
    shape_ {shape},
    ranking_ {std::move(ranking)},
    abstract_ {ranking_, shape_},
    feasible_ {abstract_, shape_, centerAtom_, ranking_, graph},
    assignmentOption_ {boost::none},
    shapePositionMap_ {},
    thermalized_ {thermalized(
      graph,
      centerAtom,
      shape_,
      ranking_
    )}
{}

/* Modification */
void AtomStereopermutator::Impl::assign(boost::optional<unsigned> assignment) {
  if(assignment && assignment.value() >= feasible_.indices.size()) {
    throw std::out_of_range("Supplied assignment index is out of range");
  }

  // Store new assignment
  assignmentOption_ = std::move(assignment);

  /* save a mapping of next neighbor indices to symmetry positions after
   * assigning (AtomIndex -> unsigned).
   */
  if(assignmentOption_) {
    shapePositionMap_ = siteToShapeVertexMap(
      abstract_.permutations.list.at(
        feasible_.indices.at(
          assignmentOption_.value()
        )
      ),
      abstract_.canonicalSites,
      ranking_.links
    );
  } else { // Wipe the map
    shapePositionMap_.clear();
  }
}

void AtomStereopermutator::Impl::assign(std::vector<Shapes::Vertex> vertexMapping) {
  if(vertexMapping.size() != Shapes::size(shape_) + 1) {
    throw std::invalid_argument("Vertex mapping for assignment has wrong length");
  }

  /* Drop the centroid, making the remaining sequence viable as an index mapping
   * within the set of shape rotations, i.e. for use in stereopermutation
   * functionality
   */
  vertexMapping.pop_back();

  /* Ok, so: We have a mapping from site positions to shape vertices from the
   * shape classification algorithm. Site positions are ordered identically to
   * the sites, so we essentially have a mapping from sites to shape vertices.
   *
   * Now we need to find the (feasible) stereopermutation that matches it.
   *
   * Can we transform the site -> shape vertex mapping into a
   * stereopermutation, generate all its rotations and then set-membership
   * check each feasible permutation?
   */
  auto soughtStereopermutation = stereopermutationFromSiteToShapeVertexMap(
    SiteToShapeVertexMap {vertexMapping},
    ranking_.links,
    abstract_.canonicalSites
  );

  auto soughtRotations = Stereopermutations::generateAllRotations(soughtStereopermutation, shape_);

  /* Although the stereopermutations in abstract_ are sorted by their index of
   * permutation and we could potentially binary search all of our sought
   * rotations within that set, this linear search will never be time-critical
   * as long as we use continuous shape measure-based shape classification.
   */
  boost::optional<unsigned> foundStereopermutation;
  const unsigned A = feasible_.indices.size();
  for(unsigned a = 0; a < A; ++a) {
    const auto& feasiblePermutation = abstract_.permutations.list.at(
      feasible_.indices.at(a)
    );
    auto findIter = std::find(
      std::begin(soughtRotations),
      std::end(soughtRotations),
      feasiblePermutation
    );
    if(findIter != std::end(soughtRotations)) {
      foundStereopermutation = a;
      break;
    }
  }

  if(foundStereopermutation) {
    assignmentOption_ = *foundStereopermutation;
    shapePositionMap_ = ShapeMap {vertexMapping};
  }
}

void AtomStereopermutator::Impl::assignRandom(Random::Engine& engine) {
  const unsigned A = numAssignments();
  if(A == 0) {
    throw std::logic_error("Cannot randomly assign a stereopermutator without feasible stereopermutations");
  }

  if(A == 1) {
    assign(0);
  } else {
    assign(
      Temple::Random::pickDiscrete(
        // Map the feasible permutations onto their weights
        Temple::map(
          feasible_.indices,
          [&](const unsigned permutationIndex) -> unsigned {
            return abstract_.permutations.weights.at(permutationIndex);
          }
        ),
        engine
      )
    );
  }
}

void AtomStereopermutator::Impl::applyPermutation(const std::vector<AtomIndex>& permutation) {
  // RankingInformation changes (lots of atom indices)
  ranking_.applyPermutation(permutation);

  // centerAtom_ must change
  centerAtom_ = permutation.at(centerAtom_);

  /* Neither shape nor assignment change. Also, although ranking and central
   * atom are implicated in their creation, abstract and feasible's states are
   * independent of atom indices.
   */
}

boost::optional<AtomStereopermutator::PropagatedState> AtomStereopermutator::Impl::propagate(
  const Graph& graph,
  RankingInformation newRanking,
  boost::optional<Shapes::Shape> shapeOption
) {
#ifndef NDEBUG
  /* Check preconditions! */
  for(const auto& siteAtomList : boost::range::join(ranking_.sites, newRanking.sites)) {
    assert(
      std::is_sorted(
        std::begin(siteAtomList),
        std::end(siteAtomList)
      )
    );
  }
#endif

  /* There are a lot of changes that can occur prior to a call to propagate that
   * we need to correct. Multiple of these can apply at any given time!
   *
   * - Substituent count change
   *   - A new substituent may be added to a site or removed
   *   - Can happen at the same time as a site removal (three sites to two, one
   *     newly haptic)
   * - Ranking order change
   *   - Order of ranked sites may change
   *   - Implies a stereopermutation change!
   * - Site count change
   *   - Determined by comparing rankings
   *   - Implies a shape change
   *   - Need to make a map between old and new site indices (this can be
   *     complicated by substituent-level changes)
   * - Shape change
   *   - Triggered by a site count change only
   *   - Need to determine a map between old and new shape vertices
   *
   * Sketch of the following procedure
   * - Determine mapping between sites, taking care to account for site and
   *   substituent count changes
   * - Determine mapping of shape vertices between shapes if a shape change
   *   happens
   * - Use shape map of current assignment to place old site indices at old shape vertices
   * - Use mapping between shape vertices to enter the new shape
   * - Use mapping between site indices to place new sites at the new shape
   * - Make a stereopermutation from the new shape map and look for it in the
   *   new feasibles
   */

  // If nothing changes, nothing changes, and nothing is propagated
  if(newRanking == ranking_) {
    return boost::none;
  }

  const int siteCountChange = static_cast<int>(newRanking.sites.size() - ranking_.sites.size());
  // Complain if more than one site is changed either way
  if(std::abs(siteCountChange) > 1) {
    throw std::logic_error("propagateGraphChange should only ever handle a single site addition or removal at once");
  }

  auto siteMapping = SiteMapping::from(ranking_, newRanking);

  /* Decide the new shape */
  Shapes::Shape newShape = shapeOption.value_or_eval(
    [&]() {
      if(siteCountChange == +1) {
        return up(shape_);
      }

      if(siteCountChange == 0) {
        return shape_;
      }

      assert(siteCountChange == -1);
      assert(siteMapping.changedSite);

      /* We can only figure out a mapping if the stereocenter is assigned
       * since otherwise cache_.symmetryPositionMap is empty and the shape
       * vertex being removed is a necessary argument to down.
       */
      if(assignmentOption_) {
        return down(
          shape_,
          shapePositionMap_.at(siteMapping.changedSite.value())
        );
      }

      // Just return the most symmetric symmetry of the target size
      return Shapes::Properties::mostSymmetric(
        newRanking.sites.size()
      );
    }
  );

  /* Generate new assignments */
  Stereopermutators::Abstract newAbstract {
    newRanking,
    newShape
  };

  Stereopermutators::Feasible newFeasible {
    newAbstract,
    newShape,
    centerAtom_,
    newRanking,
    graph
  };

  boost::optional<unsigned> newStereopermutationOption;

  /* Now we will attempt to propagate our assignment to the new set of
   * assignments. This is only necessary in the case that the stereopermutator
   * is currently assigned.
   *
   * In the case of an equal number of sites, it is only possible to
   * propagate chiral state if the new number of assignments is smaller or
   * equal to the amount we have currently. This is because say we have AABCDE
   * in octahedral, we carry chiral state, and a ranking change leads to
   * ABCDEF. There will be multiple ways to split A to F!
   */
  if(assignmentOption_ && numStereopermutations() > 1) {
    const auto removedVertexOptional = Temple::Optionals::flatMap(
      siteMapping.changedSite,
      [this, siteCountChange](const SiteIndex changedSite) -> boost::optional<Shapes::Vertex> {
        if(siteCountChange == -1) {
          return shapePositionMap_.at(changedSite);
        }

        return boost::none;
      }
    );

    const auto shapeMapping = Temple::Optionals::flatMap(
      Shapes::getMapping(shape_, newShape, removedVertexOptional),
      [](const auto& mapping) {
        return selectTransitionMapping(mapping, Options::chiralStatePreservation);
      }
    );

    const auto appliedShapeMapping = Temple::Optionals::map(
      shapeMapping,
      [&](const std::vector<Shapes::Vertex>& vertexMap) {
        if(siteCountChange == +1) {
          return Shapes::Properties::applyIndexMapping(newShape, vertexMap);
        }

        return vertexMap;
      }
    );

    if(appliedShapeMapping) {
      const auto& shapeMappingValue = appliedShapeMapping.value();

      // Use shape map of current assignment to place old site indices at old shape vertices
      using InvertedMap = Temple::StrongIndexFlatMap<Shapes::Vertex, SiteIndex>;
      const InvertedMap oldSiteIndicesAtOldVertices = shapePositionMap_.invert();

      // Use shape mapping if available to transfer old vertices to new vertices
      const InvertedMap oldSiteIndicesAtNewVertices = InvertedMap(
        Temple::map(
          Temple::iota<Shapes::Vertex>(Shapes::size(newShape)),
          [&](const Shapes::Vertex newVertex) -> SiteIndex {
            const Shapes::Vertex oldVertex = shapeMappingValue.at(newVertex);

            if(siteCountChange == +1 && oldVertex == Shapes::size(newShape) - 1)  {
              return siteMapping.changedSite.value();
            }

            return oldSiteIndicesAtOldVertices.at(oldVertex);
          }
        )
      );

      // Use site mapping to transfer old site indices to new site indices
      const InvertedMap sitesAtNewShapeVertices = InvertedMap(
        Temple::map(
          oldSiteIndicesAtNewVertices,
          [&](const SiteIndex oldSite) -> SiteIndex {
            auto findIter = siteMapping.map.find(oldSite);

            if(findIter != std::end(siteMapping.map)) {
              return findIter->second;
            }

            assert(siteCountChange == +1);
            return oldSite;
          }
        )
      );

      assert(
        Temple::all_of(
          sitesAtNewShapeVertices,
          [newShape](const SiteIndex i) { return i < Shapes::size(newShape); }
        )
      );

      // Replace the site indices by their new ranking characters
      const auto newStereopermutationCharacters = Stereopermutators::Abstract::makeStereopermutationCharacters(
        newAbstract.canonicalSites,
        newAbstract.symbolicCharacters,
        sitesAtNewShapeVertices
      );

      // Create a new assignment with those characters
      const auto trialStereopermutation = Stereopermutations::Stereopermutation(
        newStereopermutationCharacters,
        newAbstract.selfReferentialLinks
      );

      // Generate all rotations of this trial assignment
      auto allTrialRotations = Stereopermutations::generateAllRotations(trialStereopermutation, newShape);

      // Find out which of the new assignments has a rotational equivalent
      for(unsigned i = 0; i < newAbstract.permutations.list.size(); ++i) {
        auto findIter = std::find(
          std::begin(allTrialRotations),
          std::end(allTrialRotations),
          newAbstract.permutations.list.at(i)
        );
        if(findIter != std::end(allTrialRotations)) {
          newStereopermutationOption = i;
          break;
        }
      }
    }
  }

  /* If we have discovered a stereopermutation to map to, we must still figure
   * out its assignment index (assuming it is feasible)
   */
  auto newAssignmentOption = Temple::Optionals::flatMap(
    newStereopermutationOption,
    [&](const unsigned stereopermutationIndex) -> boost::optional<unsigned> {
      auto assignmentFindIter = std::find(
        std::begin(newFeasible.indices),
        std::end(newFeasible.indices),
        stereopermutationIndex
      );

      if(assignmentFindIter == std::end(newFeasible.indices)) {
        // The stereopermutation is infeasible
        return boost::none;
      }

      return assignmentFindIter - std::begin(newFeasible.indices);
    }
  );

  // Extract old state from the class
  auto oldStateTuple = std::make_tuple(
    std::move(ranking_),
    std::move(abstract_),
    std::move(feasible_),
    std::move(shapePositionMap_)
  );

  // Overwrite the class state
  shape_ = newShape;
  ranking_ = std::move(newRanking);
  abstract_ = std::move(newAbstract);
  feasible_ = std::move(newFeasible);
  thermalized_ = thermalized(
    graph,
    centerAtom_,
    shape_,
    ranking_
  );
  assign(newAssignmentOption);

  return {std::move(oldStateTuple)};
}

void AtomStereopermutator::Impl::propagateVertexRemoval(const AtomIndex removedIndex) {
  /* This function replaces any occurrences of the atom index that is being
   * removed in the global state with a placeholder of the same type and updates
   * any invalidated atom indices.
   */

  /* If the central atom is being removed, just drop this stereopermutator
   * beforehand in caller. This would just be unnecessary work.
   */
  assert(centerAtom_ != removedIndex);

  // Define some helper functions
  auto updateIndexInplace = [&removedIndex](AtomIndex& index) -> void {
    if(index > removedIndex) {
      --index;
    } else if(index == removedIndex) {
      index = PrivateGraph::removalPlaceholder;
    }
  };

  auto updateIndex = [&removedIndex](const AtomIndex index) -> AtomIndex {
    if(index > removedIndex) {
      return index - 1;
    }

    if(index == removedIndex) {
      return PrivateGraph::removalPlaceholder;
    }

    return index;
  };

  /* Update indices in RankingInformation */
  for(auto& equalPrioritySet : ranking_.substituentRanking) {
    for(auto& index : equalPrioritySet) {
      updateIndexInplace(index);
    }
  }

  for(auto& siteAtomList : ranking_.sites) {
    for(auto& atomIndex : siteAtomList) {
      updateIndexInplace(atomIndex);
    }
  }

  for(auto& link : ranking_.links) {
    link.cycleSequence = Temple::map(link.cycleSequence, updateIndex);
  }
}

const Stereopermutators::Abstract& AtomStereopermutator::Impl::getAbstract() const {
  return abstract_;
}

const Stereopermutators::Feasible& AtomStereopermutator::Impl::getFeasible() const {
  return feasible_;
}

const RankingInformation& AtomStereopermutator::Impl::getRanking() const {
  return ranking_;
}

Shapes::Shape AtomStereopermutator::Impl::getShape() const {
  return shape_;
}

const AtomStereopermutator::ShapeMap& AtomStereopermutator::Impl::getShapePositionMap() const {
  if(assignmentOption_ == boost::none) {
    throw std::logic_error(
      "The AtomStereopermutator is unassigned, sites are not assigned to "
      "symmetry positions"
    );
  }

  return shapePositionMap_;
}

boost::optional<AtomStereopermutator::ShapeMap>
AtomStereopermutator::Impl::fit(
  const Graph& graph,
  const AngstromPositions& angstromWrapper
) {
  const unsigned S = Shapes::size(shape_);
  assert(S == ranking_.sites.size());

  // Save stereopermutator state to return to if no fit is viable
  const Shapes::Shape priorShape = shape_;
  const boost::optional<unsigned> priorStereopermutation  = assignmentOption_;

  // For all atoms making up a site, decide on the spatial average position
  Eigen::Matrix<double, 3, Eigen::Dynamic> sitePositions(3, S + 1);
  for(unsigned i = 0; i < S; ++i) {
    sitePositions.col(i) = Cartesian::averagePosition(angstromWrapper.positions, ranking_.sites.at(i));
  }
  // Add the putative center
  sitePositions.col(S) = angstromWrapper.positions.row(centerAtom_);

  // Classify the shape and set it
  Shapes::Shape fittedShape;
  std::vector<Shapes::Vertex> matchingMapping;
  std::tie(fittedShape, matchingMapping) = classifyShape(sitePositions);
  setShape(fittedShape, graph);
  assign(matchingMapping);

  // Return to prior state if no feasible stereopermutations found
  if(assignmentOption_ == boost::none) {
    setShape(priorShape, graph);
    assign(priorStereopermutation);
    return boost::none;
  }

  matchingMapping.pop_back();
  return ShapeMap(matchingMapping);
}

/* Information */
double AtomStereopermutator::Impl::angle(
  const SiteIndex i,
  const SiteIndex j
) const {
  if(!assignmentOption_) {
    throw std::runtime_error("Stereopermutator is unassigned, angles are unknown!");
  }

  const unsigned S = Shapes::size(shape_);
  if(i >= S || j >= S) {
    throw std::out_of_range("Site index is out of range");
  }

  return Shapes::angleFunction(shape_)(
    shapePositionMap_.at(i),
    shapePositionMap_.at(j)
  );
}

boost::optional<unsigned> AtomStereopermutator::Impl::assigned() const {
  if(thermalized_) {
    return Temple::Optionals::map(assignmentOption_, [](unsigned /* a */) { return 0u; });
  }

  return assignmentOption_;
}

AtomIndex AtomStereopermutator::Impl::placement() const {
  return centerAtom_;
}

boost::optional<unsigned> AtomStereopermutator::Impl::indexOfPermutation() const {
  if(thermalized_) {
    return Temple::Optionals::map(assignmentOption_, [](unsigned /* a */) { return 0u; });
  }

  return Temple::Optionals::map(
    assignmentOption_,
    Temple::Functor::at(feasible_.indices)
  );
}

std::vector<AtomStereopermutator::MinimalChiralConstraint>
AtomStereopermutator::Impl::minimalChiralConstraints(const bool enforce) const {
  /* It only makes sense to emit these minimal representations of chiral
   * constraints if the stereopermutator is assigned.
   *
   * As long as the permutator is achiral, there is little reason to burden
   * refinement with error terms that achieve nothing for the final
   * conformation. Yet, if any sort of BondStereopermutator is placed on even
   * an achiral AtomStereopermutator, assumptions are made about the spatial
   * placement of symmetry positions.
   *
   * Since refinement has many terms, it can be difficult to reliably minimize
   * these assumptions on symmetry position spatial locations against partially
   * relaxed distance constraints without the clear ordering priority of chiral
   * constraints enforcing your spatial reasoning assumptions.
   *
   * So for now, until a better solution arises, emit chiral constraints if a
   * bond stereopermutator relies upon symmetry position assumptions even if
   * this stereopermutator might be achiral.
   */
  if(
    assignmentOption_
    && (numAssignments() > 1 || enforce)
  ) {
    /* Invert neighborSymmetryPositionMap_, we need a mapping of
     *  (vertex in shape) -> site index
     */
    const auto invertedShapeMap = shapePositionMap_.invert();

    return Temple::map(
      Shapes::tetrahedra(shape_),
      [&](const auto& tetrahedron) -> MinimalChiralConstraint {
        return Temple::map_stl(
          tetrahedron,
          [&](const auto& shapeVertexOptional) -> boost::optional<SiteIndex> {
            return Temple::Optionals::map(
              shapeVertexOptional,
              Temple::Functor::at(invertedShapeMap)
            );
          }
        );
      }
    );
  }

  return {};
}

std::string AtomStereopermutator::Impl::info() const {
  std::string returnString = std::to_string(centerAtom_) + ": "s + Shapes::name(shape_) +", "s;

  const auto& characters = abstract_.symbolicCharacters;
  std::copy(
    characters.begin(),
    characters.end(),
    std::back_inserter(returnString)
  );

  for(const auto& link : abstract_.selfReferentialLinks) {
    returnString += ", "s + characters.at(link.first) + "-"s + characters.at(link.second);
  }

  returnString += ", "s;

  if(assignmentOption_) {
    returnString += std::to_string(assignmentOption_.value());
  } else {
    returnString += "u";
  }

  const unsigned A = numAssignments();
  returnString += " ("s + std::to_string(A);

  const unsigned P = numStereopermutations();
  if(P != A) {
    returnString += ", "s + std::to_string(P);
  }

  returnString += ")";

  return returnString;
}

std::string AtomStereopermutator::Impl::rankInfo() const {
  /* rankInfo is specifically geared towards RankingTree's consumption,
   * and MUST use indices of permutation
   */
  return (
    "A-"s + std::to_string(static_cast<unsigned>(shape_))
    + "-"s + std::to_string(numStereopermutations())
    + "-"s + (
      indexOfPermutation()
      ? std::to_string(indexOfPermutation().value())
      : "u"s
    )
  );
}

std::vector<std::vector<SiteIndex>> AtomStereopermutator::Impl::siteGroups() const {
  if(!assignmentOption_) {
    throw std::logic_error("Stereopermutator is not assigned");
  }

  return Temple::map(
    Shapes::Properties::positionGroups(shape_),
    [this](const std::vector<Shapes::Vertex>& interconvertibleVertices) {
      return Temple::map(
        interconvertibleVertices,
        [this](const Shapes::Vertex v) -> SiteIndex {
          return shapePositionMap_.indexOf(v);
        }
      );
    }
  );
}

unsigned AtomStereopermutator::Impl::numAssignments() const {
  if(thermalized_) {
    return 1;
  }

  return feasible_.indices.size();
}

unsigned AtomStereopermutator::Impl::numStereopermutations() const {
  if(thermalized_) {
    return 1;
  }

  return abstract_.permutations.list.size();
}

void AtomStereopermutator::Impl::setShape(
  const Shapes::Shape shape,
  const Graph& graph
) {
  if(shape_ == shape) {
    // If the symmetry doesn't actually change, then nothing does
    return;
  }

  shape_ = shape;

  abstract_ = Stereopermutators::Abstract {
    ranking_,
    shape_
  };

  feasible_ = Stereopermutators::Feasible {
    abstract_,
    shape_,
    centerAtom_,
    ranking_,
    graph
  };

  thermalized_ = thermalized(
    graph,
    centerAtom_,
    shape_,
    ranking_
  );

  // Dis-assign the stereopermutator
  assign(boost::none);
}

} // namespace Molassembler
} // namespace Scine


/*********************************************************************************/
/*                                                                               */
/*     Command-line tool to access support vector machine                        */
/*                          classification functionality.                        */
/*                                                                               */
/*     Alle Meije Wink                                                           */
/*                                                                               */
/*********************************************************************************/

/*
  Update history

  Who    When       What
  AMW    13-10-12   creation

*/

#include "combisDesign.hpp"
#include <dlib/svm.h>

using namespace dlib;

namespace cb = canabis;                       // computer analysis of brain image statistics

int combisICA(cb::bisNiftiImage<intensity> *currentImage, std::string designfile)
{
    // classification table
    std::vector<size_t>               mask(1);
    std::vector<intensity>               y(1);
    std::vector<std::vector <intensity>> x(1);

    // read the design
    auto design_isbinary=combisDesign(designfile, &x, &y, &mask, currentImage);                    // read text / binary design

    double noiseStdDev = 0.25; //0.175;
    size_t nReplicates = 10;   //30;
    size_t nAngles     = 10;   //150;
    size_t nSweeps     = 0;

    auto xs  = x.size();
    auto x0s = x[0].size();
    arma::mat matX ( x0s, xs ); // column major

    for ( size_t i=0; i<xs; i++ )
        for (size_t j=0; j<x0s; j++)
            matX(j,i)=x[i][j];

    // initialise radical
    nSweeps = matX.n_rows - 1;
    //mlpack::radical::Radical rad(noiseStdDev, nReplicates, nAngles, nSweeps);

    // run radical
    arma::mat matY;
    arma::mat matW;
    //rad.DoRadical(matX.t(), matY, matW);

    std::cout << "completed";
}


#ifndef RANDOM_HPP
#define RANDOM_HPP

#include <iostream>
#include <boost/random.hpp>

//=================================================
// singleton class for random number generator
class Random {
public:
  static void Init(uint64_t seed, int num_threads = 1) {
    m_rnds.clear();

    for(int i=0; i <num_threads; i++) {
      m_rnds.push_back( boost::random::mt19937(seed+i) );
    }
  }
  static double Rand01(int thread_num) {
    boost::random::uniform_01<> uniform;
    return uniform(m_rnds[thread_num]);
  }
  static double Gaussian(int thread_num) {
    boost::random::normal_distribution<> gaussian;
    return gaussian(m_rnds[thread_num]);
  }
private:
  static std::vector<boost::random::mt19937> m_rnds;
};

#endif


//
// Created by cedric on 2021-07-25.
//

#define BOOST_TEST_DYN_LINK
#define BOOST_TEST_MODULE EightBitArithmeticTest

#include <boost/test/included/unit_test.hpp>
#include <exception>
#include "cpu/CPU.h"
#include "util/testing_utilities.h"

BOOST_AUTO_TEST_SUITE(EightBitArithmeticTest)

    BOOST_AUTO_TEST_CASE(validate0x04) {
        CPU cpu;
        set_flags(&cpu, FLAGS_NONE);
        cpu.load_register_immediate(REGISTER_B_INDEX, 0x3E);

        run_next_opcode(&cpu, {0x04});
        BOOST_CHECK(cpu.get_registers()[REGISTER_B_INDEX] == 0x3F);
        BOOST_CHECK(!cpu.is_zero_flag_on());
        BOOST_CHECK(!cpu.is_subtract_flag_on());
        BOOST_CHECK(!cpu.is_half_carry_flag_on());
        BOOST_CHECK(!cpu.is_carry_flag_on());

        run_next_opcode(&cpu, {0x04});
        BOOST_CHECK(cpu.get_registers()[REGISTER_B_INDEX] == 0x40);
        BOOST_CHECK(!cpu.is_zero_flag_on());
        BOOST_CHECK(!cpu.is_subtract_flag_on());
        BOOST_CHECK(cpu.is_half_carry_flag_on());
        BOOST_CHECK(!cpu.is_carry_flag_on());

        cpu.load_register_immediate(REGISTER_B_INDEX, 0xFF);
        run_next_opcode(&cpu, {0x04});
        BOOST_CHECK(cpu.get_registers()[REGISTER_B_INDEX] == 0x00);
        BOOST_CHECK(cpu.is_zero_flag_on());
        BOOST_CHECK(!cpu.is_subtract_flag_on());
        BOOST_CHECK(cpu.is_half_carry_flag_on());
        BOOST_CHECK(!cpu.is_carry_flag_on());
    }

    BOOST_AUTO_TEST_CASE(validate0x0C) {
        CPU cpu;
        set_flags(&cpu, FLAGS_NONE);
        cpu.load_register_immediate(REGISTER_C_INDEX, 0x1E);

        run_next_opcode(&cpu, {0x0C});
        BOOST_CHECK(cpu.get_registers()[REGISTER_C_INDEX] == 0x1F);
        BOOST_CHECK(!cpu.is_zero_flag_on());
        BOOST_CHECK(!cpu.is_subtract_flag_on());
        BOOST_CHECK(!cpu.is_half_carry_flag_on());
        BOOST_CHECK(!cpu.is_carry_flag_on());

        run_next_opcode(&cpu, {0x0C});
        BOOST_CHECK(cpu.get_registers()[REGISTER_C_INDEX] == 0x20);
        BOOST_CHECK(!cpu.is_zero_flag_on());
        BOOST_CHECK(!cpu.is_subtract_flag_on());
        BOOST_CHECK(cpu.is_half_carry_flag_on());
        BOOST_CHECK(!cpu.is_carry_flag_on());

        cpu.load_register_immediate(REGISTER_C_INDEX, 0xFF);
        run_next_opcode(&cpu, {0x0C});
        BOOST_CHECK(cpu.get_registers()[REGISTER_C_INDEX] == 0x00);
        BOOST_CHECK(cpu.is_zero_flag_on());
        BOOST_CHECK(!cpu.is_subtract_flag_on());
        BOOST_CHECK(cpu.is_half_carry_flag_on());
        BOOST_CHECK(!cpu.is_carry_flag_on());
    }

    BOOST_AUTO_TEST_CASE(validate0x14) {
        CPU cpu;
        set_flags(&cpu, FLAGS_NONE);
        cpu.load_register_immediate(REGISTER_D_INDEX, 0x0E);

        run_next_opcode(&cpu, {0x14});
        BOOST_CHECK(cpu.get_registers()[REGISTER_D_INDEX] == 0x0F);
        BOOST_CHECK(!cpu.is_zero_flag_on());
        BOOST_CHECK(!cpu.is_subtract_flag_on());
        BOOST_CHECK(!cpu.is_half_carry_flag_on());
        BOOST_CHECK(!cpu.is_carry_flag_on());

        run_next_opcode(&cpu, {0x14});
        BOOST_CHECK(cpu.get_registers()[REGISTER_D_INDEX] == 0x10);
        BOOST_CHECK(!cpu.is_zero_flag_on());
        BOOST_CHECK(!cpu.is_subtract_flag_on());
        BOOST_CHECK(cpu.is_half_carry_flag_on());
        BOOST_CHECK(!cpu.is_carry_flag_on());

        cpu.load_register_immediate(REGISTER_D_INDEX, 0xFF);
        run_next_opcode(&cpu, {0x14});
        BOOST_CHECK(cpu.get_registers()[REGISTER_D_INDEX] == 0x00);
        BOOST_CHECK(cpu.is_zero_flag_on());
        BOOST_CHECK(!cpu.is_subtract_flag_on());
        BOOST_CHECK(cpu.is_half_carry_flag_on());
        BOOST_CHECK(!cpu.is_carry_flag_on());
    }

    BOOST_AUTO_TEST_CASE(validate0x1C) {
        CPU cpu;
        set_flags(&cpu, FLAGS_NONE);
        cpu.load_register_immediate(REGISTER_E_INDEX, 0x9E);

        run_next_opcode(&cpu, {0x1C});
        BOOST_CHECK(cpu.get_registers()[REGISTER_E_INDEX] == 0x9F);
        BOOST_CHECK(!cpu.is_zero_flag_on());
        BOOST_CHECK(!cpu.is_subtract_flag_on());
        BOOST_CHECK(!cpu.is_half_carry_flag_on());
        BOOST_CHECK(!cpu.is_carry_flag_on());

        run_next_opcode(&cpu, {0x1C});
        BOOST_CHECK(cpu.get_registers()[REGISTER_E_INDEX] == 0xA0);
        BOOST_CHECK(!cpu.is_zero_flag_on());
        BOOST_CHECK(!cpu.is_subtract_flag_on());
        BOOST_CHECK(cpu.is_half_carry_flag_on());
        BOOST_CHECK(!cpu.is_carry_flag_on());

        cpu.load_register_immediate(REGISTER_E_INDEX, 0xFF);
        run_next_opcode(&cpu, {0x1C});
        BOOST_CHECK(cpu.get_registers()[REGISTER_E_INDEX] == 0x00);
        BOOST_CHECK(cpu.is_zero_flag_on());
        BOOST_CHECK(!cpu.is_subtract_flag_on());
        BOOST_CHECK(cpu.is_half_carry_flag_on());
        BOOST_CHECK(!cpu.is_carry_flag_on());
    }

    BOOST_AUTO_TEST_CASE(validate0x24) {
        CPU cpu;
        set_flags(&cpu, FLAGS_NONE);
        cpu.load_register_immediate(REGISTER_H_INDEX, 0x6E);

        run_next_opcode(&cpu, {0x24});
        BOOST_CHECK(cpu.get_registers()[REGISTER_H_INDEX] == 0x6F);
        BOOST_CHECK(!cpu.is_zero_flag_on());
        BOOST_CHECK(!cpu.is_subtract_flag_on());
        BOOST_CHECK(!cpu.is_half_carry_flag_on());
        BOOST_CHECK(!cpu.is_carry_flag_on());

        run_next_opcode(&cpu, {0x24});
        BOOST_CHECK(cpu.get_registers()[REGISTER_H_INDEX] == 0x70);
        BOOST_CHECK(!cpu.is_zero_flag_on());
        BOOST_CHECK(!cpu.is_subtract_flag_on());
        BOOST_CHECK(cpu.is_half_carry_flag_on());
        BOOST_CHECK(!cpu.is_carry_flag_on());

        cpu.load_register_immediate(REGISTER_H_INDEX, 0xFF);
        run_next_opcode(&cpu, {0x24});
        BOOST_CHECK(cpu.get_registers()[REGISTER_H_INDEX] == 0x00);
        BOOST_CHECK(cpu.is_zero_flag_on());
        BOOST_CHECK(!cpu.is_subtract_flag_on());
        BOOST_CHECK(cpu.is_half_carry_flag_on());
        BOOST_CHECK(!cpu.is_carry_flag_on());
    }

    BOOST_AUTO_TEST_CASE(validate0x37) {
        CPU cpu;
        set_flags(&cpu, FLAGS_NONE);
        run_next_opcode(&cpu, {0x37});
        BOOST_CHECK(cpu.is_carry_flag_on());
    }

    BOOST_AUTO_TEST_CASE(validate0x2C) {
        CPU cpu;
        set_flags(&cpu, FLAGS_NONE);
        cpu.load_register_immediate(REGISTER_L_INDEX, 0xBE);

        run_next_opcode(&cpu, {0x2C});
        BOOST_CHECK(cpu.get_registers()[REGISTER_L_INDEX] == 0xBF);
        BOOST_CHECK(!cpu.is_zero_flag_on());
        BOOST_CHECK(!cpu.is_subtract_flag_on());
        BOOST_CHECK(!cpu.is_half_carry_flag_on());
        BOOST_CHECK(!cpu.is_carry_flag_on());

        run_next_opcode(&cpu, {0x2C});
        BOOST_CHECK(cpu.get_registers()[REGISTER_L_INDEX] == 0xC0);
        BOOST_CHECK(!cpu.is_zero_flag_on());
        BOOST_CHECK(!cpu.is_subtract_flag_on());
        BOOST_CHECK(cpu.is_half_carry_flag_on());
        BOOST_CHECK(!cpu.is_carry_flag_on());

        cpu.load_register_immediate(REGISTER_L_INDEX, 0xFF);
        run_next_opcode(&cpu, {0x2C});
        BOOST_CHECK(cpu.get_registers()[REGISTER_L_INDEX] == 0x00);
        BOOST_CHECK(cpu.is_zero_flag_on());
        BOOST_CHECK(!cpu.is_subtract_flag_on());
        BOOST_CHECK(cpu.is_half_carry_flag_on());
        BOOST_CHECK(!cpu.is_carry_flag_on());
    }

    BOOST_AUTO_TEST_CASE(validate0x34) {
        CPU cpu;
        set_flags(&cpu, FLAGS_NONE);
        cpu.store_memory_immediate(0x21F2, 0xAE);
        cpu.load_16bit_register_immediate(REGISTER_HL_INDEX, 0x21F2);

        run_next_opcode(&cpu, {0x34});
        BOOST_CHECK(cpu.get_byte_memory_indirect(REGISTER_HL_INDEX) == 0xAF);
        BOOST_CHECK(!cpu.is_zero_flag_on());
        BOOST_CHECK(!cpu.is_subtract_flag_on());
        BOOST_CHECK(!cpu.is_half_carry_flag_on());
        BOOST_CHECK(!cpu.is_carry_flag_on());

        run_next_opcode(&cpu, {0x34});
        BOOST_CHECK(cpu.get_byte_memory_indirect(REGISTER_HL_INDEX) == 0xB0);
        BOOST_CHECK(!cpu.is_zero_flag_on());
        BOOST_CHECK(!cpu.is_subtract_flag_on());
        BOOST_CHECK(cpu.is_half_carry_flag_on());
        BOOST_CHECK(!cpu.is_carry_flag_on());

        cpu.store_memory_immediate(0x21F2, 0xFF);
        run_next_opcode(&cpu, {0x34});
        BOOST_CHECK(cpu.get_byte_memory_indirect(REGISTER_HL_INDEX) == 0x00);
        BOOST_CHECK(cpu.is_zero_flag_on());
        BOOST_CHECK(!cpu.is_subtract_flag_on());
        BOOST_CHECK(cpu.is_half_carry_flag_on());
        BOOST_CHECK(!cpu.is_carry_flag_on());
    }

    BOOST_AUTO_TEST_CASE(validate0x3C) {
        CPU cpu;
        set_flags(&cpu, FLAGS_NONE);
        cpu.load_register_immediate(REGISTER_A_INDEX, 0x5E);

        run_next_opcode(&cpu, {0x3C});
        BOOST_CHECK(cpu.get_registers()[REGISTER_A_INDEX] == 0x5F);
        BOOST_CHECK(!cpu.is_zero_flag_on());
        BOOST_CHECK(!cpu.is_subtract_flag_on());
        BOOST_CHECK(!cpu.is_half_carry_flag_on());
        BOOST_CHECK(!cpu.is_carry_flag_on());

        run_next_opcode(&cpu, {0x3C});
        BOOST_CHECK(cpu.get_registers()[REGISTER_A_INDEX] == 0x60);
        BOOST_CHECK(!cpu.is_zero_flag_on());
        BOOST_CHECK(!cpu.is_subtract_flag_on());
        BOOST_CHECK(cpu.is_half_carry_flag_on());
        BOOST_CHECK(!cpu.is_carry_flag_on());

        cpu.load_register_immediate(REGISTER_A_INDEX, 0xFF);
        run_next_opcode(&cpu, {0x3C});
        BOOST_CHECK(cpu.get_registers()[REGISTER_A_INDEX] == 0x00);
        BOOST_CHECK(cpu.is_zero_flag_on());
        BOOST_CHECK(!cpu.is_subtract_flag_on());
        BOOST_CHECK(cpu.is_half_carry_flag_on());
        BOOST_CHECK(!cpu.is_carry_flag_on());
    }

    BOOST_AUTO_TEST_CASE(validate0xF6) {
        CPU cpu;
        set_flags(&cpu, FLAGS_NONE);
        cpu.load_register_immediate(REGISTER_A_INDEX, 0x00);

        run_next_opcode(&cpu, {0xF6, 0x00});
        BOOST_CHECK(cpu.get_registers()[REGISTER_A_INDEX] == 0x00);
        BOOST_CHECK(cpu.is_zero_flag_on());
        BOOST_CHECK(!cpu.is_subtract_flag_on());
        BOOST_CHECK(!cpu.is_half_carry_flag_on());
        BOOST_CHECK(!cpu.is_carry_flag_on());

        set_flags(&cpu, FLAGS_NONE);
        cpu.load_register_immediate(REGISTER_A_INDEX, 0xA3);
        run_next_opcode(&cpu, {0xF6, 0x18});
        BOOST_CHECK(cpu.get_registers()[REGISTER_A_INDEX] == 0xBB);
        BOOST_CHECK(!cpu.is_zero_flag_on());
        BOOST_CHECK(!cpu.is_subtract_flag_on());
        BOOST_CHECK(!cpu.is_half_carry_flag_on());
        BOOST_CHECK(!cpu.is_carry_flag_on());

        cpu.load_register_immediate(REGISTER_A_INDEX, 0xFF);
        run_next_opcode(&cpu, {0x3C});
        BOOST_CHECK(cpu.get_registers()[REGISTER_A_INDEX] == 0x00);
        BOOST_CHECK(cpu.is_zero_flag_on());
        BOOST_CHECK(!cpu.is_subtract_flag_on());
        BOOST_CHECK(cpu.is_half_carry_flag_on());
        BOOST_CHECK(!cpu.is_carry_flag_on());
    }

BOOST_AUTO_TEST_SUITE_END()


/* ----------------------------------------------------------------------------
 * Copyright 2017, Massachusetts Institute of Technology,
 * Cambridge, MA 02139
 * All Rights Reserved
 * Authors: Luca Carlone, et al. (see THANKS for the full author list)
 * See LICENSE for the license information
 * -------------------------------------------------------------------------- */

/**
 * @file   StereoCamera.cpp
 * @brief  Class describing a stereo camera.
 * @author Antoni Rosinol
 */

#include "kimera-vio/frontend/StereoCamera.h"

#include <Eigen/Core>

#include <opencv2/calib3d.hpp>
#include <opencv2/core.hpp>

#include <boost/utility.hpp>  // for tie

#include <gtsam/geometry/Cal3_S2.h>
#include <gtsam/geometry/Cal3_S2Stereo.h>
#include <gtsam/geometry/Pose3.h>
#include <gtsam/geometry/StereoCamera.h>

#include <glog/logging.h>

#include "kimera-vio/frontend/UndistorterRectifier.h"
#include "kimera-vio/frontend/CameraParams.h"
#include "kimera-vio/frontend/StereoFrame.h"
#include "kimera-vio/frontend/StereoMatchingParams.h"
#include "kimera-vio/utils/Macros.h"

namespace VIO {

StereoCamera::StereoCamera(const CameraParams& left_cam_params,
                           const CameraParams& right_cam_params,
                           const StereoMatchingParams& stereo_matching_params)
    : stereo_camera_impl_(),
      stereo_calibration_(nullptr),
      baseline_(0.0),
      left_cam_params_(left_cam_params),
      right_cam_params_(right_cam_params),
      left_cam_undistort_rectifier_(nullptr),
      right_cam_undistort_rectifier_(nullptr),
      stereo_matching_params_(stereo_matching_params) {
  computeRectificationParameters(left_cam_params,
                                 right_cam_params,
                                 &R1_,
                                 &R2_,
                                 &P1_,
                                 &P2_,
                                 &Q_,
                                 &ROI1_,
                                 &ROI2_);

  // Calc left camera pose after rectification
  // NOTE: OpenCV pose convention is the opposite, therefore the inverse.
  const gtsam::Rot3& camL_Rot_camLrect =
      UtilsOpenCV::cvMatToGtsamRot3(R1_).inverse();
  gtsam::Pose3 camL_Pose_camLrect(camL_Rot_camLrect, gtsam::Point3::Zero());
  B_Pose_camLrect_ = left_cam_params.body_Pose_cam_.compose(camL_Pose_camLrect);

  // Calc baseline (see L.2700 and L.2616 in
  // https://github.com/opencv/opencv/blob/master/modules/calib3d/src/calibration.cpp
  // NOTE: OpenCV pose convention is the opposite, therefore the missing -1.0
  CHECK_NE(Q_.at<double>(3, 2), 0.0);
  baseline_ = 1.0 / Q_.at<double>(3, 2);
  CHECK_GT(baseline_, 0.0);

  //! Create stereo camera calibration after rectification and undistortion.
  gtsam::Cal3_S2 left_undist_rect_cam_mat = UtilsOpenCV::Cvmat2Cal3_S2(P1_);
  stereo_calibration_ =
      boost::make_shared<gtsam::Cal3_S2Stereo>(left_undist_rect_cam_mat.fx(),
                                               left_undist_rect_cam_mat.fy(),
                                               left_undist_rect_cam_mat.skew(),
                                               left_undist_rect_cam_mat.px(),
                                               left_undist_rect_cam_mat.py(),
                                               baseline_);

  //! Create undistort rectifiers: these should be called after
  //! computeRectificationParameters.
  left_cam_undistort_rectifier_ =
      VIO::make_unique<UndistorterRectifier>(P1_, left_cam_params, R1_);
  right_cam_undistort_rectifier_ =
      VIO::make_unique<UndistorterRectifier>(P2_, right_cam_params, R2_);

  //! Create stereo camera implementation
  stereo_camera_impl_ =
      gtsam::StereoCamera(B_Pose_camLrect_, stereo_calibration_);
}

void StereoCamera::project(const LandmarksCV& lmks,
                           KeypointsCV* left_kpts,
                           KeypointsCV* right_kpts) const {
  CHECK_NOTNULL(left_kpts)->clear();
  CHECK_NOTNULL(right_kpts)->clear();
  const auto& n_lmks = lmks.size();
  left_kpts->reserve(n_lmks);
  right_kpts->reserve(n_lmks);
  for (const auto& lmk : lmks) {
    const gtsam::StereoPoint2& kp =
        stereo_camera_impl_.project2(gtsam::Point3(lmk.x, lmk.y, lmk.z));
    left_kpts->push_back(KeypointCV(kp.uL(), kp.v()));
    right_kpts->push_back(KeypointCV(kp.uR(), kp.v()));
  }
}

void StereoCamera::backProject(const KeypointsCV& kps,
                               const cv::Mat& disparity_img,
                               LandmarksCV* lmks) const {
  CHECK_NOTNULL(lmks)->clear();
  CHECK_EQ(disparity_img.type(), CV_16S);  // 16-bit signed
  lmks->reserve(kps.size());
  for (const KeypointCV& kp : kps) {
    // Some algorithms, like StereoBM or StereoSGBM compute 16-bit fixed-point
    // disparity map (where each disparity value has 4 fractional bits),
    // whereas other algorithms output 32-bit floating-point disparity map.
    auto disparity = disparity_img.at<int16_t>(kp);
    // TODO(TONI): check disparity is valid!!
    CHECK_GE(disparity, 0);
    // WARNING gtsam's stereo points are doubles, our keypoints are floats,
    // and disparity is int16_t...
    gtsam::StereoPoint2 z(kp.x, kp.x + disparity, kp.y);
    gtsam::Point3 lmk = stereo_camera_impl_.backproject2(z);
    lmks->push_back(LandmarkCV(lmk.x(), lmk.y(), lmk.z()));
  }
}

void StereoCamera::stereoDisparityReconstruction(const cv::Mat& left_img,
                                                 const cv::Mat& right_img,
                                                 cv::Mat* disparity_img) {
  CHECK_NOTNULL(disparity_img);
  CHECK_EQ(disparity_img->cols, left_img.cols);
  CHECK_EQ(right_img.cols, left_img.cols);
  CHECK_EQ(disparity_img->rows, left_img.rows);
  CHECK_EQ(right_img.rows, left_img.rows);
  CHECK_EQ(right_img.type(), left_img.type());
  CHECK_EQ(disparity_img->type(), CV_32F);

  // Setup stereo matcher
  cv::Ptr<cv::StereoMatcher> stereo_matcher;
  if (use_sgbm_) {
    int mode;
    if (use_mode_HH_) {
      mode = cv::StereoSGBM::MODE_HH;
    } else {
      mode = cv::StereoSGBM::MODE_SGBM;
    }
    stereo_matcher = cv::StereoSGBM::create(min_disparity_,
                                            num_disparities_,
                                            sad_window_size_,
                                            p1_,
                                            p2_,
                                            disp_12_max_diff_,
                                            pre_filter_cap_,
                                            uniqueness_ratio_,
                                            speckle_window_size_,
                                            speckle_range_,
                                            mode);
  } else {
    cv::Ptr<cv::StereoBM> sbm =
        cv::StereoBM::create(num_disparities_, sad_window_size_);

    sbm->setPreFilterType(pre_filter_type_);
    sbm->setPreFilterSize(pre_filter_size_);
    sbm->setPreFilterCap(pre_filter_cap_);
    sbm->setMinDisparity(min_disparity_);
    sbm->setTextureThreshold(texture_threshold_);
    sbm->setUniquenessRatio(uniqueness_ratio_);
    sbm->setSpeckleRange(speckle_range_);
    sbm->setSpeckleWindowSize(speckle_window_size_);
    if (!ROI1_.empty() && !ROI2_.empty()) {
      sbm->setROI1(ROI1_);
      sbm->setROI2(ROI2_);
    } else {
      LOG(WARNING) << "ROIs are empty.";
    }

    stereo_matcher = sbm;
  }

  // Reconstruct scene
  stereo_matcher->compute(left_img, right_img, *disparity_img);

  // Optionally, post-filter disparity
  if (post_filter_disparity_) {
    // Use disparity post-filter
    // wls_filter = createDisparityWLSFilter(left_matcher);
    // Ptr<StereoMatcher> right_matcher = createRightMatcher(left_matcher);
    // See
    // https://docs.opencv.org/3.3.1/d3/d14/tutorial_ximgproc_disparity_filtering.html#gsc.tab=0
  }

  // Optionally, smooth the disparity image
  if (median_blur_disparity_) {
    cv::medianBlur(*disparity_img, *disparity_img, 5);
  }

  static constexpr bool debug = false;
  if (debug) {
    // cv::Mat raw_disp_vis;
    // cv::ximgproc::getDisparityVis(left_disp,raw_disp_vis,vis_mult);
    // cv::namedWindow("raw disparity", WINDOW_AUTOSIZE);
    // cv::imshow("raw disparity", raw_disp_vis);
    // cv::Mat filtered_disp_vis;
    // cv::ximgproc::getDisparityVis(filtered_disp,filtered_disp_vis,vis_mult);
    // cv::namedWindow("filtered disparity", WINDOW_AUTOSIZE);
    // cv::imshow("filtered disparity", filtered_disp_vis);
    // cv::waitKey();
  }
}

void StereoCamera::backProjectDisparityTo3D(const cv::Mat& disparity_img,
                                            cv::Mat* depth) {
  CHECK_NOTNULL(depth);
  // Check disparity img is CV_32F, if it is CV_16S it should have been
  // divided by 16, see StereoBM/stereoSGBM output.
  CHECK(!disparity_img.empty());
  CHECK_EQ(disparity_img.type(), CV_32F)
      << "Wrong type for disparity img, mind that if the disparity image is "
         "of type CV_16S, you might need to divide the values by 16.";
  // We use the Q_ pre-computed during undistortRectification, so check it has
  // been computed: in principle, we calculate that at ctor time, so it is
  // always computed.
  CHECK_NE(Q_.at<double>(3, 2), 0.0);
  static constexpr bool kHandleMissingValues = true;
  cv::reprojectImageTo3D(
      disparity_img, *depth, Q_, kHandleMissingValues, CV_32F);
}

void StereoCamera::undistortRectifyStereoFrame(StereoFrame* stereo_frame) {
  CHECK_NOTNULL(stereo_frame);
  LOG_IF(WARNING, stereo_frame->isRectified())
      << "Rectifying already rectified stereo frame ...";
  stereo_frame->setIsRectified(true);
  //! Rectify and undistort images using precomputed maps.
  CHECK(left_cam_undistort_rectifier_ != nullptr);
  left_cam_undistort_rectifier_->undistortRectifyImage(
      stereo_frame->getLeftFrameMutable()->img_,
      &stereo_frame->left_img_rectified_);
  CHECK(right_cam_undistort_rectifier_ != nullptr);
  right_cam_undistort_rectifier_->undistortRectifyImage(
      stereo_frame->getRightFrameMutable()->img_,
      &stereo_frame->right_img_rectified_);
}

void StereoCamera::computeRectificationParameters(
    const CameraParams& left_cam_params,
    const CameraParams& right_cam_params,
    cv::Mat* R1,
    cv::Mat* R2,
    cv::Mat* P1,
    cv::Mat* P2,
    cv::Mat* Q,
    cv::Rect* ROI1,
    cv::Rect* ROI2) {
  CHECK_NOTNULL(R1);
  CHECK_NOTNULL(R2);
  CHECK_NOTNULL(P1);
  CHECK_NOTNULL(P2);
  CHECK_NOTNULL(Q);
  CHECK_NOTNULL(ROI1);
  CHECK_NOTNULL(ROI2);

  //! Extrinsics of the stereo (not rectified) relative pose between cameras
  gtsam::Pose3 camL_Pose_camR =
      (left_cam_params.body_Pose_cam_).between(right_cam_params.body_Pose_cam_);

  // Get extrinsics in open CV format.
  // NOTE: openCV pose convention is the opposite, that's why we have to
  // invert
  cv::Mat camL_Rot_camR, camL_Tran_camR;
  boost::tie(camL_Rot_camR, camL_Tran_camR) =
      UtilsOpenCV::Pose2cvmats(camL_Pose_camR.inverse());

  // kAlpha is -1 by default, but that introduces invalid keypoints!
  // here we should use kAlpha = 0 so we get only valid pixels...
  // But that has an issue that it removes large part of the image, check:
  // https://github.com/opencv/opencv/issues/7240 for this issue with kAlpha
  // Setting to -1 to make it easy, but it should NOT be -1!
  static constexpr int kAlpha = -1;
  switch (left_cam_params.distortion_model_) {
    case DistortionModel::RADTAN: {
      cv::stereoRectify(
          // Input
          left_cam_params.K_,
          left_cam_params.distortion_coeff_mat_,
          right_cam_params.K_,
          right_cam_params.distortion_coeff_mat_,
          left_cam_params.image_size_,
          camL_Rot_camR,
          camL_Tran_camR,
          // Output
          *R1,
          *R2,
          *P1,
          *P2,
          *Q,
          cv::CALIB_ZERO_DISPARITY,
          kAlpha,
          cv::Size(),
          ROI1,
          ROI2);
    } break;
    case DistortionModel::EQUIDISTANT: {
      cv::fisheye::stereoRectify(
          // Input
          left_cam_params.K_,
          left_cam_params.distortion_coeff_mat_,
          right_cam_params.K_,
          right_cam_params.distortion_coeff_mat_,
          left_cam_params.image_size_,
          camL_Rot_camR,
          camL_Tran_camR,
          // Output
          *R1,
          *R2,
          *P1,
          *P2,
          *Q,
          // TODO: Flag to maximise area???
          cv::CALIB_ZERO_DISPARITY);
    } break;
    default: {
      LOG(FATAL) << "Unknown DistortionModel: "
                 << VIO::to_underlying(left_cam_params.distortion_model_);
    }
  }
}

}  // namespace VIO


#ifndef PYTHONIC_INCLUDE_NUMPY_FMAX_HPP
#define PYTHONIC_INCLUDE_NUMPY_FMAX_HPP

#include "pythonic/include/utils/functor.hpp"
#include "pythonic/include/types/ndarray.hpp"
#include "pythonic/include/utils/numpy_traits.hpp"

#include <boost/simd/function/max.hpp>

PYTHONIC_NS_BEGIN

namespace numpy
{
#define NUMPY_NARY_FUNC_NAME fmax
#define NUMPY_NARY_FUNC_SYM boost::simd::max
#include "pythonic/include/types/numpy_nary_expr.hpp"
}
PYTHONIC_NS_END

#endif


/*
 * Runs the target recognition algorithm on webcam feed or an image stream fed
 * from stdin.
 * Reports information on target position, orientation and scale as
 * comma-separated values (ASCII) in a specified UNIX socket.
 * This socket acts like a UDP one -- no buffering; if you don't read data when
 * it's available it won't keep it. Multiple programs can consume the data
 * (broadcast/multicast socket).
 */

#include <iostream>
#include <opencv2/opencv.hpp>
#include "TargetFinder_orig.h"
#include <boost/property_tree/ptree.hpp>
#include <boost/property_tree/ini_parser.hpp>

using namespace std;
using namespace cv;

void set_targetfinder_config(boost::property_tree::ptree pt) {
    #define pg(param) pt.get<double>("params." #param)
    CenterElipticalRatio = pg(center_elliptical_ratio);
    CenterDistancesRatio = pg(center_distances_ratio);
    MINDISTANCERATIO = pg(min_distance_ratio);
    MAXDISTANCERATIO = pg(max_distance_ratio);
    TargetSizeRatio = pg(target_size_ratio);
    MINTARGETSIZERATIO = pg(min_target_size_ratio);
    MAXTARGETSIZERATIO = pg(max_target_size_ratio);
    LRPairRatioThreshold = pg(lr_pair_ratio_threshold);
    pairCenterRatioThreshold = pg(pair_center_ratio_threshold);
    LRRatioThreshold = pg(lr_ratio_threshold);
    LRBlacktoCenterRatioThreshold = pg(lr_black_to_center_ratio_threshold);
    VALID_ROWS_TO_CENTER_RATIO = pg(valid_rows_to_center_ratio);
    CONTRAST = pg(contrast);
    DECAY_RATE = pg(decay_rate);
    SAMPLE_EVERY = pg(sample_every);
    UPDATE_EVERY = pg(update_every);
}

int main() {
    boost::property_tree::ptree pt;
    char* config_path = getenv("QUADTARGET_CONFIG");
    if(config_path == NULL) {
        boost::property_tree::ini_parser::read_ini("config.ini", pt);
    } else {
        boost::property_tree::ini_parser::read_ini(config_path, pt);
    }

    bool headless = pt.get<bool>("gui.headless");
    bool saveTmp = pt.get<bool>("gui.saveTmp");
    bool saveToFiles = pt.get<bool>("gui.saveToFiles");
    bool showProcessed = pt.get<bool>("gui.showProcessed");
    bool loop_video = false;
    int wait_key = pt.get<int>("gui.waitKey");
    string video_file = pt.get<string>("camera.file");

    VideoCapture cap;
    if(video_file.compare("none") == 0) {
        cap = VideoCapture(pt.get<int>("camera.index"));
    } else {
        cap = VideoCapture(video_file);
        loop_video = true;
    }
    if(!cap.isOpened()) {
        cerr << "Couldn't open webcam or video" << endl;
        return -1;
    }
    cap.set(CV_CAP_PROP_FRAME_WIDTH, pt.get<int>("camera.width"));
    cap.set(CV_CAP_PROP_FRAME_HEIGHT, pt.get<int>("camera.height"));
    cap.set(CV_CAP_PROP_FPS, pt.get<int>("camera.fps"));
    // cap.set(CV_CAP_PROP_CONVERT_RGB, false);

    Finder *tf;
    set_targetfinder_config(pt);
    if(pt.get<string>("tracker.mode").compare("cascade") == 0) {
        DEBUGPRINT("using cascade method");
        tf = new CascadeFinder(
                pt.get<string>("cascade.xml"),
                pt.get<double>("cascade.scaleFactor"),
                pt.get<int>("cascade.minNeighbors"),
                pt.get<int>("cascade.minSize")
        );
    } else {
        DEBUGPRINT("using original method");
        tf = new TargetFinder(headless);
    }

    int threshold_mode = 0;
    int threshold_value = pt.get<int>("threshold.value");
    int threshold_kernel = pt.get<int>("threshold.kernel_size");
    int threshold_c = pt.get<int>("threshold.c");
    if(pt.get<string>("threshold.mode").compare("otsu") == 0) {
        threshold_mode = 1;
    } else if(pt.get<string>("threshold.mode").compare("fixed") == 0) {
        threshold_mode = 2;
    } else if(pt.get<string>("threshold.mode").compare("mean") == 0) {
        threshold_mode = 3;
    } else if(pt.get<string>("threshold.mode").compare("gaussian") == 0) {
        threshold_mode = 4;
    }
    bool should_greyscale = pt.get<bool>("tracker.grey");

    bool running = true;
    Mat frame;
    // Mat yuv[3];
    Mat processed;
    Mat output;
    int count = 0;
    while(running) {
        DEBUGPRINT("count: " << count);
        if(cap.read(frame)) {
            if(should_greyscale) {
                cvtColor(frame, processed, CV_BGR2GRAY);
            } else {
                frame.copyTo(processed);
            }
            // split(frame, yuv);
            // processed = yuv[0];
            switch(threshold_mode) {
                case 1:
                    threshold(processed, processed, 128, 255, THRESH_OTSU);
                    break;
                case 2:
                    threshold(processed, processed, threshold_value, 255, THRESH_BINARY);
                    break;
                case 3:
                    adaptiveThreshold(processed, processed, 255,
                                      CV_ADAPTIVE_THRESH_MEAN_C, THRESH_BINARY,
                                      threshold_kernel, threshold_c);
                    break;
                case 4:
                    adaptiveThreshold(processed, processed, 255,
                                      CV_ADAPTIVE_THRESH_GAUSSIAN_C, THRESH_BINARY,
                                      threshold_kernel, threshold_c);
                    break;
                case 0:
                default:
                    break;
            }
            frame.copyTo(output);
            vector<FoundTarget> ft = tf->doTargetRecognition(&processed, &output);
            for(int i = 0; i < ft.size(); i++) {
                cout << "frame=" << count << ",t=" << i << ",";
                cout << ft.at(i).toString() << endl;
            }
            if(!headless) {
                imshow("input", frame);
                if(showProcessed) {
                    imshow("processed", processed);
                }
                imshow("output", output);
            }
            if(saveTmp) {
                imwrite("/tmp/target.png", frame);
            }
            if(saveToFiles) {
                imwrite("input.jpg", frame);
                if(showProcessed) {
                    imwrite("processed.jpg", processed);
                }
                imwrite("output.jpg", output);
            }
            count++;
        } else {
            DEBUGPRINT("couldn't get webcam frame? Or we've reached end of video.");
            if(loop_video) {
                cap.open(video_file);
            }
        }
        if(!headless) {
            running = !(waitKey(wait_key) == 'q');
        }
    }

    return 0;
}