#define BOOST_TEST_MODULE pcraster dal data_space_test
#include <boost/test/unit_test.hpp>
#include "dal_DataSpace.h"
#include "dal_DataSpaceAddress.h"
#include "dal_Dimension.h"
#include "dal_RasterDimensions.h"

BOOST_AUTO_TEST_CASE(test_)
{
  using namespace dal;

  {
    // scenario / time / space / space

    std::set<std::string> scenarios;
    std::vector<size_t> timeSteps;
    scenarios.insert("aap");
    timeSteps.push_back(1);
    timeSteps.push_back(100);
    timeSteps.push_back(1);

    RasterDimensions rasterDimensions(100, 100);

    DataSpace space;
    space.addDimension(Dimension(Scenarios, scenarios));
    space.addDimension(Dimension(Time, timeSteps));
    space.addDimension(Dimension(Space, RegularDiscretisation,
         rasterDimensions));
    BOOST_CHECK_EQUAL(space.rank(), size_t(3));
    BOOST_CHECK(space.isSpatial());
    BOOST_CHECK(space.hasRaster());
    BOOST_CHECK(space.hasTime());
    BOOST_CHECK_EQUAL(space.nrWideDimensions(), size_t(2));

    DataSpaceAddress address = space.address();
    address.setCoordinate<std::string>(0, "aap");
    address.setCoordinate<size_t>(1, 50);
    address.setCoordinate<SpatialCoordinate>(2, SpatialCoordinate(44.4, -55.5));

    DataSpace space2(space, address);
    BOOST_CHECK_EQUAL(space2.size(), space.size());
    BOOST_CHECK_EQUAL(space2.nrWideDimensions(), size_t(0));

    for(size_t i = 0; i < space2.size(); ++i) {
      BOOST_CHECK_EQUAL(space2.dimension(i).coordinateType(),
         space.dimension(i).coordinateType());
      BOOST_CHECK_EQUAL(space2.dimension(i).meaning(),
         space.dimension(i).meaning());
      BOOST_CHECK_EQUAL(space2.dimension(i).discretisation(),
         space.dimension(i).discretisation());
    }

    BOOST_CHECK_EQUAL(space2.dimension(0).nrValues(), size_t(1));
    BOOST_CHECK_EQUAL(space2.dimension(0).value<std::string>(0), "aap");
    BOOST_CHECK_EQUAL(space2.dimension(1).nrValues(), size_t(3));
    BOOST_CHECK_EQUAL(space2.dimension(1).value<size_t>(0), size_t(50));
    BOOST_CHECK_EQUAL(space2.dimension(1).value<size_t>(1), size_t(50));
    BOOST_CHECK_EQUAL(space2.dimension(1).value<size_t>(2), size_t(1));
    BOOST_CHECK(space2.dimension(2).value<RasterDimensions>(0) ==
         RasterDimensions(1, 1, 1.0, 44.0, -55.0));
  }
}

BOOST_AUTO_TEST_CASE(merge)
{
  using namespace dal;

  std::set<std::string> scenarios;
  scenarios.insert("aap");

  RasterDimensions rasterDimensions(100, 60);

  Dimension scenariosDimension(Scenarios, scenarios);
  Dimension samplesDimension(Samples, size_t(1), size_t(1000), size_t(1));
  Dimension cummulativeProbabilitiesDimension(CumulativeProbabilities,
         float(0.01), float(0.99), float(0.01));
  Dimension timeDimension(Time, size_t(1), size_t(1000), size_t(1));
  Dimension spaceDimension(Space, RegularDiscretisation, rasterDimensions);

  {
    DataSpace space1, space2, space3;

    space1.addDimension(scenariosDimension);
    space1.addDimension(timeDimension);
    space1.addDimension(spaceDimension);

    BOOST_CHECK(space2.isEmpty());
    BOOST_CHECK(space3.isEmpty());

    space2 |= space1;
    BOOST_CHECK_EQUAL(space2.rank(), size_t(3));
    BOOST_CHECK(space2.hasScenarios());
    BOOST_CHECK(space2.hasTime());
    BOOST_CHECK(space2.isSpatial());
    BOOST_CHECK(space2.hasRaster());
    BOOST_CHECK(space2 == space1);

    // Again.
    space2 |= space1;
    BOOST_CHECK_EQUAL(space2.rank(), size_t(3));
    BOOST_CHECK(space2.hasScenarios());
    BOOST_CHECK(space2.hasTime());
    BOOST_CHECK(space2.isSpatial());
    BOOST_CHECK(space2.hasRaster());
    BOOST_CHECK(space2 == space1);

    // Other way around.
    space1 |= space2;
    BOOST_CHECK_EQUAL(space1.rank(), size_t(3));
    BOOST_CHECK(space1.hasScenarios());
    BOOST_CHECK(space1.hasTime());
    BOOST_CHECK(space1.isSpatial());
    BOOST_CHECK(space1.hasRaster());
    BOOST_CHECK(space1 == space2);

    // Merge with empty space.
    space1 |= space3;
    BOOST_CHECK(space3.isEmpty());
    BOOST_CHECK_EQUAL(space1.rank(), size_t(3));
    BOOST_CHECK(space1.hasScenarios());
    BOOST_CHECK(space1.hasTime());
    BOOST_CHECK(space1.isSpatial());
    BOOST_CHECK(space1.hasRaster());
    BOOST_CHECK(space1 == space2);
  }

  // Without Scenarios, Samples and CumulativeProbabilities.
  // - Time (time series)
  // - Space (raster)
  // - Time/Space (dynamic raster: stack)
  {
    DataSpace space1, space2;

    space1.addDimension(timeDimension);
    space2.addDimension(spaceDimension);

    space1 |= space2;
    BOOST_CHECK_EQUAL(space1.dimension(0).meaning(), Time);
    BOOST_CHECK_EQUAL(space1.dimension(1).meaning(), Space);
    BOOST_CHECK(space1.dimension(1).value<RasterDimensions>(0) ==
         rasterDimensions);

    // The other way around.
    space1.clear();
    space1.addDimension(timeDimension);

    space2 |= space1;
    BOOST_CHECK_EQUAL(space2.dimension(0).meaning(), Time);
    BOOST_CHECK_EQUAL(space2.dimension(1).meaning(), Space);
    BOOST_CHECK(space2.dimension(1).value<RasterDimensions>(0) ==
         rasterDimensions);
  }

  // With Scenarios.
  // - Scenarios/Space (for each scenario a raster)
  // - Scenarios/Time (for each scenario a time series)
  // - Scenarios/Time/Space (for each scenario a stack)
  {
    DataSpace space1, space2;

    space1.addDimension(scenariosDimension);
    space2.addDimension(timeDimension);

    space1 |= space2;
    BOOST_CHECK_EQUAL(space1.dimension(0).meaning(), Scenarios);
    BOOST_CHECK_EQUAL(space1.dimension(1).meaning(), Time);

    space1.clear();
    space1.addDimension(scenariosDimension);

    space2 |= space1;
    BOOST_CHECK_EQUAL(space2.dimension(0).meaning(), Scenarios);
    BOOST_CHECK_EQUAL(space2.dimension(1).meaning(), Time);
  }

  // With Samples.
  // - Scenarios/Samples/Space
  // - Scenarios/Samples/Time
  // - Scenarios/Samples/Time/Space
  {
    DataSpace space1, space2, space3;

    space1.addDimension(scenariosDimension);
    space2.addDimension(samplesDimension);
    space3.addDimension(timeDimension);

    space1 |= space2;
    BOOST_CHECK_EQUAL(space1.dimension(0).meaning(), Scenarios);
    BOOST_CHECK_EQUAL(space1.dimension(1).meaning(), Samples);

    space1 |= space3;
    BOOST_CHECK_EQUAL(space1.dimension(0).meaning(), Scenarios);
    BOOST_CHECK_EQUAL(space1.dimension(1).meaning(), Samples);
    BOOST_CHECK_EQUAL(space1.dimension(2).meaning(), Time);

    space1.clear();
    space1.addDimension(scenariosDimension);

    space2 |= space1;
    BOOST_CHECK_EQUAL(space2.dimension(0).meaning(), Scenarios);
    BOOST_CHECK_EQUAL(space2.dimension(1).meaning(), Samples);

    space2 |= space3;
    BOOST_CHECK_EQUAL(space2.dimension(0).meaning(), Scenarios);
    BOOST_CHECK_EQUAL(space2.dimension(1).meaning(), Samples);
    BOOST_CHECK_EQUAL(space2.dimension(2).meaning(), Time);

    space2.clear();
    space2.addDimension(samplesDimension);

    space3 |= space1;
    BOOST_CHECK_EQUAL(space3.dimension(0).meaning(), Scenarios);
    BOOST_CHECK_EQUAL(space3.dimension(1).meaning(), Time);

    space3 |= space2;
    BOOST_CHECK_EQUAL(space3.dimension(0).meaning(), Scenarios);
    BOOST_CHECK_EQUAL(space3.dimension(1).meaning(), Samples);
    BOOST_CHECK_EQUAL(space3.dimension(2).meaning(), Time);
  }

  // With CumulativeProbabilities.
  // - Scenarios/CumulativeProbabilities/Space
  // - Scenarios/Time/CumulativeProbabilities
  // - Scenarios/Time/CumulativeProbabilities/Space
  {
    DataSpace space1, space2, space3;

    space1.addDimension(scenariosDimension);
    space2.addDimension(timeDimension);
    space3.addDimension(cummulativeProbabilitiesDimension);

    space1 |= space2;
    BOOST_CHECK_EQUAL(space1.dimension(0).meaning(), Scenarios);
    BOOST_CHECK_EQUAL(space1.dimension(1).meaning(), Time);

    space1 |= space3;
    BOOST_CHECK_EQUAL(space1.dimension(0).meaning(), Scenarios);
    BOOST_CHECK_EQUAL(space1.dimension(1).meaning(), Time);
    BOOST_CHECK_EQUAL(space1.dimension(2).meaning(), CumulativeProbabilities);

    space1.clear();
    space1.addDimension(scenariosDimension);

    space2 |= space1;
    BOOST_CHECK_EQUAL(space2.dimension(0).meaning(), Scenarios);
    BOOST_CHECK_EQUAL(space2.dimension(1).meaning(), Time);

    space2 |= space3;
    BOOST_CHECK_EQUAL(space2.dimension(0).meaning(), Scenarios);
    BOOST_CHECK_EQUAL(space2.dimension(1).meaning(), Time);
    BOOST_CHECK_EQUAL(space2.dimension(2).meaning(), CumulativeProbabilities);

    space2.clear();
    space2.addDimension(timeDimension);

    space3 |= space1;
    BOOST_CHECK_EQUAL(space3.dimension(0).meaning(), Scenarios);
    BOOST_CHECK_EQUAL(space3.dimension(1).meaning(), CumulativeProbabilities);

    space3 |= space2;
    BOOST_CHECK_EQUAL(space3.dimension(0).meaning(), Scenarios);
    BOOST_CHECK_EQUAL(space3.dimension(1).meaning(), Time);
    BOOST_CHECK_EQUAL(space3.dimension(2).meaning(), CumulativeProbabilities);
  }

  {
    // Because of bug (only in release mode, see arrow below):
    // space1: /scenarios{0}
    //         /time[1986, 2003, 1]
    //         /cumulative probabilities[0.01, 0.99, 0.01]
    //         /space[1, 87, 1]/[1, 55, 1]
    // space2: /scenarios{5000}
    //         /time[1986, 2003, 1]
    //         /cumulative probabilities[0.01, 0.99, 0.01]
    //         /space[1, 87, 1]/[1, 55, 1]
    // result: /scenarios{0, 5000}
    //         /time[1986, 2003, 1]
    //         /cumulative probabilities[0.01, 0.99, -2.14748e-07]   <---------
    //         /space[1, 87, 1]/[1, 55, 1]
    DataSpace space1, space2;

    std::set<std::string> scenarios;
    std::vector<size_t> timeSteps;
    std::vector<float>quantiles;
    scenarios.insert("0");
    timeSteps.push_back(1986);
    timeSteps.push_back(2003);
    timeSteps.push_back(1);
    quantiles.push_back(0.01f);
    quantiles.push_back(0.99f);
    quantiles.push_back(0.01f);

    RasterDimensions rasterDimensions(87, 55);

    space1.addDimension(Dimension(Scenarios, scenarios));
    space1.addDimension(Dimension(Time, timeSteps));
    space1.addDimension(Dimension(CumulativeProbabilities, quantiles));
    space1.addDimension(Dimension(Space, RegularDiscretisation,
         rasterDimensions));

    scenarios.clear();
    scenarios.insert("5000");
    space2.addDimension(Dimension(Scenarios, scenarios));
    space2.addDimension(Dimension(Time, timeSteps));
    space2.addDimension(Dimension(CumulativeProbabilities, quantiles));
    space2.addDimension(Dimension(Space, RegularDiscretisation,
         rasterDimensions));

    space1 |= space2;
    BOOST_CHECK_EQUAL(space1.size(), size_t(4));
    BOOST_CHECK_EQUAL(space1.dimension(2).nrValues(), size_t(3));
    BOOST_CHECK(comparable<float>(space1.dimension(2).value<float>(2), 0.01f));
  }

  // Merge space with different discretizations (raster and vector).
  // Merging a raster and a vector data space should result in a data space
  // with info about both types of dimensions. Raster info must come before
  // vector info.
  {
    DataSpace resultSpace, rasterSpace, vectorSpace;

    RasterDimensions rasterDimensions(87, 55);
    rasterSpace.addDimension(Dimension(Space, RegularDiscretisation,
         rasterDimensions));

    SpaceDimensions featureDimensions(555.5, 222.2, 666.6, -111.1);
    vectorSpace.addDimension(Dimension(Space, BorderedDiscretisation,
         featureDimensions));

    {
      resultSpace = rasterSpace | vectorSpace;
      BOOST_CHECK_EQUAL(resultSpace.size(), size_t(2));
      BOOST_CHECK(resultSpace.dimension(0).value<RasterDimensions>(0) ==
         rasterDimensions);
      BOOST_CHECK(resultSpace.dimension(1).value<SpaceDimensions>(0) ==
         featureDimensions);
    }

    {
      resultSpace = vectorSpace | rasterSpace;
      BOOST_CHECK_EQUAL(resultSpace.size(), size_t(2));
      BOOST_CHECK(resultSpace.dimension(0).value<RasterDimensions>(0) ==
         rasterDimensions);
      BOOST_CHECK(resultSpace.dimension(1).value<SpaceDimensions>(0) ==
         featureDimensions);
    }
  }

  // Real world example than failed.
  // space1: scenarios/time/probabilities/raster/features
  // space2: scenarios/time/probabilities/features
  // result: scenarios/time/probabilities/raster/features
  {
    DataSpace resultSpace, space1, space2;

    std::set<std::string> scenarios;
    scenarios.insert("1000");
    SpaceDimensions featureDimensions(464260, 6.18012e+06, 791405, 5.68706e+06);
    RasterDimensions rasterDimensions(87, 55, 5000.0, 464000.0, 6131500.0);

    space2.addDimension(Dimension(Scenarios, scenarios));
    space2.addDimension(Dimension(Time, size_t(1986), size_t(2003), size_t(1)));
    space2.addDimension(Dimension(CumulativeProbabilities, float(0.01),
         float(0.99), float(0.01)));
    space2.addDimension(Dimension(Space, BorderedDiscretisation,
         featureDimensions));

    scenarios.insert("0");
    scenarios.insert("5000");
    scenarios.insert("25000");

    space1.addDimension(Dimension(Scenarios, scenarios));
    space1.addDimension(Dimension(Time, size_t(1986), size_t(2003), size_t(1)));
    space1.addDimension(Dimension(CumulativeProbabilities, float(0.01),
         float(0.99), float(0.01)));
    space1.addDimension(Dimension(Space, RegularDiscretisation,
         rasterDimensions));
    space1.addDimension(Dimension(Space, BorderedDiscretisation,
         featureDimensions));

    resultSpace = space1 | space2;

    BOOST_CHECK_EQUAL(resultSpace.size(), size_t(5));
    BOOST_CHECK_EQUAL(resultSpace.dimension(0).meaning(), Scenarios);
    BOOST_CHECK_EQUAL(resultSpace.dimension(1).meaning(), Time);
    BOOST_CHECK_EQUAL(resultSpace.dimension(2).meaning(),
         CumulativeProbabilities);
    BOOST_CHECK_EQUAL(resultSpace.dimension(3).meaning(), Space);
    BOOST_CHECK_EQUAL(resultSpace.dimension(3).discretisation(),
         RegularDiscretisation);
    BOOST_CHECK_EQUAL(resultSpace.dimension(4).meaning(), Space);
    BOOST_CHECK_EQUAL(resultSpace.dimension(4).discretisation(),
         BorderedDiscretisation);

    // And the other way around.
    resultSpace = space2 | space1;

    BOOST_CHECK_EQUAL(resultSpace.size(), size_t(5));
    BOOST_CHECK_EQUAL(resultSpace.dimension(0).meaning(), Scenarios);
    BOOST_CHECK_EQUAL(resultSpace.dimension(1).meaning(), Time);
    BOOST_CHECK_EQUAL(resultSpace.dimension(2).meaning(),
         CumulativeProbabilities);
    BOOST_CHECK_EQUAL(resultSpace.dimension(3).meaning(), Space);
    BOOST_CHECK_EQUAL(resultSpace.dimension(3).discretisation(),
         RegularDiscretisation);
    BOOST_CHECK_EQUAL(resultSpace.dimension(4).meaning(), Space);
    BOOST_CHECK_EQUAL(resultSpace.dimension(4).discretisation(),
         BorderedDiscretisation);
  }
}

BOOST_AUTO_TEST_CASE(intersect)
{
  using namespace dal;

  std::set<std::string> scenarios;
  scenarios.insert("aap");
  scenarios.insert("noot");
  scenarios.insert("mies");

  std::vector<size_t> timeSteps;
  timeSteps.push_back(size_t(1));
  timeSteps.push_back(size_t(100));
  timeSteps.push_back(size_t(1));

  RasterDimensions rasterDimensions(100, 100);

  // Result of intersecting an empty space with whatever other space is
  // the empty space.
  {
    DataSpace space1;
    DataSpace space2;
    space1 &= space2;
    BOOST_CHECK(space1.isEmpty());

    space2.addDimension(Dimension(Scenarios, scenarios));
    space2.addDimension(Dimension(Time, timeSteps));
    space2.addDimension(Dimension(Space, RegularDiscretisation,
         rasterDimensions));

    space1 &= space2;
    BOOST_CHECK(space1.isEmpty());
  }

  // Result of intersecting whatever space with an empty space is the
  // empty space.
  {
    DataSpace space1;
    DataSpace space2;

    space1.addDimension(Dimension(Scenarios, scenarios));
    space1.addDimension(Dimension(Time, timeSteps));
    space1.addDimension(Dimension(Space, RegularDiscretisation,
         rasterDimensions));

    space1 &= space2;
    BOOST_CHECK(space1.isEmpty());
  }

  // Result of intersecting two spaces with different dimensions is an
  // empty space.
  {
    DataSpace space1;
    space1.addDimension(Dimension(Scenarios, scenarios));

    DataSpace space2;
    space2.addDimension(Dimension(Time, timeSteps));

    space1 &= space2;
    BOOST_CHECK(space1.isEmpty());
  }

  // Result of intersecting two space with the same dimensions is this
  // set of equal dimensions.
  {
    DataSpace space1;
    DataSpace space2;

    space1.addDimension(Dimension(Scenarios, scenarios));
    space1.addDimension(Dimension(Time, timeSteps));
    space1.addDimension(Dimension(Space, RegularDiscretisation,
         rasterDimensions));

    space2.addDimension(Dimension(Scenarios, scenarios));

    space1 &= space2;
    BOOST_CHECK_EQUAL(space1.size(), size_t(1));
    BOOST_CHECK_EQUAL(space1.dimension(0).meaning(), Scenarios);
  }

  {
    DataSpace space1;
    DataSpace space2;

    space1.addDimension(Dimension(Scenarios, scenarios));
    space1.addDimension(Dimension(Time, timeSteps));
    space1.addDimension(Dimension(Space, RegularDiscretisation,
         rasterDimensions));

    space2.addDimension(Dimension(Time, timeSteps));

    space1 &= space2;
    BOOST_CHECK_EQUAL(space1.size(), size_t(1));
    BOOST_CHECK_EQUAL(space1.dimension(0).meaning(), Time);
  }

  {
    DataSpace space1;
    DataSpace space2;

    space1.addDimension(Dimension(Scenarios, scenarios));
    space1.addDimension(Dimension(Time, timeSteps));
    space1.addDimension(Dimension(Space, RegularDiscretisation,
         rasterDimensions));

    space2.addDimension(Dimension(Space, RegularDiscretisation,
         rasterDimensions));

    space1 &= space2;
    BOOST_CHECK_EQUAL(space1.size(), size_t(1));
    BOOST_CHECK_EQUAL(space1.dimension(0).meaning(), Space);
  }

  //
  {
  }
}

BOOST_AUTO_TEST_CASE(trim)
{
  using namespace dal;

  std::set<std::string> scenarios;
  std::vector<size_t> timeSteps;
  scenarios.insert("aap");
  timeSteps.push_back(1);
  timeSteps.push_back(100);
  timeSteps.push_back(1);

  RasterDimensions rasterDimensions(100, 100);

  DataSpace appSpace;
  appSpace.addDimension(Dimension(Scenarios, scenarios));
  appSpace.addDimension(Dimension(Time, timeSteps));
  appSpace.addDimension(Dimension(Space, RegularDiscretisation,
         rasterDimensions));

  {
    // Trim an address in scenario / time / space / space to
    // address in space / space.

    RasterDimensions rasterDimensions(10, 55);

    DataSpace dataSpace;
    dataSpace.addDimension(Dimension(Space, RegularDiscretisation,
         rasterDimensions));

    DataSpaceAddress address(3);
    address.setCoordinate<std::string>(0, "aap");
    address.setCoordinate<size_t>(1, 5);
    address.setCoordinate<SpatialCoordinate>(2, SpatialCoordinate(5.5, 6.6));

    address = dataSpace.trim(appSpace, address);
    BOOST_CHECK_EQUAL(address.size(), size_t(1));
    BOOST_CHECK(address.coordinate<SpatialCoordinate>(0) ==
         SpatialCoordinate(5.5, 6.6));
  }

  {
    // Trim an address in scenario / time / space / space to
    // address in scenario / time.

    std::set<std::string> scenarios;
    std::vector<size_t> timeSteps;
    scenarios.insert("aap");
    timeSteps.push_back(5);
    timeSteps.push_back(10);
    timeSteps.push_back(1);

    DataSpace dataSpace;
    dataSpace.addDimension(Dimension(Scenarios, scenarios));
    dataSpace.addDimension(Dimension(Time, timeSteps));

    DataSpaceAddress address(3);
    address.setCoordinate<std::string>(0, "aap");
    address.setCoordinate<size_t>(1, 5);
    address.setCoordinate<SpatialCoordinate>(2, SpatialCoordinate(5.5, 6.6));

    address = dataSpace.trim(appSpace, address);
    BOOST_CHECK_EQUAL(address.size(), size_t(2));
    BOOST_CHECK_EQUAL(address.coordinate<std::string>(0), "aap");
    BOOST_CHECK_EQUAL(address.coordinate<size_t>(1), size_t(5));
  }

  {
    // Trim an address in scenario / time / space / space to
    // address in time.

    std::vector<size_t> timeSteps;
    timeSteps.push_back(9);
    timeSteps.push_back(20);
    timeSteps.push_back(1);

    DataSpace dataSpace;
    dataSpace.addDimension(Dimension(Time, timeSteps));

    DataSpaceAddress address(3);
    address.setCoordinate<std::string>(0, "aap");
    address.setCoordinate<size_t>(1, 5);
    address.setCoordinate<SpatialCoordinate>(2, SpatialCoordinate(5.5, 6.6));

    address = dataSpace.trim(appSpace, address);
    BOOST_CHECK_EQUAL(address.size(), size_t(1));
    BOOST_CHECK_EQUAL(address.coordinate<size_t>(0), size_t(5));
  }

  {
    // Trim an address to a space with less dimensions.

    std::vector<size_t> timeSteps, cells;
    timeSteps.push_back(9);
    timeSteps.push_back(11);
    timeSteps.push_back(1);

    RasterDimensions rasterDimensions(10, 10);

    DataSpace fromSpace, toSpace;
    fromSpace.addDimension(Dimension(Time, timeSteps));
    fromSpace.addDimension(Dimension(Space, RegularDiscretisation,
         rasterDimensions));
    toSpace = fromSpace;
    toSpace.eraseDimension(Space);

    DataSpaceAddress address(2);
    address.setCoordinate<size_t>(0, 9);
    address.setCoordinate<SpatialCoordinate>(1, SpatialCoordinate(5.5, 6.6));
    address = toSpace.trim(fromSpace, address);
    BOOST_CHECK_EQUAL(address.size(), size_t(1));

    BOOST_CHECK(address.isValid(0));
    BOOST_CHECK_EQUAL(address.coordinate<size_t>(0), size_t(9));
  }

  {
    // NOTE: Tests updated to new approach to spatial dimensions. Comment
    // NOTE: left, but refers to previous approach.

    // A test from an Aguila use case.
    // space:
    //   /scenarios{0, 10000, 25000, 5000}
    //     /time[1986, 2003, 1]
    //       /cumulative probabilities[0.02, 0.98, 0.01]
    //         /space[1, 87, 1]
    //         /space[1, 55, 1]
    // address: (0, 1986, 0.5, 1, 1)
    //
    // trim to:
    // space:
    //   /space[1, 87, 1]
    //   /space[1, 55, 1]
    //
    // must result in:
    // address: (1, 1)

    std::set<std::string> scenarios;
    scenarios.insert("0");
    scenarios.insert("100000");
    scenarios.insert("25000");
    scenarios.insert("5000");

    std::vector<size_t> timeSteps;
    timeSteps.push_back(1986);
    timeSteps.push_back(2003);
    timeSteps.push_back(1);

    std::vector<float>quantiles;
    quantiles.push_back(0.01f);
    quantiles.push_back(0.99f);
    quantiles.push_back(0.01f);

    RasterDimensions rasterDimensions(87, 55);

    DataSpace fromSpace;
    fromSpace.addDimension(Dimension(Scenarios, scenarios));
    fromSpace.addDimension(Dimension(Time, timeSteps));
    fromSpace.addDimension(Dimension(CumulativeProbabilities, quantiles));
    fromSpace.addDimension(Dimension(Space, RegularDiscretisation,
         rasterDimensions));

    DataSpace toSpace;
    toSpace.addDimension(Dimension(Space, RegularDiscretisation,
         rasterDimensions));

    // address: (0, 1986, 0.5, (5.5, 6.6))
    DataSpaceAddress address(4);
    address.setCoordinate<std::string>(0, "0");
    address.setCoordinate<size_t>(1, 1986);
    address.setCoordinate<float>(2, 0.5f);
    address.setCoordinate<SpatialCoordinate>(3, SpatialCoordinate(5.5, 6.6));
    address = toSpace.trim(fromSpace, address);

    BOOST_REQUIRE(toSpace.isValid(address));
    BOOST_CHECK_EQUAL(address.size(), size_t(1));
    BOOST_REQUIRE(address.isValid(0));
    BOOST_CHECK(address.coordinate<SpatialCoordinate>(0) ==
         SpatialCoordinate(5.5, 6.6));

    // Additional space dimensions.
    SpaceDimensions featureDimension(5.68706e+06, 791405.0, 6.18012e+06,
         464260.0);

    // Same source space, add these:
    //   space[464260, 791405]
    //   space[5.68706e+06, 6.18012e+06]
    fromSpace.addDimension(Dimension(Space, BorderedDiscretisation,
         featureDimension));

    toSpace.clear();
    toSpace.addDimension(Dimension(Space, BorderedDiscretisation,
         featureDimension));

    // trim to: /space[464260, 791405]/space[5.68706e+06, 6.18012e+06]
    // address: (0, 1986, 0.5, (1, 1), (464260, 5687057))
    address.resize(5);
    address.setCoordinate<std::string>(0, "0");
    address.setCoordinate<size_t>(1, 1986);
    address.setCoordinate<float>(2, 0.5f);
    address.setCoordinate<SpatialCoordinate>(3, SpatialCoordinate(464260.0,
         5687057.0));
    address.setCoordinate<SpatialCoordinate>(4, SpatialCoordinate(464260.0,
         5687057.0));

    address = toSpace.trim(fromSpace, address);

    BOOST_REQUIRE(toSpace.isValid(address));
    BOOST_CHECK_EQUAL(address.size(), size_t(1));
    BOOST_REQUIRE(address.isValid(0));
    BOOST_CHECK(address.coordinate<SpatialCoordinate>(0) ==
         SpatialCoordinate(464260.0, 5687057.0));
  }
}

BOOST_AUTO_TEST_CASE(initialise_invalid_coordinates)
{
  using namespace dal;

  DataSpace space;
  DataSpaceAddress address;

  {
    // Empty space.
    address = space.address();
    BOOST_CHECK(space.isValid(address));
    address = space.initialiseInvalidCoordinates(address);
    BOOST_CHECK(space.isValid(address));
    BOOST_CHECK_EQUAL(address.size(), size_t(0));
  }

  std::set<std::string> scenarios;
  scenarios.insert("aap");
  scenarios.insert("noot");
  scenarios.insert("mies");
  space.addDimension(Dimension(Scenarios, scenarios));

  {
    // One scenario dimension added.
    address = space.address();
    BOOST_CHECK_EQUAL(address.size(), size_t(1));
    BOOST_CHECK(!space.isValid(address));
    BOOST_CHECK(!address.isValid(0));

    address = space.initialiseInvalidCoordinates(address);
    BOOST_CHECK_EQUAL(address.size(), size_t(1));
    BOOST_CHECK(space.isValid(address));
    BOOST_CHECK(address.isValid(0));
    BOOST_CHECK_EQUAL(address.coordinate<std::string>(0), "aap");
  }

  std::vector<size_t> samples;
  samples.push_back(3);
  samples.push_back(9);
  samples.push_back(2);
  space.addDimension(Dimension(Samples, samples));

  {
    // One samples dimension added.
    address = space.address();
    BOOST_CHECK_EQUAL(address.size(), size_t(2));
    BOOST_CHECK(!space.isValid(address));
    BOOST_CHECK(!address.isValid(0));
    BOOST_CHECK(!address.isValid(1));

    address = space.initialiseInvalidCoordinates(address);
    BOOST_CHECK_EQUAL(address.size(), size_t(2));
    BOOST_CHECK(space.isValid(address));
    BOOST_CHECK(address.isValid(0));
    BOOST_CHECK(address.isValid(1));
    BOOST_CHECK_EQUAL(address.coordinate<std::string>(0), "aap");
    BOOST_CHECK_EQUAL(address.coordinate<size_t>(1), size_t(3));
  }

  std::vector<size_t> timeSteps;
  timeSteps.push_back(1);
  timeSteps.push_back(100);
  timeSteps.push_back(1);
  space.addDimension(Dimension(Time, timeSteps));

  {
    // One time dimension added.
    address = space.address();
    BOOST_CHECK_EQUAL(address.size(), size_t(3));
    BOOST_CHECK(!space.isValid(address));
    BOOST_CHECK(!address.isValid(0));
    BOOST_CHECK(!address.isValid(1));
    BOOST_CHECK(!address.isValid(2));

    address = space.initialiseInvalidCoordinates(address);
    BOOST_CHECK_EQUAL(address.size(), size_t(3));
    BOOST_CHECK(space.isValid(address));
    BOOST_CHECK(address.isValid(0));
    BOOST_CHECK(address.isValid(1));
    BOOST_CHECK(address.isValid(2));
    BOOST_CHECK_EQUAL(address.coordinate<std::string>(0), "aap");
    BOOST_CHECK_EQUAL(address.coordinate<size_t>(1), size_t(3));
    BOOST_CHECK_EQUAL(address.coordinate<size_t>(2), size_t(1));
  }

  space.eraseDimension(Samples);

  std::vector<float> quantiles;
  quantiles.push_back(0.01f);
  quantiles.push_back(0.99f);
  quantiles.push_back(0.01f);
  space.addDimension(Dimension(CumulativeProbabilities, quantiles));

  {
    // One cumulative probability dimension added.
    address = space.address();
    BOOST_CHECK_EQUAL(address.size(), size_t(3));
    BOOST_CHECK(!space.isValid(address));
    BOOST_CHECK(!address.isValid(0));
    BOOST_CHECK(!address.isValid(1));
    BOOST_CHECK(!address.isValid(2));

    address = space.initialiseInvalidCoordinates(address);
    BOOST_CHECK_EQUAL(address.size(), size_t(3));
    BOOST_CHECK(space.isValid(address));
    BOOST_CHECK(address.isValid(0));
    BOOST_CHECK(address.isValid(1));
    BOOST_CHECK(address.isValid(2));
    BOOST_CHECK_EQUAL(address.coordinate<std::string>(0), "aap");
    BOOST_CHECK_EQUAL(address.coordinate<size_t>(1), size_t(1));
    BOOST_CHECK_CLOSE(address.coordinate<float>(2), float(0.50), 0.001f);
  }

  RasterDimensions rasterDimensions(96, 95, 3.0, 5.5, 6.6);
  space.addDimension(Dimension(Space, RegularDiscretisation, rasterDimensions));

  {
    // One space dimension added.
    address = space.address();
    BOOST_CHECK_EQUAL(address.size(), size_t(4));
    BOOST_CHECK(!space.isValid(address));
    BOOST_CHECK(!address.isValid(0));
    BOOST_CHECK(!address.isValid(1));
    BOOST_CHECK(!address.isValid(2));
    BOOST_CHECK(!address.isValid(3));

    address = space.initialiseInvalidCoordinates(address);
    BOOST_CHECK_EQUAL(address.size(), size_t(4));
    BOOST_CHECK(space.isValid(address));
    BOOST_CHECK(address.isValid(0));
    BOOST_CHECK(address.isValid(1));
    BOOST_CHECK(address.isValid(2));
    BOOST_CHECK(address.isValid(3));
    BOOST_CHECK_EQUAL(address.coordinate<std::string>(0), "aap");
    BOOST_CHECK_EQUAL(address.coordinate<size_t>(1), size_t(1));
    BOOST_CHECK_CLOSE(address.coordinate<float>(2), float(0.50), 0.001f);
    BOOST_CHECK(address.coordinate<SpatialCoordinate>(3) ==
         SpatialCoordinate(5.5, 6.6));
  }

  {
    DataSpace space;
    DataSpaceAddress address;
    std::vector<float> quantiles;
    quantiles.push_back(0.90f);
    quantiles.push_back(0.95f);
    quantiles.push_back(0.99f);
    space.addDimension(Dimension(CumulativeProbabilities, quantiles));

    address = space.address();
    address = space.initialiseInvalidCoordinates(address);
    BOOST_CHECK_EQUAL(address.size(), size_t(1));
    BOOST_CHECK(space.isValid(address));
    BOOST_CHECK(address.isValid(0));
    BOOST_CHECK_EQUAL(address.coordinate<float>(0), float(0.90));
  }

  {
    DataSpace space;
    DataSpaceAddress address;
    std::vector<float> quantiles;
    quantiles.push_back(0.01f);
    quantiles.push_back(0.99f);
    quantiles.push_back(0.01f);
    space.addDimension(Dimension(CumulativeProbabilities, quantiles));

    address = space.address();
    address = space.initialiseInvalidCoordinates(address);
    BOOST_CHECK_EQUAL(address.size(), size_t(1));
    BOOST_CHECK(space.isValid(address));
    BOOST_CHECK(address.isValid(0));
    BOOST_CHECK_CLOSE(address.coordinate<float>(0), float(0.5), 0.001f);
  }
}

BOOST_AUTO_TEST_CASE(equals)
{
  using namespace dal;

  DataSpace space1, space2;

  RasterDimensions rasterDimensions(96, 95, 3.0, 5.5, 6.6);

  space1.addDimension(Dimension(Space, RegularDiscretisation,
         rasterDimensions));

  BOOST_CHECK(space1 == space1);
  BOOST_CHECK(space2 == space2);
  BOOST_CHECK(space1 != space2);

  space2 = space1;

  BOOST_CHECK(space1 == space2);
}

BOOST_AUTO_TEST_CASE(replace_dimension)
{
  using namespace dal;

  DataSpace space;

  std::vector<size_t> timeSteps;
  timeSteps.push_back(1);
  timeSteps.push_back(100);
  timeSteps.push_back(1);
  space.addDimension(Dimension(Time, timeSteps));
  BOOST_CHECK_EQUAL(space.dimension(0).meaning(), Time);

  RasterDimensions rasterDimensions(96, 95, 3.0, 5.5, 6.6);
  space.replaceDimension(0,
         Dimension(Space, RegularDiscretisation, rasterDimensions));

  BOOST_CHECK_EQUAL(space.dimension(0).meaning(), Space);
}

BOOST_AUTO_TEST_CASE(contains)
{
  using namespace dal;

  {
    std::vector<size_t> timeSteps;
    timeSteps.push_back( 9);
    timeSteps.push_back(21);
    timeSteps.push_back( 1);

    RasterDimensions rasterDimensions(3, 2, 3.0);

    DataSpace space;
    space.addDimension(Dimension(Time, timeSteps));
    space.addDimension(Dimension(Space, RegularDiscretisation,
         rasterDimensions));

    DataSpaceAddress address(space.address());
    BOOST_CHECK_EQUAL(address.size(), size_t(2));
    BOOST_CHECK(!space.contains(address));

    address.setCoordinate<size_t>(0, 9);
    BOOST_CHECK(!space.contains(address));

    // North west corner.
    address.setCoordinate<SpatialCoordinate>(1, SpatialCoordinate(0.0, 0.0));
    BOOST_CHECK(space.contains(address));

    // South east corner.
    address.setCoordinate<SpatialCoordinate>(1, SpatialCoordinate(6.0, -9.0));
    BOOST_CHECK(space.contains(address));

    // Outside of area.
    address.setCoordinate<SpatialCoordinate>(1, SpatialCoordinate(-1.0, -1.0));
    BOOST_CHECK(!space.contains(address));

    address.setCoordinate<SpatialCoordinate>(1, SpatialCoordinate(0.0, 0.0));
    BOOST_CHECK(space.contains(address));

    address.setCoordinate<size_t>(0, 8);
    BOOST_CHECK(!space.contains(address));

    address.setCoordinate<size_t>(0, 11);
    BOOST_CHECK(space.contains(address));
  }

  {
    DataSpace space;
    std::vector<size_t> timeSteps, rows, cols;
    timeSteps.push_back( 9);
    timeSteps.push_back(21);
    timeSteps.push_back( 2);
    space.addDimension(Dimension(Time, timeSteps));

    DataSpaceAddress address(space.address());
    BOOST_CHECK(!space.contains(address));

    address.setCoordinate<size_t>(0, 9);
    BOOST_CHECK(space.contains(address));

    address.setCoordinate<size_t>(0, 10);
    BOOST_CHECK(!space.contains(address));

    address.setCoordinate<size_t>(0, 11);
    BOOST_CHECK(space.contains(address));
  }

  {
    DataSpace space;
    std::vector<float> quantiles;
    quantiles.push_back(float(0.1));
    quantiles.push_back(float(0.9));
    quantiles.push_back(float(0.01));
    space.addDimension(Dimension(CumulativeProbabilities, quantiles));

    DataSpaceAddress address(space.address());
    BOOST_CHECK(!space.contains(address));

    address.setCoordinate<float>(0, 0.1f);
    BOOST_CHECK(space.contains(address));
  }
}


#ifndef IMAGE_REPROJECTION_PLUGINS_DEM_SURFACE_MODEL_HPP
#define IMAGE_REPROJECTION_PLUGINS_DEM_SURFACE_MODEL_HPP

#include <limits>
#include <string>

#include <image_reprojection/surface_model.hpp>
#include <image_reprojection/transform.hpp>
#include <nav_msgs/OccupancyGrid.h>
#include <ros/node_handle.h>
#include <tf/transform_datatypes.h>
#include <topic_tools/shape_shifter.h>

#include <opencv2/core/core.hpp>

#include <boost/thread/locks.hpp>
#include <boost/thread/shared_mutex.hpp>

namespace image_reprojection_plugins {

class DEMSurfaceModel : public image_reprojection::SurfaceModel {
public:
  DEMSurfaceModel() {}

  virtual ~DEMSurfaceModel() {}

  virtual void update(const topic_tools::ShapeShifter &surface) {
    const nav_msgs::OccupancyGridConstPtr dem(surface.instantiate< nav_msgs::OccupancyGrid >());
    CV_Assert(dem);
    update(*dem);
  }

  void update(const nav_msgs::OccupancyGrid &dem) {
    CV_Assert(dem.info.width > 0);
    CV_Assert(dem.info.height > 0);
    CV_Assert(dem.data.size() == dem.info.width * dem.info.height);

    boost::unique_lock< boost::shared_mutex > write_lock(mutex_);

    // frame id of arguments of intersection()
    tf_frame_id_ = dem.header.frame_id;

    // frame conversion between DEM and tf frames
    tf::poseMsgToTF(dem.info.origin, dem2tf_);
    tf2dem_ = dem2tf_.inverse();

    // resolution of DEM (xy)
    resolution_ = dem.info.resolution;

    // elevation data (z)
    data_.create(dem.info.width, dem.info.height, CV_64FC1);
    const double data_scale((max_data_ - min_data_) / 100.);
    for (int xid = 0; xid < dem.info.width; ++xid) {
      for (int yid = 0; yid < dem.info.height; ++yid) {
        const int did(xid + yid * dem.info.width);
        data_.at< double >(yid, xid) = min_data_ + data_scale * dem.data[did];
      }
    }
  }

  virtual std::string getFrameId() const {
    boost::shared_lock< boost::shared_mutex > read_lock(mutex_);
    return tf_frame_id_;
  }

private:
  virtual void onInit() {
    boost::unique_lock< boost::shared_mutex > write_lock(mutex_);

    ros::NodeHandle &pnh(getPrivateNodeHandle());
    min_data_ = pnh.param("min_data", 0.);
    max_data_ = pnh.param("max_data", 1.);
  }

  virtual void onIntersection(const cv::Vec3f &src_origin, const cv::Mat &src_direction,
                              cv::Mat &dst, cv::Mat &mask) const {
    namespace ir = image_reprojection;

    boost::shared_lock< boost::shared_mutex > read_lock(mutex_);

    multirayDEMIntersection(ir::transform(src_origin, tf2dem_),
                            ir::transform(src_direction, tf2dem_.getBasis(), mask), dst, mask);
    dst = ir::transform(dst, dem2tf_, mask);
  }

  void multirayDEMIntersection(const cv::Vec3f &src_origin, const cv::Mat &src_direction,
                               cv::Mat &dst, cv::Mat &mask) const {
    dst.create(src_direction.size(), CV_32FC3);
    for (int x = 0; x < src_direction.size().width; ++x) {
      for (int y = 0; y < src_direction.size().height; ++y) {
        unsigned char &m(mask.at< unsigned char >(y, x));
        const cv::Vec3f &sd(src_direction.at< cv::Vec3f >(y, x));
        cv::Vec3f &d(dst.at< cv::Vec3f >(y, x));
        m = (m != 0 && rayDEMIntersection(src_origin, sd, d)) ? 1 : 0;
      }
    }
  }

  bool rayDEMIntersection(const cv::Vec3f &src_origin, const cv::Vec3f &src_direction,
                          cv::Vec3f &dst) const {
    bool intersects(false);
    double min_distance;
    for (int xid = 0; xid < data_.size().width; ++xid) {
      for (int yid = 0; yid < data_.size().height; ++yid) {
        cv::Vec3f this_intersection;
        if (rayBoxIntersection(src_origin, src_direction, xid, yid, this_intersection)) {
          const double this_distance(cv::norm(src_origin, this_intersection));
          if (!intersects || (intersects && this_distance < min_distance)) {
            intersects = true;
            min_distance = this_distance;
            dst = this_intersection;
          }
        }
      }
    }

    return intersects;
  }

  bool rayBoxIntersection(const cv::Vec3f &src_origin, const cv::Vec3f &src_direction,
                          const int xid, const int yid, cv::Vec3f &dst) const {
    // box specs
    const cv::Vec3d box_origin(xid * resolution_, yid * resolution_, 0.);
    const cv::Vec3d box_length(resolution_, resolution_, data_.at< double >(yid, xid));

    // intersection point (p) can bescribed as
    //   p = r + t * d
    //   (r: ray origin, d: ray direction, t >= 0)

    // calculate possible range of t in each dimension
    cv::Vec3d t_1d_min, t_1d_max;
    for (int i = 0; i < 3; ++i) {
      if (!raySegmentIntersection(src_origin[i], src_direction[i], box_origin[i],
                                  box_origin[i] + box_length[i], t_1d_min[i], t_1d_max[i])) {
        return false;
      }
    }

    // check if overlap of range of t exists
    double t_min, t_max;
    cv::minMaxIdx(t_1d_min, NULL, &t_min);
    cv::minMaxIdx(t_1d_max, &t_max);
    if (/* no overlap */ t_min > t_max) {
      return false;
    }

    // return the nearest intersection point in ray direction
    dst = src_origin + t_min * src_direction;
    return true;
  }

  bool raySegmentIntersection(const double src_origin, const double src_direction,
                              const double seg_min, const double seg_max, double &t_min,
                              double &t_max) const {
    // o: ray origin
    // ->, <-: ray direction
    // m, M: segment (m <= M)

    if (src_direction == 0.) {
      // o --- m --- M  OR  m --- M --- o
      if (src_origin < seg_min || src_origin > seg_max) {
        return false;
      }
      // m --- o --- M
      t_min = 0.;
      t_max = std::numeric_limits< double >::max();
      return true;
    } else if (src_direction > 0.) {
      // m --- M ---- o->
      if (src_origin > seg_max) {
        return false;
      }
      // m --- o-> --- M
      if (src_origin > seg_min) {
        t_min = 0.;
        t_max = (seg_max - src_origin) / src_direction;
        return true;
      }
      // o-> --- m ---- M
      t_min = (seg_min - src_origin) / src_direction;
      t_max = (seg_max - src_origin) / src_direction;
      return true;
    } else /* src_direction < 0. */ {
      // <-o --- m --- M
      if (src_origin < seg_min) {
        return false;
      }
      // m --- <-o --- M
      if (src_origin < seg_max) {
        t_min = 0.;
        t_max = (seg_min - src_origin) / src_direction;
        return true;
      }
      // m --- M --- <-o
      t_min = (seg_max - src_origin) / src_direction;
      t_max = (seg_min - src_origin) / src_direction;
      return true;
    }

    static const bool NEVER_REACH_HERE_BUG(false);
    CV_Assert(NEVER_REACH_HERE_BUG);
  }

private:
  mutable boost::shared_mutex mutex_;

  // info from nav_msgs::OccupancyGrid
  std::string tf_frame_id_;
  tf::Transform dem2tf_, tf2dem_;
  double resolution_;
  cv::Mat data_;

  // additional info from rosparam
  double min_data_, max_data_;
};

} // namespace image_reprojection_plugins

#endif


// __BEGIN_LICENSE__
//  Copyright (c) 2018, United States Government as represented by the
//  Administrator of the National Aeronautics and Space Administration. All
//  rights reserved.
//
//  The NGT platform is 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.
// __END_LICENSE__

// Joint n-cloud joint alignment. Reimplementing the Matlab code at
// https://searchcode.com/file/13619767/Code/matlab/GlobalProcrustesICP/globalProcrustes.m
// One tiny change is that that code had a little bug, when invoked
// with 10 iterations it would actually do only 9. So, before
// comparing this code's result to that one, un-comment the line which
// says "original incorrect logic".  Based on the paper: Global
// registration of multiple point clouds embedding the Generalized
// Procrustes Analysis into an ICP framework by Toldo, Roberto and
// Beinat, Alberto and Crosilla, Fabio.

// Here some code is used that is part of the Fast Global Registration
// software, https://github.com/IntelVCL/FastGlobalRegistration which
// is released under the MIT License (MIT).

#include <vector>
#include <iostream>
#include <fstream>
#include <flann/flann.hpp>
#include <Eigen/Core>
#include <Eigen/Geometry>

#include <vw/Core/Stopwatch.h>
#include <asp/Core/Common.h>
#include <asp/Core/Macros.h>
#include <asp/Core/PointUtils.h>
#include <asp/Core/EigenUtils.h>
#include <asp/Tools/pc_align_utils.h>
#include <boost/filesystem.hpp>
#include <boost/algorithm/string.hpp>

namespace fs = boost::filesystem;
namespace po = boost::program_options;

using namespace vw;
using namespace std;
using namespace vw::cartography;
using namespace asp;

typedef double RealT; // We will use doubles in libpointmatcher.
typedef PointMatcher<RealT> PM;
typedef PM::DataPoints DP;

// TODO: Figure out these magic numbers
#define FIFTEEN 15
#define ONE_TWO_EIGHT 128

typedef flann::Index<flann::L2<double> > KDTree_double;

/// Options container
struct Options : public vw::cartography::GdalWriteOptions {
  // Input
  string in_prefix, in_transforms, datum, csv_format_str, csv_proj4_str;
  int    num_iter, max_num_points;
  double semi_major_axis, semi_minor_axis, rel_error_tol;
  bool   save_transformed_clouds, align_to_first_cloud, verbose;
  std::vector<std::string> cloud_files;
  // Output
  string out_prefix;

  Options(){}
};

void handle_arguments( int argc, char *argv[], Options& opt ) {
  po::options_description general_options("");
  general_options.add_options()
    ("num-iterations",           po::value(&opt.num_iter)->default_value(100),
     "Maximum number of iterations.")
    ("max-num-points", po::value(&opt.max_num_points)->default_value(1000000),
     "Maximum number of (randomly picked) points from each cloud to use.")
    ("csv-format",               po::value(&opt.csv_format_str)->default_value(""), asp::csv_opt_caption().c_str())
    ("csv-proj4",                po::value(&opt.csv_proj4_str)->default_value(""),
     "The PROJ.4 string to use to interpret the entries in input CSV files.")
    ("datum",                    po::value(&opt.datum)->default_value(""),
     "Use this datum for CSV files instead of auto-detecting it. Options: WGS_1984, D_MOON (1,737,400 meters), D_MARS (3,396,190 meters), MOLA (3,396,000 meters), NAD83, WGS72, and NAD27. Also accepted: Earth (=WGS_1984), Mars (=D_MARS), Moon (=D_MOON).")
    ("semi-major-axis",          po::value(&opt.semi_major_axis)->default_value(0),
     "Explicitly set the datum semi-major axis in meters.")
    ("semi-minor-axis",          po::value(&opt.semi_minor_axis)->default_value(0),
     "Explicitly set the datum semi-minor axis in meters.")
    ("output-prefix,o",          po::value(&opt.out_prefix)->default_value(""),
     "Specify the output prefix. The computed alignment transforms and, if desired, the transformed clouds, will be saved to names starting with this prefix.")
    ("save-transformed-clouds", po::bool_switch(&opt.save_transformed_clouds)->default_value(false)->implicit_value(true),
     "Apply the obtained alignment transforms to the input clouds and save them.")
    ("initial-transforms-prefix", po::value(&opt.in_prefix)->default_value(""),
     "The prefix of the transforms to be used as initial guesses. The naming convention is the same as for the transforms written on output.")
    ("initial-transforms", po::value(&opt.in_transforms)->default_value(""),
     "Specify the initial transforms as a list of files separated by spaces and in quotes, that is, as 'trans1.txt ... trans_n.txt'.")
    ("relative-error-tolerance", po::value(&opt.rel_error_tol)->default_value(1e-10),
     "Stop when the change in the error divided by the error itself is less than this.")
    ("align-to-first-cloud", po::bool_switch(&opt.align_to_first_cloud)->default_value(false)->implicit_value(true),
     "Align the other clouds to the first one, rather than to their common centroid.")
    ("verbose", po::bool_switch(&opt.verbose)->default_value(false)->implicit_value(true),
     "Print the alignment error after each iteration.");

  general_options.add( vw::cartography::GdalWriteOptionsDescription(opt) );

  po::options_description positional("");
  positional.add_options()
    ("cloud-files", po::value<std::vector<std::string> >(&opt.cloud_files));
  po::positional_options_description positional_desc;
  positional_desc.add("cloud-files", -1);
  std::string usage("[options] <cloud-files> -o <output prefix>");

  bool allow_unregistered = false;
  std::vector<std::string> unregistered;
  po::variables_map vm =
    asp::check_command_line(argc, argv, opt, general_options, general_options,
                            positional, positional_desc, usage,
                            allow_unregistered, unregistered);

  if ( opt.out_prefix.empty() )
    vw_throw( ArgumentErr() << "Missing output prefix.\n" << usage << general_options );

  if ( opt.num_iter < 0 )
    vw_throw( ArgumentErr() << "The number of iterations must be non-negative.\n"
              << usage << general_options );

  if ( (opt.semi_major_axis != 0 && opt.semi_minor_axis == 0) ||
       (opt.semi_minor_axis != 0 && opt.semi_major_axis == 0) ){
    vw_throw( ArgumentErr() << "One of the semi-major or semi-minor axes"
              << " was specified, but not the other one.\n"
              << usage << general_options );
  }

  if (opt.semi_major_axis < 0 || opt.semi_minor_axis < 0){
    vw_throw( ArgumentErr() << "The semi-major and semi-minor axes cannot "
                            << "be negative.\n" << usage << general_options );
  }

  if (opt.datum != "" && opt.semi_major_axis != 0 && opt.semi_minor_axis != 0 ){
    vw_throw( ArgumentErr() << "Both the datum string and datum semi-axes were "
                            << "specified. At most one needs to be set.\n"
                            << usage << general_options );
  }

  if (opt.cloud_files.size() < 2)
    vw_throw( ArgumentErr() << "Must have at least two clouds to align.\n"
              << usage << general_options );

  if (opt.in_prefix != "" && opt.in_transforms != "") {
    vw_throw( ArgumentErr() << "Cannot specify both the initial transforms prefix "
              << "and the list of initial transforms.\n" << usage << general_options );

  }

}

bool triplet_less(vw::Vector3 const& p, vw::Vector3 const& q){

  if (p.x() < q.x()) return true;
  if (p.x() > q.x()) return false;

  if (p.y() < q.y()) return true;
  if (p.y() > q.y()) return false;

  if (p.z() < q.z()) return true;
  if (p.z() > q.z()) return false;

  return false;
}

// Define equality for pairs of integers
bool operator==(std::pair<int, int> const& p, std::pair<int, int> const& q){
  return (p.first == q.first) && (p.second == q.second);
}

// An operator used to sort pairs of integers
struct CustomCompare {
  bool operator()(std::pair<int, int> const& p, std::pair<int, int> const& q){
    if (p.first < q.first) return true;
    if (p.first > q.first) return false;

    if (p.second < q.second) return true;
    if (p.second > q.second) return false;

    return false;
  }
};

bool vector_less(Eigen::VectorXd const& p, Eigen::VectorXd const& q){
  for (int i = 0; i < p.size(); i++) {
    if (p[i] < q[i]) return true;
    if (p[i] > q[i]) return false;
  }
  return false;
}

// Convert a cloud from libpointmatcher's format to a vector of points
void convert_cloud(DP const& in_cloud, std::vector<vw::Vector3> & out_cloud){
  out_cloud.clear();
  for (int col = 0; col < in_cloud.features.cols(); col++) {
    vw::Vector3 p;
    for (int row = 0; row < 3; row++)
      p[row] = in_cloud.features(row, col);
    out_cloud.push_back(p);
  }
}

// Read a cloud stored in a plain text file
void read_cloud(std::string const& file, std::vector<vw::Vector3> & cloud, double shift){
  cloud.clear();
  std::ifstream ifs(file.c_str());
  double x, y, z;
  vw::Vector3 p;
  while (ifs >> x >> y >> z) {
    p = vw::Vector3(x + shift, y + shift, z + shift);
    cloud.push_back(p);
  }
  ifs.close();
}

// Apply a rotation + transform to a cloud
void apply_transform_to_cloud( std::vector<vw::Vector3> & cloud, Eigen::MatrixXd const& T){
  for (size_t pointIter = 0; pointIter < cloud.size(); pointIter++) {
    Eigen::VectorXd v(4);
    v << cloud[pointIter][0], cloud[pointIter][1], cloud[pointIter][2], 1.0;
    v = T*v;
    for (int it = 0; it < 3; it++)
      cloud[pointIter][it] = v[it];
  }
}

void print_cloud(std::vector<vw::Vector3> const& cloud){
  for (size_t it = 0; it < cloud.size(); it++) {
    vw_out() << cloud[it].x() << ' ' << cloud[it].y() << ' ' << cloud[it].z() << std::endl;
  }
}

void sort_and_make_unqiue(std::vector<vw::Vector3> & cloud){
  std::sort(cloud.begin(), cloud.end(), triplet_less);
  std::vector<vw::Vector3>::iterator it = std::unique(cloud.begin(), cloud.end());
  cloud.resize(std::distance(cloud.begin(), it));
}

void BuildKDTree_double(std::vector<vw::Vector3> const& cloud, KDTree_double* tree){
  int rows = cloud.size();
  int dim = vw::Vector3().size();

  // Need to have this to avoid a seg fault later
  if (rows*dim <= 0)
    vw_throw( ArgumentErr() << "Cannot operate on empty clouds.\n" );

  std::vector<double> dataset(rows * dim);
  flann::Matrix<double> dataset_mat(&dataset[0], rows, dim);

  for (int i = 0; i < rows; i++) {
    for (int j = 0; j < dim; j++) {
      dataset[i * dim + j] = cloud[i][j];
    }
  }

  // Using the default leaf_max_size = 10 in flann::KDTreeSingleIndexParams()
  KDTree_double temp_tree(dataset_mat, flann::KDTreeSingleIndexParams(FIFTEEN));
  temp_tree.buildIndex();
  *tree = temp_tree;
}

void SearchKDTree_double(KDTree_double* tree, const vw::Vector3& input,
                         std::vector<int>& indices,
                         std::vector<double>& dists, int nn){
  int rows_t = 1;
  int dim = input.size();

  std::vector<double> query;
  query.resize(rows_t*dim);
  for (int i = 0; i < dim; i++)
    query[i] = input[i];
  flann::Matrix<double> query_mat(&query[0], rows_t, dim);

  indices.resize(rows_t*nn);
  dists.resize(rows_t*nn);
  flann::Matrix<int> indices_mat(&indices[0], rows_t, nn);
  flann::Matrix<double> dists_mat(&dists[0], rows_t, nn);

  tree->knnSearch(query_mat, indices_mat, dists_mat, nn, flann::SearchParams(ONE_TWO_EIGHT));
}

std::string transform_file(std::string const& out_prefix, int index){
  std::ostringstream os;
  os << out_prefix << "-transform-" << index << ".txt";
  std::string transFile = os.str();
  return transFile;
}

/// Save the transforms
void write_transforms(std::vector<Eigen::MatrixXd> const& transVec,
		     std::string const& out_prefix){

  for (size_t it = 0; it < transVec.size(); it++) {
    std::string transFile = transform_file(out_prefix, it);
    vw_out() << "Writing: " << transFile << std::endl;
    write_transform(transVec[it], transFile);
  }
}

/// Read the transforms. We assume we know their number.
void read_transforms(std::vector<Eigen::MatrixXd> & transVec,
		     std::string const& in_prefix){

  for (size_t it = 0; it < transVec.size(); it++) {
    std::string transFile = transform_file(in_prefix, it);
    read_transform(transVec[it], transFile);
  }
}

/// Read the transforms from a list. We assume we know their numbers.
void read_transforms_from_list(std::vector<Eigen::MatrixXd> & transVec,
                               std::string const& transform_list){
  std::istringstream is(transform_list);
  for (size_t it = 0; it < transVec.size(); it++) {
    std::string transFile;
    if ( !(is >> transFile) )
      vw_throw( ArgumentErr() << "Cannot parse enough transform files from: " << transform_list);
    read_transform(transVec[it], transFile);
  }
}

int main(int argc, char *argv[]){

  Options opt;
  try {
    handle_arguments(argc, argv, opt);

    // Parse the csv format string and csv projection string
    asp::CsvConv csv_conv;
    csv_conv.parse_csv_format(opt.csv_format_str, opt.csv_proj4_str);

    // Create the output directory
    vw::create_out_dir(opt.out_prefix);

    // Turn on logging to file
    asp::log_to_file(argc, argv, "", opt.out_prefix);

    // Try to read the georeference/datum info
    GeoReference geo;
    read_georef(opt.cloud_files, opt.datum, opt.csv_proj4_str,
                opt.semi_major_axis, opt.semi_minor_axis,
                opt.csv_format_str,  csv_conv, geo);

    int numClouds = opt.cloud_files.size();

    // The vector of transforms among the clouds
    std::vector<Eigen::MatrixXd> transVec(numClouds);
    for (int cloudIter = 0; cloudIter < numClouds; cloudIter++) {
      transVec[cloudIter] = Eigen::MatrixXd::Identity(4, 4);
    }

    std::vector< std::vector<vw::Vector3> > clouds(numClouds);

    // Load the first subsampled point cloud. Calculate the shift to apply to all clouds.
    vw::Vector3 shift;
    bool   calc_shift = true; // Shift points so the first point is (0,0,0)
    bool   is_lola_rdr_format = false;   // may get overwritten
    double mean_ref_longitude    = 0.0;  // may get overwritten
    double mean_source_longitude = 0.0;  // may get overwritten
    BBox2 empty_box;
    DP in_cloud;
    bool verbose = true;
    load_cloud(opt.cloud_files[0], opt.max_num_points, empty_box,
               calc_shift, shift, geo, csv_conv, is_lola_rdr_format,
               mean_ref_longitude, verbose, in_cloud);
    convert_cloud(in_cloud, clouds[0]);

    vw_out() << "Data shifted internally by subtracting: " << shift << std::endl;

    calc_shift = false; // We will use the same shift from here on

    for (int cloudIter = 1; cloudIter < numClouds; cloudIter++) {
      load_cloud(opt.cloud_files[cloudIter], opt.max_num_points, empty_box,
                 calc_shift, shift, geo, csv_conv, is_lola_rdr_format,
                 mean_ref_longitude, verbose, in_cloud);
      convert_cloud(in_cloud, clouds[cloudIter]);
    }

    // Read the clouds
    for (int cloudIter = 0; cloudIter < numClouds; cloudIter++) {
      //read_cloud(opt.cloud_files[cloudIter], clouds[cloudIter], shifts[cloudIter]);
      sort_and_make_unqiue(clouds[cloudIter]);
    }

    // Read any initial transforms, either using a prefix or an explicit list
    bool has_init_transform = false;
    if (opt.in_prefix != "" || opt.in_transforms != "") {
      has_init_transform = true;
      if (opt.in_prefix != "")
        read_transforms(transVec, opt.in_prefix);
      else if (opt.in_transforms != "")
        read_transforms_from_list(transVec, opt.in_transforms);

      // The point clouds are shifted, so shift the initial transforms as well.
      for (int cloudIter = 0; cloudIter < numClouds; cloudIter++)
        transVec[cloudIter] = apply_shift(transVec[cloudIter], shift);
    }

    // Apply the initial transforms to the clouds
    if (has_init_transform) {
      for (int cloudIter = 0; cloudIter < numClouds; cloudIter++)
        apply_transform_to_cloud(clouds[cloudIter], transVec[cloudIter]);
    }

    // Build the trees
    std::vector< boost::shared_ptr<KDTree_double> > Trees;
    for (int it = 0; it < numClouds; it++) {
      boost::shared_ptr<KDTree_double>
        tree(new KDTree_double(flann::KDTreeSingleIndexParams(FIFTEEN)));
      BuildKDTree_double(clouds[it], tree.get());
      Trees.push_back(tree);
    }

    std::string errCaption = std::string("Computing the error, defined as the mean of ") +
      "pairwise distances from each cloud to the centroid cloud.\n";
    if (opt.verbose) vw_out() << errCaption;

    // Mean error from the centroid cloud points to all matching
    // points in all clouds.  We don't use the mean square error like
    // the Matlab code, which had bugs in how it was computed anyway.
    double initError = -1, errAfter = -1, prevError = -1;

    Stopwatch sw1;
    sw1.start();

    vw::TerminalProgressCallback tpc("asp", "Performing alignment\t--> ");
    double inc_amount = 1.0 / std::max(opt.num_iter, 1);

    // Don't display the progress in verbose mode, as then it interferes with
    // printing of errors.
    if (!opt.verbose)
      tpc.report_progress(0);

    int step = 0; // Starting step
    // int step = 1; // original incorrect logic in the Matlab code
    while (step < opt.num_iter){

      bool firstStep = (step == 0);
      bool lastStep  = (step == opt.num_iter - 1);

      std::vector<int> modelSpan(numClouds+1);
      int numOfPoints = 0;
      for (int it = 0; it < numClouds; it++) {
        modelSpan[it] = numOfPoints;
        numOfPoints += clouds[it].size();
      }
      modelSpan[numClouds] = numOfPoints;

      // This will record for each point in each cloud which point in
      // every other cloud is closest to it. This matrix will store the
      // indices of these points.
      std::vector<Eigen::VectorXd> CentroidPtsBelMod(numOfPoints);
      for (size_t row = 0; row < CentroidPtsBelMod.size(); row++) {
        CentroidPtsBelMod[row] = Eigen::VectorXd(numClouds);
        for (int cloudIter = 0; cloudIter < numClouds; cloudIter++) {
          CentroidPtsBelMod[row][cloudIter] = -1;
        }
      }

      for (size_t i = 0; i < numClouds; i++) {

        //spanI = modelSpan(i)+1:modelSpan(i+1);
        std::vector<int> spanI;
        int beg = modelSpan[i], end = modelSpan[i+1] - 1;
        for (int it = beg; it <= end; it++) spanI.push_back(it);

        //CentroidPtsBelMod(spanI,i) = 1:length(spanI);
        for (size_t it = 0; it < spanI.size(); it++) {
          CentroidPtsBelMod[spanI[it]][i] = it;
        }

        for (size_t j = i + 1; j < numClouds; j++) {

          //spanJ = modelSpan(j)+1:modelSpan(j+1);
          std::vector<int> spanJ;
          int beg = modelSpan[j], end = modelSpan[j+1] - 1;
          for (int it = beg; it <= end; it++) spanJ.push_back(it);

          std::vector<int> match, match2;
          std::vector<double> dist;
          typedef std::set<std::pair<int, int>, CustomCompare> PairType;

          // For each point in cloud i, find a match in cloud j
          PairType Corr1;
          for (size_t index_i = 0; index_i < clouds[i].size(); index_i++){
            SearchKDTree_double(Trees[j].get(), clouds[i][index_i], match, dist, 1);
            if (match.empty()) continue; // should not happen
            int index_j = match[0];
            Corr1.insert(std::pair<int, int>(index_j, index_i));
          }

          // Now do it in reverse
          PairType Corr2;
          for (size_t index_j = 0; index_j < clouds[j].size(); index_j++){
            SearchKDTree_double(Trees[i].get(), clouds[j][index_j], match, dist, 1);
            if (match.empty()) continue; // should not happen
            int index_i = match[0];
            Corr2.insert(std::pair<int, int>(index_j, index_i));
          }

          std::vector< std::pair<int, int> > Corr;
          for (PairType::iterator it = Corr1.begin(); it != Corr1.end(); it++) {
            PairType::iterator it2 = Corr2.find(*it);
            if (it2 == Corr2.end()) continue;

            Corr.push_back(*it);
          }

          for (size_t it = 0; it < Corr.size(); it++) {
            // CentroidPtsBelMod(spanI(Corr(:,2)),j) = Corr(:,1)';
            CentroidPtsBelMod[spanI[Corr[it].second]][j] = Corr[it].first;

            // CentroidPtsBelMod(spanJ(Corr(:,1)),i) = Corr(:,2)';
            CentroidPtsBelMod[spanJ[Corr[it].first]][i] = Corr[it].second;
          }

        }
      }

      std::sort(CentroidPtsBelMod.begin(), CentroidPtsBelMod.end(), vector_less);

      // Remove non-unique elements, and remove entries which exist only in one cloud
      int pos = 0;
      for (size_t row = 0; row < CentroidPtsBelMod.size(); row++) {

        int num_good = 0;
        for (int cloudIter = 0; cloudIter < CentroidPtsBelMod[row].size(); cloudIter++) {
          if (CentroidPtsBelMod[row][cloudIter] >= 0)
            num_good++;
        }
        if (num_good < 2) continue;
        if (row > 0 && CentroidPtsBelMod[row - 1] == CentroidPtsBelMod[row])
          continue;

        CentroidPtsBelMod[pos] = CentroidPtsBelMod[row];
        pos++;
      }
      CentroidPtsBelMod.resize(pos);

      // We connected the clouds. Find the average cloud.
      std::vector<vw::Vector3> centroid(CentroidPtsBelMod.size());
      for (size_t row = 0; row < CentroidPtsBelMod.size(); row++) {
        vw::Vector3 pt(0, 0, 0);
        int num = 0;
        for (int cloudIter = 0; cloudIter < CentroidPtsBelMod[row].size(); cloudIter++) {
          int idx = CentroidPtsBelMod[row][cloudIter];
          if (idx < 0) continue;
          vw::Vector3 curr = clouds[cloudIter][idx];
          pt += curr;
          num++;
        }
        pt /= num;
        centroid[row] = pt;
      }

      double errBefore = 0;
      errAfter = 0;
      int numErrors = 0;

      // Find the transform from each cloud to the centroid, and apply it to each cloud
      for (size_t cloudIter = 0; cloudIter < numClouds; cloudIter++) {

        std::vector<Eigen::Vector3d> src, dst;
        Eigen::Matrix3d rot;
        Eigen::Vector3d trans;

        for (int row = 0; row < CentroidPtsBelMod.size(); row++) {
          int pointIter = CentroidPtsBelMod[row][cloudIter];
          if (pointIter < 0) continue;

          vw::Vector3 curr = clouds[cloudIter][pointIter];
          vw::Vector3 ctr  = centroid[row];

          Eigen::Vector3d p;
          p[0] = curr[0]; p[1] = curr[1]; p[2] = curr[2];
          src.push_back(p);

          p[0] = ctr[0]; p[1] = ctr[1]; p[2] = ctr[2];
          dst.push_back(p);

	  errBefore += norm_2(curr - ctr);

	  numErrors++;
        }

        computeRigidTransform(src, dst, rot, trans);

        // Update the output transforms
        Eigen::MatrixXd currT = Eigen::MatrixXd::Identity(4, 4);;
        for (int row = 0; row < 3; row++)
          for (int col = 0; col < 3; col++)
            currT(row, col) = rot(row, col);
        for (int row = 0; row < 3; row++)
          currT(row, 3) = trans(row);
        transVec[cloudIter] = currT*transVec[cloudIter];

        // Move the clouds to the new location for the next iteration
        apply_transform_to_cloud(clouds[cloudIter], currT);

	// Compute the error after the transform is applied
        for (int row = 0; row < CentroidPtsBelMod.size(); row++) {
          int pointIter = CentroidPtsBelMod[row][cloudIter];
          if (pointIter < 0) continue;
          vw::Vector3 curr = clouds[cloudIter][pointIter];
          vw::Vector3 ctr  = centroid[row];
	  errAfter += norm_2(curr - ctr);
        }

      }

      errBefore /= numErrors;
      errAfter /= numErrors;

      if (firstStep) {
        initError = errBefore;
        if (opt.verbose) {
          vw_out() << "Error before alignment: " << initError  << std::endl;
        }
      }

      if (opt.verbose)
        vw_out() << "Error at iteration " << step << ": " << errAfter << "\n";
      else
        tpc.report_incremental_progress(inc_amount);

      if (!firstStep) {
        if (errAfter > 0 && std::abs(errAfter - prevError)/errAfter < opt.rel_error_tol) {
          break;
        }
      }

      prevError = errAfter;

      step++;
    } // End of iterations refining the transforms

    if (!verbose) {
      tpc.report_finished();
    }

    sw1.stop();

    if (!opt.verbose) {
      vw_out() << "\n" << errCaption;
      vw_out() << "Error before alignment: " << initError  << std::endl;
    }
    vw_out() << "Error after alignment:  " << errAfter << std::endl;

    vw_out() << "Performed: " << step + 1 << " iterations.\n";

    vw_out() << "\nAlignment took " << sw1.elapsed_seconds() << " [s]" << endl;

    if (opt.align_to_first_cloud) {
      // Make the first transform be the identity
      for (int cloudIter = 1; cloudIter < numClouds; cloudIter++)
        transVec[cloudIter] = transVec[0].inverse() * transVec[cloudIter];
      transVec[0] = Eigen::MatrixXd::Identity(4, 4);
    }

    // Undo the shift when the clouds were read
    for (int cloudIter = 0; cloudIter < numClouds; cloudIter++)
      transVec[cloudIter] = apply_shift(transVec[cloudIter], -shift);

    for (int cloudIter = 0; cloudIter < numClouds; cloudIter++)
      vw_out() << "Final transform for cloud: " << cloudIter << ":\n"
               << transVec[cloudIter] << std::endl;
    write_transforms(transVec, opt.out_prefix);

    if (opt.save_transformed_clouds) {
      for (int cloudIter = 0; cloudIter < numClouds; cloudIter++) {
        std::ostringstream os;
        os << opt.out_prefix << "-trans_cloud-" << cloudIter;
        std::string trans_prefix = os.str();
        save_trans_point_cloud(opt, opt.cloud_files[cloudIter], trans_prefix,
                               geo, csv_conv, transVec[cloudIter]);
      }
    }

  } ASP_STANDARD_CATCHES;

  return 0;
}


/*
* Software License Agreement (BSD License)
*
*  Copyright (c) 2015, Fabian Girrbach, Social Robotics Lab, University of Freiburg
*  All rights reserved.
*
*  Redistribution and use in source and binary forms, with or without
*  modification, are permitted provided that the following conditions are met:
*
*  * Redistributions of source code must retain the above copyright notice, this
*    list of conditions and the following disclaimer.
*  * Redistributions in binary form must reproduce the above copyright notice,
*    this list of conditions and the following disclaimer in the documentation
*    and/or other materials provided with the distribution.
*  * Neither the name of the copyright holder nor the names of its contributors
*    may be used to endorse or promote products derived from this software
*    without specific prior written permission.
*
*  THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
*  AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
*  IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
*  DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
*  FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
*  DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
*  SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
*  CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
*  OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
*  OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
*/

#include <srl_nearest_neighbor_tracker/data_association/global_nearest_neighbor_data_association.h>

#include <srl_nearest_neighbor_tracker/base/defs.h>
#include <srl_nearest_neighbor_tracker/base/stl_helpers.h>
#include <srl_nearest_neighbor_tracker/base/lap.h>
#include <srl_nearest_neighbor_tracker/ros/params.h>
#include <boost/foreach.hpp>
#define foreach BOOST_FOREACH

#include <fstream>

#include <Eigen/LU>
#include <limits>
#include <map>

namespace srl_nnt {

void GlobalNearestNeighborDataAssociation::initializeDataAssociation(const ros::NodeHandle& nodeHandle, const ros::NodeHandle& privateNodeHandle) {
    m_nodeHandle = nodeHandle;
    m_privateNodeHandle = privateNodeHandle;
}

Pairings GlobalNearestNeighborDataAssociation::performDataAssociation(Tracks& tracks, const Observations& observations){

    ROS_DEBUG("Performing data association");

    //
    // Step 0: Make sure observations are in a good state
    //
    foreach(Observation::Ptr observation, observations) {
        observation->matched = false;
    }

    foreach(Track::Ptr track, tracks) {
        track->observation.reset();
    }

    Pairings compatiblePairings;

    //
    // Step 1: go through all possible associations and store compatible ones
    //

    //
    // Step 1:  go through all possible associations and store them in a array
    //          we fill a cost matrix with all possible associations
    //

    if (tracks.size() > 0 && observations.size() > 0)
    {
        LAPSolver<double> linearAssignmentProblem;
        const double MATRIX_LN_EPS = -1e8;
        const double MAX_GATING_DISTANCE = Params::get<double>("max_gating_distance", 1.0);

        Eigen::MatrixXd costMatrix =  Eigen::MatrixXd::Constant(tracks.size(), observations.size(), BIG_COST);
        Pairing pairingArray[tracks.size()][observations.size()];
        for(size_t t=0; t < tracks.size() ; t++)
        {
            Track::Ptr track = tracks.at(t);
            Pairings trackPairings;
            for(size_t ob=0; ob < observations.size() ; ob++)
            {
                Observation::Ptr observation = observations.at(ob);
                // Create a new pairing
                Pairing& pairingRef (pairingArray[t][ob]);

                // Calculate innovation v and inverse of innovation covariance S
                ObsVector v = observation->z - track->state->zp();
                // do simple gating here to save expensive operations
                if (v.norm() < MAX_GATING_DISTANCE)
                {
                    pairingRef.v = v;
                    pairingRef.track = track;
                    pairingRef.observation = observation;
                    pairingRef.validated = false;

                    ObsMatrix S = track->state->H() * track->state->Cp() * track->state->H().transpose() + observation->R;
                    Eigen::FullPivLU<ObsMatrix> lu(S);
                    double ln_det_S = 0;
                    ln_det_S = log(lu.determinant());
                    // Calculate inverse of innovation covariance if possible
                    if (ln_det_S > MATRIX_LN_EPS)
                    {
                        pairingRef.Sinv = lu.inverse();
                        pairingRef.d = (pairingRef.v.transpose() * pairingRef.Sinv * pairingRef.v)(0,0);
                        pairingRef.singular = pairingRef.d < 0.0;
                    }
                    else
                    {
                        pairingRef.Sinv = ObsMatrix::Constant(OBS_DIM, OBS_DIM, numeric_limits<double>::quiet_NaN());
                        pairingRef.d = numeric_limits<double>::quiet_NaN();
                        pairingRef.singular = true;

                        ROS_WARN_STREAM("Singular pairing encountered!\nTrack " << track->id << " measurement prediction:\n" << track->state->zp() << "\nTrack covariance prediction:\n" << track->state->Cp()
                                        << "\nObservation " << observation->id << " mean:\n" << observation->z << "\nObservation covariance:\n" << observation->R );
                    }

                    // Perform gating
                    if((!pairingRef.singular) && (pairingRef.d < CHI2INV_99[OBS_DIM])) {
                        // Store in list of compatible pairings
                        if (Params::get<bool>("use_correlation_log", false))
                            costMatrix(t, ob) = (pairingRef.d + ln_det_S)/track->detectionProbability;
                        else
                            costMatrix(t, ob) = (pairingRef.d)/track->detectionProbability;
                        //ROS_INFO_STREAM("Current cost = " << costMatrix(t, ob) << " d " << pairing->d << " lnS " << ln_det_S);
                    }
                }
            }
        }

        ROS_DEBUG("%zu compatible pairings have been found for %zu existing tracks and %zu new observations!", compatiblePairings.size(), tracks.size(), observations.size() );

        Eigen::VectorXi trackAssignments, obsAssignments;

        double totalCost = linearAssignmentProblem.calculateAssignment(costMatrix, trackAssignments, obsAssignments);

        int matrixSize = tracks.size() * observations.size();
        static int maxMatrixSize = 0;
        if(matrixSize > maxMatrixSize) {
            maxMatrixSize = matrixSize;
            std::ofstream file("cost_matrix.txt");
            file << "COST MATRIX DUMP:" << std::endl << std::endl;
            file << costMatrix << std::endl;
            file << "TOTAL COST: " << totalCost << std::endl;
            file << "TRACK ASSIGNMENTS:" << trackAssignments << std::endl;
            file << "OBS ASSIGNMENTS:" << obsAssignments << std::endl;
        }

        // ROS_INFO_STREAM("Cost Matrix " << costMatrix);
        // ROS_INFO_STREAM("Track Assignments " << trackAssignments);
        // ROS_INFO_STREAM("Observation Assignments " << obsAssignments);

        for (size_t i = 0 ; i < tracks.size(); i++)
        {
            Track::Ptr track = tracks.at(i);
            double currentCost = costMatrix(i, obsAssignments(i));
            if (currentCost < BIG_COST)
            {
                compatiblePairings.push_back(boost::make_shared<Pairing>(pairingArray[i][obsAssignments(i)]));
                //ROS_DEBUG_STREAM("Assigning track " << i << "(" << track->id << ")"<< " to observation " << obsAssignments(i));
                markAsMatched(track, compatiblePairings.back());
            }
            else if (track->trackStatus == Track::OCCLUDED){
                // Track is occluded
                track->observation.reset();
                track->trackStatus = Track::OCCLUDED;
                track->numberOfConsecutiveOcclusions++;
            }
            else{
                // Track is missed
                track->observation.reset();
                track->trackStatus = Track::MISSED;
                track->numberOfConsecutiveMisses++;
            }
        }
    }

    return compatiblePairings;

}

}


//==================================================================================================
/*!
  @file

  @copyright 2016 NumScale SAS
  @copyright 2016 J.T. Lapreste

  Distributed under the Boost Software License, Version 1.0.
  (See accompanying file LICENSE.md or copy at http://boost.org/LICENSE_1_0.txt)
*/
//==================================================================================================
#ifndef BOOST_SIMD_FUNCTION_SIMD_SIN_INCLUDED
#define BOOST_SIMD_FUNCTION_SIMD_SIN_INCLUDED

#include <boost/simd/function/scalar/sin.hpp>
#include <boost/simd/arch/common/generic/function/autodispatcher.hpp>
#include <boost/simd/arch/common/generic/function/sin.hpp>

#endif


#include <boost/test/unit_test.hpp>
#include <boost/concept_check.hpp>

#include "che/score/ev_alignment.h" // header to test
#include "che/io/file_fasta.h"

using namespace biosim;

BOOST_AUTO_TEST_SUITE(suite_ev_alignment)

BOOST_AUTO_TEST_CASE(ev_alignment_evaluate) {
  che::score::ev_alignment score;
  BOOST_CHECK(score.evaluate(che::alignment(che::structure())) == 0);

  che::structure s(che::io::file_fasta::read("../test/data/1a00_a.fasta").get_first_structure("A"));
  che::ps seq(s.get_ps());
  che::ps seq_shortened;
  seq_shortened.insert(seq_shortened.end(), seq.begin(), seq.begin() + 6);
  che::structure struct_shortened(s.get_identifier(), s.get_storage(), seq_shortened);
  BOOST_CHECK(score.evaluate(che::alignment(struct_shortened)) == 0);

  std::vector<che::structure> s_vec2(2, struct_shortened);
  che::alignment a2(s_vec2, {0, 0}, {struct_shortened.get_length(), struct_shortened.get_length()});
  BOOST_CHECK_CLOSE(score.evaluate(a2), 33.5659, 1e-3);

  std::vector<che::structure> s_vec3(3, struct_shortened);
  che::alignment a3(s_vec3, {0, 0, 0}, std::vector<size_t>(3, struct_shortened.get_length()));
  BOOST_CHECK_CLOSE(score.evaluate(a3), 33.5659, 1e-3);

  che::ps seq_profile;
  seq_profile.emplace_back(che::cc("ASX"));
  che::structure struct_profile(s.get_storage(), s.get_identifier(), seq_profile);
  std::vector<che::structure> profile_vec(2, struct_profile);
  che::alignment a(profile_vec, {0, 0}, {struct_profile.get_length(), struct_profile.get_length()});
  BOOST_CHECK_CLOSE(score.evaluate(a), 3.49272, 1e-3);

  che::ps seq_non_profile;
  seq_non_profile.insert(seq_non_profile.end(), seq.begin(), seq.begin() + 1);
  che::structure struct_non_profile(s.get_storage(), s.get_identifier(), seq_non_profile);
  std::vector<che::structure> sequence_and_profile{struct_profile, struct_non_profile};
  a = che::alignment(sequence_and_profile, {0, 0}, {struct_profile.get_length(), struct_non_profile.get_length()});
  BOOST_CHECK_CLOSE(score.evaluate(a), -2.60469, 1e-3);

  che::ps seq_gap;
  seq_gap.emplace_back(che::cc('-'));
  che::structure struct_gap(s.get_storage(), s.get_identifier(), seq_gap);
  std::vector<che::structure> sequence_and_gap{struct_non_profile, struct_gap};
  a = che::alignment(sequence_and_gap, {0, 0}, {struct_non_profile.get_length(), struct_gap.get_length()});
  BOOST_CHECK_CLOSE(score.evaluate(a), -4.2143, 1e-3);

  che::ps seq_n;
  seq_n.emplace_back(che::cc("ASN"));
  che::structure struct_n(s.get_storage(), s.get_identifier(), seq_n);
  std::vector<che::structure> structs{struct_profile, struct_n};
  a = che::alignment(structs, {0, 0}, {struct_profile.get_length(), struct_n.get_length()});
  BOOST_CHECK_CLOSE(score.evaluate(a), 3.46247, 1e-3);

  structs = std::vector<che::structure>(2, struct_n);
  a = che::alignment(structs, {0, 0}, {struct_n.get_length(), struct_n.get_length()});
  BOOST_CHECK_CLOSE(score.evaluate(a), 5.65324, 1e-3);

  che::ps seq_x;
  seq_x.emplace_back(che::cc('X'));
  che::structure struct_x(s.get_storage(), s.get_identifier(), seq_x);
  structs = std::vector<che::structure>({struct_n, struct_x});
  a = che::alignment(structs, {0, 0}, {struct_n.get_length(), struct_x.get_length()});
  BOOST_CHECK_CLOSE(score.evaluate(a), -0.991768, 1e-3);

  structs = std::vector<che::structure>(2, struct_x);
  a = che::alignment(structs, {0, 0}, {struct_x.get_length(), struct_x.get_length()});
  BOOST_CHECK_CLOSE(score.evaluate(a), -1.10221, 1e-3);

  che::ps seq_i;
  seq_i.emplace_back(che::cc('I'));
  che::structure struct_i(s.get_storage(), s.get_identifier(), seq_i);
  structs = std::vector<che::structure>(2, struct_i);
  a = che::alignment(structs, {0, 0}, {struct_i.get_length(), struct_i.get_length()});
  BOOST_CHECK_CLOSE(score.evaluate(a), 3.99851, 1e-3);

  che::ps seq_iv;
  seq_iv.emplace_back(che::cc(che::cc::weight_map({{"ILE", 1.0}, {"VAL", 1.0}})));
  che::structure struct_iv(s.get_storage(), s.get_identifier(), seq_iv);
  structs = std::vector<che::structure>({struct_i, struct_iv});
  a = che::alignment(structs, {0, 0}, {struct_i.get_length(), struct_iv.get_length()});
  BOOST_CHECK_CLOSE(score.evaluate(a), 3.27278, 1e-3);

  structs = std::vector<che::structure>(2, struct_iv);
  a = che::alignment(structs, {0, 0}, {struct_iv.get_length(), struct_iv.get_length()});
  BOOST_CHECK_CLOSE(score.evaluate(a), 3.21538, 1e-3);
}

BOOST_AUTO_TEST_SUITE_END()


// Boost.Geometry (aka GGL, Generic Geometry Library)

// Copyright (c) 2007-2015 Barend Gehrels, Amsterdam, the Netherlands.
// Copyright (c) 2008-2015 Bruno Lalande, Paris, France.
// Copyright (c) 2009-2015 Mateusz Loskot, London, UK.

// This file was modified by Oracle on 2015.
// Modifications copyright (c) 2015, Oracle and/or its affiliates.

// Contributed and/or modified by Menelaos Karavelas, on behalf of Oracle

// Parts of Boost.Geometry are redesigned from Geodan's Geographic Library
// (geolib/GGL), copyright (c) 1995-2010 Geodan, Amsterdam, the Netherlands.

// Use, modification and distribution is subject to the Boost Software License,
// Version 1.0. (See accompanying file LICENSE_1_0.txt or copy at
// http://www.boost.org/LICENSE_1_0.txt)

#ifndef BOOST_GEOMETRY_ALGORITHMS_DETAIL_ENVELOPE_BOX_HPP
#define BOOST_GEOMETRY_ALGORITHMS_DETAIL_ENVELOPE_BOX_HPP

#include <boost/geometry/core/cs.hpp>
#include <boost/geometry/core/coordinate_system.hpp>
#include <boost/geometry/core/tags.hpp>

#include <boost/geometry/strategies/strategy_transform.hpp>

#include <boost/geometry/algorithms/convert.hpp>
#include <boost/geometry/algorithms/transform.hpp>
#include <boost/geometry/algorithms/detail/normalize.hpp>

#include <boost/geometry/algorithms/dispatch/envelope.hpp>

namespace boost { namespace geometry
{

#ifndef DOXYGEN_NO_DETAIL
namespace detail { namespace envelope
{

struct envelope_box_on_spheroid
{
    template<typename BoxIn, typename BoxOut>
    static inline void apply(BoxIn const& box_in, BoxOut& mbr)
    {
        BoxIn box_in_normalized = detail::return_normalized<BoxIn>(box_in);

        geometry::transform(box_in_normalized, mbr);
    }
};

template <typename CSTag>
struct envelope_box
{
    template<typename BoxIn, typename BoxOut>
    static inline void apply(BoxIn const& box_in, BoxOut& mbr)
    {
        geometry::convert(box_in, mbr);
    }
};

template <>
struct envelope_box<spherical_equatorial_tag>
    : envelope_box_on_spheroid
{};

template <>
struct envelope_box<geographic_tag>
    : envelope_box_on_spheroid
{};

}} // namespace detail::envelope
#endif // DOXYGEN_NO_DETAIL

#ifndef DOXYGEN_NO_DISPATCH
namespace dispatch
{

template <typename Box>
struct envelope<Box, box_tag>
    : detail::envelope::envelope_box<typename cs_tag<Box>::type>
{};

} // namespace dispatch
#endif

}} // namespace boost::geometry

#endif // BOOST_GEOMETRY_ALGORITHMS_DETAIL_ENVELOPE_BOX_HPP


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

#ifndef CRYPTO3_BLOCK_KASUMI_HPP
#define CRYPTO3_BLOCK_KASUMI_HPP

#include <boost/endian/arithmetic.hpp>
#include <boost/endian/conversion.hpp>

#include <nil/crypto3/block/detail/kasumi/kasumi_functions.hpp>

#include <nil/crypto3/block/detail/block_stream_processor.hpp>
#include <nil/crypto3/block/detail/cipher_modes.hpp>

namespace nil {
    namespace crypto3 {
        namespace block {
            /*!
             * @brief Kasumi. A 64-bit cipher used in 3GPP mobile phone protocols.
             * There is no reason to use it outside of this context.
             *
             * @ingroup block
             */
            class kasumi {
            protected:
                typedef detail::kasumi_functions policy_type;

                constexpr static const std::size_t key_schedule_size = policy_type::key_schedule_size;
                typedef typename policy_type::key_schedule_type key_schedule_type;

            public:
                constexpr static const std::size_t rounds = policy_type::rounds;

                constexpr static const std::size_t word_bits = policy_type::word_bits;
                typedef typename policy_type::word_type word_type;

                constexpr static const std::size_t block_bits = policy_type::block_bits;
                constexpr static const std::size_t block_words = policy_type::block_words;
                typedef typename policy_type::block_type block_type;

                constexpr static const std::size_t key_bits = policy_type::key_bits;
                constexpr static const std::size_t key_words = policy_type::key_words;
                typedef typename policy_type::key_type key_type;

                template<class Mode, typename StateAccumulator, std::size_t ValueBits>
                struct stream_processor {
                    struct params_type {

                        constexpr static const std::size_t value_bits = ValueBits;
                        constexpr static const std::size_t length_bits = policy_type::word_bits * 2;
                    };

                    typedef block_stream_processor<Mode, StateAccumulator, params_type> type;
                };

                typedef typename stream_endian::little_octet_big_bit endian_type;

                kasumi(const key_type &key) {
                    schedule_key(key_schedule, key);
                }

                ~kasumi() {
                    key_schedule.fill(0);
                }

                inline block_type encrypt(const block_type &plaintext) const {
                    return encrypt_block(plaintext, key_schedule);
                }

                inline block_type decrypt(const block_type &ciphertext) const {
                    return decrypt_block(ciphertext, key_schedule);
                }

            protected:
                inline block_type encrypt_block(const block_type &plaintext,
                                                const key_schedule_type &key_schedule) const {
                    word_type B0 = boost::endian::native_to_big(plaintext[0]);
                    word_type B1 = boost::endian::native_to_big(plaintext[1]);
                    word_type B2 = boost::endian::native_to_big(plaintext[2]);
                    word_type B3 = boost::endian::native_to_big(plaintext[3]);

                    for (size_t j = 0; j != rounds; j += 2) {
                        const word_type *K = &key_schedule[8 * j];

                        word_type R = B1 ^ (policy_type::template rotl<1>(B0) & K[0]);
                        word_type L = B0 ^ (policy_type::template rotl<1>(R) | K[1]);

                        L = policy_type::FI(L ^ K[2], K[3]) ^ R;
                        R = policy_type::FI(R ^ K[4], K[5]) ^ L;
                        L = policy_type::FI(L ^ K[6], K[7]) ^ R;

                        R = B2 ^= R;
                        L = B3 ^= L;

                        R = policy_type::FI(R ^ K[10], K[11]) ^ L;
                        L = policy_type::FI(L ^ K[12], K[13]) ^ R;
                        R = policy_type::FI(R ^ K[14], K[15]) ^ L;

                        R ^= (policy_type::template rotl<1>(L) & K[8]);
                        L ^= (policy_type::template rotl<1>(R) | K[9]);

                        B0 ^= L;
                        B1 ^= R;
                    }

                    return {boost::endian::big_to_native(B0), boost::endian::big_to_native(B1),
                            boost::endian::big_to_native(B2), boost::endian::big_to_native(B3)};
                }

                inline block_type decrypt_block(const block_type &ciphertext,
                                                const key_schedule_type &key_schedule) const {
                    word_type B0 = boost::endian::native_to_big(ciphertext[0]);
                    word_type B1 = boost::endian::native_to_big(ciphertext[1]);
                    word_type B2 = boost::endian::native_to_big(ciphertext[2]);
                    word_type B3 = boost::endian::native_to_big(ciphertext[3]);

                    for (size_t j = 0; j != rounds; j += 2) {
                        const word_type *K = &key_schedule[8 * (6 - j)];

                        word_type L = B2, R = B3;

                        L = policy_type::FI(L ^ K[10], K[11]) ^ R;
                        R = policy_type::FI(R ^ K[12], K[13]) ^ L;
                        L = policy_type::FI(L ^ K[14], K[15]) ^ R;

                        L ^= (policy_type::template rotl<1>(R) & K[8]);
                        R ^= (policy_type::template rotl<1>(L) | K[9]);

                        R = B0 ^= R;
                        L = B1 ^= L;

                        L ^= (policy_type::template rotl<1>(R) & K[0]);
                        R ^= (policy_type::template rotl<1>(L) | K[1]);

                        R = policy_type::FI(R ^ K[2], K[3]) ^ L;
                        L = policy_type::FI(L ^ K[4], K[5]) ^ R;
                        R = policy_type::FI(R ^ K[6], K[7]) ^ L;

                        B2 ^= L;
                        B3 ^= R;
                    }

                    return {boost::endian::big_to_native(B0), boost::endian::big_to_native(B1),
                            boost::endian::big_to_native(B2), boost::endian::big_to_native(B3)};
                }

                key_schedule_type key_schedule;

                void schedule_key(key_schedule_type &key_schedule, const key_type &key) {
                    std::array<word_type, 16> K = {0};
                    for (size_t i = 0; i != rounds; ++i) {
                        K[i] = boost::endian::native_to_big(key[i]);
                        K[i + 8] = K[i] ^ policy_type::round_constants[i];
                    }

                    for (size_t i = 0; i != rounds; ++i) {
                        key_schedule[8 * i] = policy_type::template rotl<2>(K[(i + 0) % 8]);
                        key_schedule[8 * i + 1] = policy_type::template rotl<1>(K[(i + 2) % 8 + 8]);
                        key_schedule[8 * i + 2] = policy_type::template rotl<5>(K[(i + 1) % 8]);
                        key_schedule[8 * i + 3] = K[(i + 4) % 8 + 8];
                        key_schedule[8 * i + 4] = policy_type::template rotl<8>(K[(i + 5) % 8]);
                        key_schedule[8 * i + 5] = K[(i + 3) % 8 + 8];
                        key_schedule[8 * i + 6] = policy_type::template rotl<13>(K[(i + 6) % 8]);
                        key_schedule[8 * i + 7] = K[(i + 7) % 8 + 8];
                    }

                    K.fill(0);
                }
            };
        }    // namespace block
    }        // namespace crypto3
}    // namespace nil
#endif


// Copyright (c) 2021 Graphcore Ltd. All rights reserved.
#include "CTCInferenceCodeletTestConnection.hpp"

#include <poplar/Engine.hpp>
#include <poplar/Graph.hpp>
#include <poplar/IPUModel.hpp>
#include <poplar/Type.hpp>

#include <poplibs_support/LogArithmetic.hpp>
#include <poplibs_support/TestDevice.hpp>
#include <poplibs_test/CTCInference.hpp>
#include <poplibs_test/Util.hpp>
#include <poputil/VertexTemplates.hpp>

#include <boost/multi_array.hpp>

using namespace poplar;
using namespace poplar::program;
using namespace poplibs_test::ctc;
using namespace poplibs_test;
using namespace poplibs_test::util;
using namespace poplibs_support;
using namespace poputil;

namespace poplibs_test {
namespace ctc {

template <typename InputType, typename PartialsType>
std::vector<Candidate<PartialsType>> runGenerateCandidatesCodelet(
    Graph &graph, TestDevice &device, DeviceType deviceType, Type inType,
    Type partialsType, const boost::multi_array<InputType, 2> &logProbsIn,
    unsigned timestep,
    const std::vector<BeamProbability<PartialsType>> &beamProbs,
    const BeamHistory &beamHistory, unsigned classToMakeAddend, unsigned beam,
    unsigned blankClass, bool testGenerateCopyVertex, bool profile) {
  const auto target = graph.getTarget();

  const auto maxT = logProbsIn.size();
  const auto numClassesIncBlank = logProbsIn[0].size();
  const auto beamwidth = beamProbs.size();

  auto logProbs =
      graph.addVariable(inType, {maxT, numClassesIncBlank}, "logProbs");

  // Codelet tests don't add an extra timestep = 0 initial state like the
  // main implementation does. The effect of this is that we need to provide
  // timestep + 1 here to compensate
  auto currentTimestep = graph.addConstant(UNSIGNED_INT, {}, timestep + 1);
  auto complete = graph.addConstant(UNSIGNED_INT, {}, 0u);

  graph.setTileMapping(logProbs, 0);
  graph.setTileMapping(currentTimestep, 0);
  graph.setTileMapping(complete, 0);

  auto cs = graph.addComputeSet("cs");
  const auto vertexName = testGenerateCopyVertex
                              ? "popnn::CTCGenerateCopyCandidates"
                              : "popnn::CTCGenerateExtendCandidates";
  auto vertex = graph.addVertex(
      cs, templateVertex(vertexName, inType, partialsType, UNSIGNED_INT));
  graph.setTileMapping(vertex, 0);

  graph.connect(vertex["logProbs"], logProbs.flatten());
  graph.connect(vertex["currentTimestep"], currentTimestep);
  graph.connect(vertex["complete"], complete);

  graph.setInitialValue(vertex["numClassesIncBlank"], numClassesIncBlank);

  if (testGenerateCopyVertex) {
    graph.setInitialValue(vertex["beamIdx"], beam);
    graph.setInitialValue(vertex["blankClass"], blankClass);
  } else {
    graph.setInitialValue(vertex["startBeam"], 0);
    graph.setInitialValue(vertex["endBeam"], beamwidth);
    graph.setInitialValue(vertex["addendSymbol"], classToMakeAddend);
  }
  Sequence uploadProg, downloadProg;
  std::vector<std::pair<std::string, char *>> tmap;

  CandidateHandles rawCandidates;
  if (testGenerateCopyVertex) {
    rawCandidates =
        createAndConnectCandidates(graph, vertex, "candidate", partialsType, {},
                                   uploadProg, downloadProg, tmap);

  } else {
    rawCandidates = createAndConnectCandidates(graph, vertex, "extendCandidate",
                                               partialsType, {beamwidth},
                                               uploadProg, downloadProg, tmap);
  }

  // Inputs
  std::unique_ptr<char[]> rawLogProbs;
  rawLogProbs = allocateHostMemoryForTensor(logProbs, "logProbs", graph,
                                            uploadProg, downloadProg, tmap);

  auto rawBeamProbs = createAndConnectBeamProbs(
      graph, vertex, partialsType, {beamwidth},
      (testGenerateCopyVertex ? BeamScalars::NON_BLANK : BeamScalars::BLANK),
      uploadProg, downloadProg, tmap);

  // Initialise inputs
  std::vector<unsigned> lastBeamOutputsIn{};
  std::vector<double> beamProbNonBlankIn{};
  std::vector<double> beamProbBlankIn{};
  std::vector<double> beamProbTotalIn{};
  for (unsigned i = 0; i < beamwidth; i++) {
    lastBeamOutputsIn.push_back(beamHistory.getLastOutput(i));
    beamProbNonBlankIn.push_back(beamProbs[i].pnb);
    beamProbBlankIn.push_back(beamProbs[i].pb);
    beamProbTotalIn.push_back(log::add(beamProbs[i].pb, beamProbs[i].pnb));
  }

  copy(target, logProbsIn, inType, rawLogProbs.get());
  copy(target, lastBeamOutputsIn, UNSIGNED_INT, rawBeamProbs.lastOutput.get());
  copy(target, beamProbNonBlankIn, partialsType, rawBeamProbs.pnb.get());
  copy(target, beamProbBlankIn, partialsType, rawBeamProbs.pb.get());
  copy(target, beamProbTotalIn, partialsType, rawBeamProbs.pTotal.get());

  // TODO Need to initialise outputs?

  OptionFlags engineOptions;
  if (profile) {
    engineOptions.set("debug.instrumentCompute", "true");
  }
  Sequence prog;
  prog.add(Execute(cs));
  Engine engine(graph, Sequence{uploadProg, prog, downloadProg}, engineOptions);
  attachStreams(engine, tmap);
  device.bind([&](const Device &d) {
    engine.load(d);
    engine.run();
  });

  const auto totalCandidates = testGenerateCopyVertex ? 1 : beamwidth;
  std::vector<unsigned> candidateParentOut(totalCandidates);
  std::vector<unsigned> candidateAddendOut(totalCandidates);
  std::vector<unsigned> validCandidatesOut(1);

  // TODO partialsType == float
  std::vector<float> candidateBeamProbBlankOut(totalCandidates);
  std::vector<float> candidateBeamProbNonBlankOut(totalCandidates);
  std::vector<float> candidateBeamProbTotalOut(totalCandidates);
  copy(target, UNSIGNED_INT, rawCandidates.parent.get(), candidateParentOut);
  copy(target, UNSIGNED_INT, rawCandidates.addend.get(), candidateAddendOut);
  copy(target, partialsType, rawCandidates.probNonBlank.get(),
       candidateBeamProbNonBlankOut);
  copy(target, partialsType, rawCandidates.probBlank.get(),
       candidateBeamProbBlankOut);
  copy(target, partialsType, rawCandidates.probTotal.get().get(),
       candidateBeamProbTotalOut);

  if (profile && deviceType != DeviceType::Cpu) {
    engine.printProfileSummary(std::cout,
                               OptionFlags{{"showExecutionSteps", "true"}});
  }

  std::vector<Candidate<float>> candidates;
  for (unsigned i = 0; i < totalCandidates; i++) {
    candidates.push_back({candidateParentOut[i], candidateAddendOut[i],
                          candidateBeamProbNonBlankOut[i],
                          candidateBeamProbBlankOut[i],
                          candidateBeamProbTotalOut[i]});
  }
  return candidates;
}

template std::vector<Candidate<float>> runGenerateCandidatesCodelet(
    Graph &graph, TestDevice &device, DeviceType deviceType, Type inType,
    Type partialsType, const boost::multi_array<float, 2> &logProbsIn,
    unsigned timestep, const std::vector<BeamProbability<float>> &beamProbs,
    const BeamHistory &beamHistory, unsigned classToMakeAddend, unsigned beam,
    unsigned blankClass, bool testGenerateCopyVertex, bool profile);

} // namespace ctc
} // namespace poplibs_test