/* Siconos is a program dedicated to modeling, simulation and control
 * of non smooth dynamical systems.
 *
 * Copyright 2016 INRIA.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
*/

#include "SiconosConfig.h"
#include <boost/numeric/bindings/ublas/vector_proxy.hpp>
#include <boost/numeric/ublas/matrix_proxy.hpp>
#include <boost/numeric/bindings/trans.hpp>
#include <boost/numeric/bindings/blas/level2.hpp>
#include <boost/numeric/bindings/ublas/vector.hpp>
#include <boost/numeric/bindings/ublas/matrix.hpp>
#include <boost/numeric/bindings/std/vector.hpp>

// Note Franck : sounds useless. It seems it's defined in bindings
// (to be checked, especially on windows)

// #define BIND_FORTRAN_LOWERCASE_UNDERSCORE
namespace siconosBindings = boost::numeric::bindings;

// for ublas::axpy_prod, ...
#include <boost/numeric/ublas/operation.hpp>

// for matrix stuff like value_type
//#include <boost/numeric/bindings/traits/ublas_matrix.hpp>

#include "SiconosVector.hpp"
#include "SimpleMatrix.hpp"
#include "BlockMatrixIterators.hpp"
#include "BlockMatrix.hpp"
#include "SiconosAlgebra.hpp"

using namespace Siconos;

// ========== Products matrix - vector //

// Computation of y = A*x
//
// Two specific functions are used to handle all the cases where x or y are blocks.
// All of their blocks can also be blocks. Then we use:
// - private_prod to "slice" A when y is block, ie according to its rows.
// - private_addprod to "slice" A when x is block, and to sum over the columns of blocks to compute y = sum subA x[i].

// The following function is private and used inside prod(...) public functions.
// It is required to deal with block vectors of blocks ( of blocks ...).
// It computes res = subA*x +res, subA being a submatrix of A (rows from startRow to startRow+sizeY and columns between startCol and startCol+sizeX).
// If x is a block vector, it call the present function for all blocks.
const SiconosVector prod(const SiconosMatrix& A, const SiconosVector& x)
{
  // To compute y = A * x
  assert(!(A.isPLUFactorized()) && "A is PLUFactorized in prod !!");

  if (A.size(1) != x.size())
    SiconosMatrixException::selfThrow("prod(matrix,vector) error: inconsistent sizes.");

  unsigned int numA = A.num();
  unsigned int numX = x.num();

  if (numA == 0) // if A is block ...
    SiconosMatrixException::selfThrow("prod(matrix,vector) error: not yet implemented for block matrix.");

  if (numA == 6) // A = 0
    return (DenseVect)(ublas::zero_vector<double>(x.size()));

  else if (numA == 7) // A = Identity
    return x;

  else
  {
    if (numX == 1)
    {
      if (numA == 1)
        return (DenseVect)(prod(*A.dense(), *x.dense()));
      else if (numA == 2)
        return (DenseVect)(prod(*A.triang(), *x.dense()));
      else if (numA == 3)
        return (DenseVect)(prod(*A.sym(), *x.dense()));
      else if (numA == 4)
        return (DenseVect)(prod(*A.sparse(), *x.dense()));
      else // if(numA==5)
        return (DenseVect)(prod(*A.banded(), *x.dense()));
    }
    else //if(numX == 4)
    {
      if (numA == 1)
        return (DenseVect)(prod(*A.dense(), *x.sparse()));
      else if (numA == 2)
        return (DenseVect)(prod(*A.triang(), *x.sparse()));
      else if (numA == 3)
        return (DenseVect)(prod(*A.sym(), *x.sparse()));
      else if (numA == 4)
        return (DenseVect)(prod(*A.sparse(), *x.sparse()));
      else // if(numA==5)
        return (DenseVect)(prod(*A.banded(), *x.sparse()));
    }
  }
}

void prod(double a, const SiconosMatrix& A, const SiconosVector& x, SiconosVector& y, bool init)
{
  // To compute y = a*A * x in an "optimized" way (in comparison with y = prod(A,x) )
  // or y += a*A*x if init = false.
  assert(!(A.isPLUFactorized()) && "A is PLUFactorized in prod !!");

  if (A.size(1) != x.size())
    SiconosMatrixException::selfThrow("prod(A,x,y) error: inconsistent sizes between A and x.");

  if (A.size(0) != y.size())
    SiconosMatrixException::selfThrow("prod(A,x,y) error: inconsistent sizes between A and y.");

  unsigned int numA = A.num();
  unsigned int numX = x.num();
  unsigned int numY = y.num();

  if (numA == 0) // If A is Block
    SiconosMatrixException::selfThrow("prod(A,x,y) error: not yet implemented for block matrices.");

  if (numA == 6) // A = 0
  {
    if (init)
      y.zero();
    //else nothing
  }

  else if (numA == 7) // A = identity
  {
    scal(a, x, y, init);
  }

  else // A is not 0 or identity
  {

    // === First case: y is not a block vector ===
    {
      {
        if (init)
        {
          if (&x != &y) // if no common memory between x and y.
          {
            if (numX == 1)
            {
              if (numY != 1)
                SiconosMatrixException::selfThrow("prod(A,x,y) error: y (output) must be a dense vector.");

              if (numA == 1)
                noalias(*y.dense()) = a * ublas::prod(*A.dense(), *x.dense());
              else if (numA == 2)
                noalias(*y.dense()) = a * ublas::prod(*A.triang(), *x.dense());
              else if (numA == 3)
                noalias(*y.dense()) = a * ublas::prod(*A.sym(), *x.dense());
              else if (numA == 4)
                noalias(*y.dense()) = a * ublas::prod(*A.sparse(), *x.dense());
              else //if(numA==5)
                noalias(*y.dense()) = a * ublas::prod(*A.banded(), *x.dense());
            }
            else //if(numX == 4)
            {
              if (numY != 1 && numA != 4)
                SiconosMatrixException::selfThrow("prod(A,x,y) error: y (output) must be a dense vector.");

              if (numA == 1)
                noalias(*y.dense()) = a * ublas::prod(*A.dense(), *x.sparse());
              else if (numA == 2)
                noalias(*y.dense()) = a * ublas::prod(*A.triang(), *x.sparse());
              else if (numA == 3)
                noalias(*y.dense()) = a * ublas::prod(*A.sym(), *x.sparse());
              else if (numA == 4)
              {
                if (numY == 1)
                  noalias(*y.dense()) = a * ublas::prod(*A.sparse(), *x.sparse());
                else
                  noalias(*y.sparse()) = a * ublas::prod(*A.sparse(), *x.sparse());
              }
              else //if(numA==5)
                noalias(*y.dense()) = a * ublas::prod(*A.banded(), *x.sparse());
            }
          }
          else // if x and y are the same object => alias
          {
            if (numX == 1)
            {
              if (numA == 1)
                *y.dense() = a * ublas::prod(*A.dense(), *x.dense());
              else if (numA == 2)
                *y.dense() = a * ublas::prod(*A.triang(), *x.dense());
              else if (numA == 3)
                *y.dense() = a * ublas::prod(*A.sym(), *x.dense());
              else if (numA == 4)
                *y.dense() = a * ublas::prod(*A.sparse(), *x.dense());
              else //if(numA==5)
                *y.dense() = a * ublas::prod(*A.banded(), *x.dense());
            }
            else //if(numX == 4)
            {
              if (numA == 1)
                *y.sparse() = a * ublas::prod(*A.dense(), *x.sparse());
              else if (numA == 2)
                *y.sparse() = a * ublas::prod(*A.triang(), *x.sparse());
              else if (numA == 3)
                *y.sparse() = a * ublas::prod(*A.sym(), *x.sparse());
              else if (numA == 4)
                *y.sparse() = a * ublas::prod(*A.sparse(), *x.sparse());
              else //if(numA==5)
                *y.sparse() = a * ublas::prod(*A.banded(), *x.sparse());
            }
          }
        }
        else // += case
        {
          if (&x != &y) // if no common memory between x and y.
          {
            if (numX == 1)
            {
              if (numY != 1)
                SiconosMatrixException::selfThrow("prod(A,x,y) error: y (output) must be a dense vector.");

              if (numA == 1)
                noalias(*y.dense()) += a * ublas::prod(*A.dense(), *x.dense());
              else if (numA == 2)
                noalias(*y.dense()) += a * ublas::prod(*A.triang(), *x.dense());
              else if (numA == 3)
                noalias(*y.dense()) += a * ublas::prod(*A.sym(), *x.dense());
              else if (numA == 4)
                noalias(*y.dense()) += a * ublas::prod(*A.sparse(), *x.dense());
              else //if(numA==5)
                noalias(*y.dense()) += a * ublas::prod(*A.banded(), *x.dense());
            }
            else //if(numX == 4)
            {
              if (numY != 1 && numA != 4)
                SiconosMatrixException::selfThrow("prod(A,x,y) error: y (output) must be a dense vector.");

              if (numA == 1)
                noalias(*y.dense()) += a * ublas::prod(*A.dense(), *x.sparse());
              else if (numA == 2)
                noalias(*y.dense()) += a * ublas::prod(*A.triang(), *x.sparse());
              else if (numA == 3)
                noalias(*y.dense()) += a * ublas::prod(*A.sym(), *x.sparse());
              else if (numA == 4)
              {
                if (numY == 1)
                  noalias(*y.dense()) += a * ublas::prod(*A.sparse(), *x.sparse());
                else
                  noalias(*y.sparse()) += a * ublas::prod(*A.sparse(), *x.sparse());
              }
              else //if(numA==5)
                noalias(*y.dense()) += a * ublas::prod(*A.banded(), *x.sparse());
            }
          }
          else // if x and y are the same object => alias
          {
            if (numX == 1)
            {
              if (numA == 1)
                *y.dense() += a * ublas::prod(*A.dense(), *x.dense());
              else if (numA == 2)
                *y.dense() += a * ublas::prod(*A.triang(), *x.dense());
              else if (numA == 3)
                *y.dense() += a * ublas::prod(*A.sym(), *x.dense());
              else if (numA == 4)
                *y.dense() += a * ublas::prod(*A.sparse(), *x.dense());
              else //if(numA==5)
                *y.dense() += a * ublas::prod(*A.banded(), *x.dense());
            }
            else //if(numX == 4)
            {
              if (numA == 1)
                *y.sparse() += a * ublas::prod(*A.dense(), *x.sparse());
              else if (numA == 2)
                *y.sparse() += a * ublas::prod(*A.triang(), *x.sparse());
              else if (numA == 3)
                *y.sparse() += a * ublas::prod(*A.sym(), *x.sparse());
              else if (numA == 4)
                *y.sparse() += a * ublas::prod(*A.sparse(), *x.sparse());
              else //if(numA==5)
                *y.sparse() += a * ublas::prod(*A.banded(), *x.sparse());
            }
          }
        }
      }
    }
  }
}

void prod(const SiconosVector& x, const SiconosMatrix& A, SiconosVector& y, bool init)
{
  // To compute y = trans(A) * x in an "optimized" way, if init = true
  // (or y = trans(A) * x + y if init = false
  assert(!(A.isPLUFactorized()) && "A is PLUFactorized in prod !!");

  if (A.size(0) != x.size())
    SiconosMatrixException::selfThrow("prod(x,A,y) error: inconsistent sizes between A and x.");

  if (A.size(1) != y.size())
    SiconosMatrixException::selfThrow("prod(x,A,y) error: inconsistent sizes between A and y.");

  unsigned int numA = A.num();
  unsigned int numX = x.num();
  unsigned int numY = y.num();

  if (numA == 0) // If A is Block
    SiconosMatrixException::selfThrow("prod(x,A,y) error: not yet implemented for block matrices.");

  if (numA == 6) // A = 0
  {
    if (init)
      y.zero();
    // else nothing
  }

  else if (numA == 7) // A = identity
  {
    if (!init)
      y += x;
    else
    {
      if (&x != &y) y = x ; // if x and y do not share memory (ie are different objects)
      // else nothing
    }
  }

  else // A is not 0 or identity
  {
    {
      if (init)
      {

        if (&x != &y) // if no common memory between x and y.
        {
          if (numX == 1)
          {
            if (numY != 1)
              SiconosMatrixException::selfThrow("prod(x,A,y) error: y (output) must be a dense vector.");

            if (numA == 1)
              noalias(*y.dense()) = ublas::prod(trans(*A.dense()), *x.dense());
            else if (numA == 2)
              noalias(*y.dense()) = ublas::prod(trans(*A.triang()), *x.dense());
            else if (numA == 3)
              noalias(*y.dense()) = ublas::prod(trans(*A.sym()), *x.dense());
            else if (numA == 4)
              noalias(*y.dense()) = ublas::prod(trans(*A.sparse()), *x.dense());
            else //if(numA==5)
              noalias(*y.dense()) = ublas::prod(trans(*A.banded()), *x.dense());
          }
          else //if(numX == 4)
          {
            if (numY != 1 && numA != 4)
              SiconosMatrixException::selfThrow("prod(x,A,y) error: y (output) must be a dense vector.");

            if (numA == 1)
              noalias(*y.dense()) = ublas::prod(trans(*A.dense()), *x.sparse());
            else if (numA == 2)
              noalias(*y.dense()) = ublas::prod(trans(*A.triang()), *x.sparse());
            else if (numA == 3)
              noalias(*y.dense()) = ublas::prod(trans(*A.sym()), *x.sparse());
            else if (numA == 4)
            {
              if (numY == 1)
                noalias(*y.dense()) = ublas::prod(trans(*A.sparse()), *x.sparse());
              else
                noalias(*y.sparse()) = ublas::prod(trans(*A.sparse()), *x.sparse());
            }
            else //if(numA==5)
              noalias(*y.dense()) = ublas::prod(trans(*A.banded()), *x.sparse());
          }
        }
        else // if x and y are the same object => alias
        {
          if (numX == 1)
          {
            if (numA == 1)
              *y.dense() = ublas::prod(trans(*A.dense()), *x.dense());
            else if (numA == 2)
              *y.dense() = ublas::prod(trans(*A.triang()), *x.dense());
            else if (numA == 3)
              *y.dense() = ublas::prod(trans(*A.sym()), *x.dense());
            else if (numA == 4)
              *y.dense() = ublas::prod(trans(*A.sparse()), *x.dense());
            else //if(numA==5)
              *y.dense() = ublas::prod(trans(*A.banded()), *x.dense());
          }
          else //if(numX == 4)
          {
            if (numA == 1)
              *y.sparse() = ublas::prod(trans(*A.dense()), *x.sparse());
            else if (numA == 2)
              *y.sparse() = ublas::prod(trans(*A.triang()), *x.sparse());
            else if (numA == 3)
              *y.sparse() = ublas::prod(trans(*A.sym()), *x.sparse());
            else if (numA == 4)
              *y.sparse() = ublas::prod(trans(*A.sparse()), *x.sparse());
            else //if(numA==5)
              *y.sparse() = ublas::prod(trans(*A.banded()), *x.sparse());
          }
        }
      }
      else // += case
      {

        if (&x != &y) // if no common memory between x and y.
        {
          if (numX == 1)
          {
            if (numY != 1)
              SiconosMatrixException::selfThrow("prod(x,A,y) error: y (output) must be a dense vector.");

            if (numA == 1)
              noalias(*y.dense()) += ublas::prod(trans(*A.dense()), *x.dense());
            else if (numA == 2)
              noalias(*y.dense()) += ublas::prod(trans(*A.triang()), *x.dense());
            else if (numA == 3)
              noalias(*y.dense()) += ublas::prod(trans(*A.sym()), *x.dense());
            else if (numA == 4)
              noalias(*y.dense()) += ublas::prod(trans(*A.sparse()), *x.dense());
            else //if(numA==5)
              noalias(*y.dense()) += ublas::prod(trans(*A.banded()), *x.dense());
          }
          else //if(numX == 4)
          {
            if (numY != 1 && numA != 4)
              SiconosMatrixException::selfThrow("prod(x,A,y) error: y (output) must be a dense vector.");

            if (numA == 1)
              noalias(*y.dense()) += ublas::prod(trans(*A.dense()), *x.sparse());
            else if (numA == 2)
              noalias(*y.dense()) += ublas::prod(trans(*A.triang()), *x.sparse());
            else if (numA == 3)
              noalias(*y.dense()) += ublas::prod(trans(*A.sym()), *x.sparse());
            else if (numA == 4)
            {
              if (numY == 1)
                noalias(*y.dense()) += ublas::prod(trans(*A.sparse()), *x.sparse());
              else
                noalias(*y.sparse()) += ublas::prod(trans(*A.sparse()), *x.sparse());
            }
            else //if(numA==5)
              noalias(*y.dense()) += ublas::prod(trans(*A.banded()), *x.sparse());
          }
        }
        else // if x and y are the same object => alias
        {
          if (numX == 1)
          {
            if (numA == 1)
              *y.dense() += ublas::prod(trans(*A.dense()), *x.dense());
            else if (numA == 2)
              *y.dense() += ublas::prod(trans(*A.triang()), *x.dense());
            else if (numA == 3)
              *y.dense() += ublas::prod(trans(*A.sym()), *x.dense());
            else if (numA == 4)
              *y.dense() += ublas::prod(trans(*A.sparse()), *x.dense());
            else //if(numA==5)
              *y.dense() += ublas::prod(trans(*A.banded()), *x.dense());
          }
          else //if(numX == 4)
          {
            if (numA == 1)
              *y.sparse() += ublas::prod(trans(*A.dense()), *x.sparse());
            else if (numA == 2)
              *y.sparse() += ublas::prod(trans(*A.triang()), *x.sparse());
            else if (numA == 3)
              *y.sparse() += ublas::prod(trans(*A.sym()), *x.sparse());
            else if (numA == 4)
              *y.sparse() += ublas::prod(trans(*A.sparse()), *x.sparse());
            else //if(numA==5)
              *y.sparse() += ublas::prod(trans(*A.banded()), *x.sparse());
          }
        }
      }
    }
  }
}

void prod(const SiconosMatrix& A, const SiconosVector& x, SiconosVector& y, bool init)
{
  // To compute y = A * x in an "optimized" way (in comparison with y = prod(A,x) )
  // or y += A*x if init = false.
  assert(!(A.isPLUFactorized()) && "A is PLUFactorized in prod !!");

  if (A.size(1) != x.size())
    SiconosMatrixException::selfThrow("prod(A,x,y) error: inconsistent sizes between A and x.");

  if (A.size(0) != y.size())
    SiconosMatrixException::selfThrow("prod(A,x,y) error: inconsistent sizes between A and y.");

  unsigned int numA = A.num();
  unsigned int numX = x.num();
  unsigned int numY = y.num();

  if (numA == 0) // If A is Block
    SiconosMatrixException::selfThrow("prod(A,x,y) error: not yet implemented for block matrices.");

  if (numA == 6) // A = 0
  {
    if (init)
      y.zero();
    //else nothing
  }

  else if (numA == 7) // A = identity
  {
    if (!init)
      y += x;
    else
    {
      if (&x != &y) y = x ; // if x and y do not share memory (ie are different objects)
      // else nothing
    }
  }

  else // A is not 0 or identity
  {

    // === First case: y is not a block vector ===
    if (init)
    {
      if (&x != &y) // if no common memory between x and y.
      {
        if (numX == 1)
        {
          if (numY != 1)
            SiconosMatrixException::selfThrow("prod(A,x,y) error: y (output) must be a dense vector.");

          assert(y.dense() != x.dense());

          if (numA == 1)
            noalias(*y.dense()) = ublas::prod(*A.dense(), *x.dense());
          else if (numA == 2)
            noalias(*y.dense()) = ublas::prod(*A.triang(), *x.dense());
          else if (numA == 3)
            noalias(*y.dense()) = ublas::prod(*A.sym(), *x.dense());
          else if (numA == 4)
            noalias(*y.dense()) = ublas::prod(*A.sparse(), *x.dense());
          else //if(numA==5)
            noalias(*y.dense()) = ublas::prod(*A.banded(), *x.dense());
        }
        else //if(numX == 4)
        {
          if (numY != 1 && numA != 4)
            SiconosMatrixException::selfThrow("prod(A,x,y) error: y (output) must be a dense vector.");

          if (numA == 1)
            noalias(*y.dense()) = ublas::prod(*A.dense(), *x.sparse());
          else if (numA == 2)
            noalias(*y.dense()) = ublas::prod(*A.triang(), *x.sparse());
          else if (numA == 3)
            noalias(*y.dense()) = ublas::prod(*A.sym(), *x.sparse());
          else if (numA == 4)
          {
            if (numY == 1)
              noalias(*y.dense()) = ublas::prod(*A.sparse(), *x.sparse());
            else
              noalias(*y.sparse()) = ublas::prod(*A.sparse(), *x.sparse());
          }
          else //if(numA==5)
            noalias(*y.dense()) = ublas::prod(*A.banded(), *x.sparse());
        }
      }
      else // if x and y are the same object => alias
      {
        if (numX == 1)
        {
          if (numA == 1)
            *y.dense() = ublas::prod(*A.dense(), *x.dense());
          else if (numA == 2)
            *y.dense() = ublas::prod(*A.triang(), *x.dense());
          else if (numA == 3)
            *y.dense() = ublas::prod(*A.sym(), *x.dense());
          else if (numA == 4)
            *y.dense() = ublas::prod(*A.sparse(), *x.dense());
          else //if(numA==5)
            *y.dense() = ublas::prod(*A.banded(), *x.dense());
        }
        else //if(numX == 4)
        {
          if (numA == 1)
            *y.sparse() = ublas::prod(*A.dense(), *x.sparse());
          else if (numA == 2)
            *y.sparse() = ublas::prod(*A.triang(), *x.sparse());
          else if (numA == 3)
            *y.sparse() = ublas::prod(*A.sym(), *x.sparse());
          else if (numA == 4)
            *y.sparse() = ublas::prod(*A.sparse(), *x.sparse());
          else //if(numA==5)
            *y.sparse() = ublas::prod(*A.banded(), *x.sparse());
        }
      }
    }
    else // += case
    {
      if (&x != &y) // if no common memory between x and y.
      {
        if (numX == 1)
        {
          if (numY != 1)
            SiconosMatrixException::selfThrow("prod(A,x,y) error: y (output) must be a dense vector.");

          if (numA == 1)
            noalias(*y.dense()) += ublas::prod(*A.dense(), *x.dense());
          else if (numA == 2)
            noalias(*y.dense()) += ublas::prod(*A.triang(), *x.dense());
          else if (numA == 3)
            noalias(*y.dense()) += ublas::prod(*A.sym(), *x.dense());
          else if (numA == 4)
            noalias(*y.dense()) += ublas::prod(*A.sparse(), *x.dense());
          else //if(numA==5)
            noalias(*y.dense()) += ublas::prod(*A.banded(), *x.dense());
        }
        else //if(numX == 4)
        {
          if (numY != 1 && numA != 4)
            SiconosMatrixException::selfThrow("prod(A,x,y) error: y (output) must be a dense vector.");

          if (numA == 1)
            noalias(*y.dense()) += ublas::prod(*A.dense(), *x.sparse());
          else if (numA == 2)
            noalias(*y.dense()) += ublas::prod(*A.triang(), *x.sparse());
          else if (numA == 3)
            noalias(*y.dense()) += ublas::prod(*A.sym(), *x.sparse());
          else if (numA == 4)
          {
            if (numY == 1)
              noalias(*y.dense()) += ublas::prod(*A.sparse(), *x.sparse());
            else
              noalias(*y.sparse()) += ublas::prod(*A.sparse(), *x.sparse());
          }
          else //if(numA==5)
            noalias(*y.dense()) += ublas::prod(*A.banded(), *x.sparse());
        }
      }
      else // if x and y are the same object => alias
      {
        if (numX == 1)
        {
          if (numA == 1)
            *y.dense() += ublas::prod(*A.dense(), *x.dense());
          else if (numA == 2)
            *y.dense() += ublas::prod(*A.triang(), *x.dense());
          else if (numA == 3)
            *y.dense() += ublas::prod(*A.sym(), *x.dense());
          else if (numA == 4)
            *y.dense() += ublas::prod(*A.sparse(), *x.dense());
          else //if(numA==5)
            *y.dense() += ublas::prod(*A.banded(), *x.dense());
        }
        else //if(numX == 4)
        {
          if (numA == 1)
            *y.sparse() += ublas::prod(*A.dense(), *x.sparse());
          else if (numA == 2)
            *y.sparse() += ublas::prod(*A.triang(), *x.sparse());
          else if (numA == 3)
            *y.sparse() += ublas::prod(*A.sym(), *x.sparse());
          else if (numA == 4)
            *y.sparse() += ublas::prod(*A.sparse(), *x.sparse());
          else //if(numA==5)
            *y.sparse() += ublas::prod(*A.banded(), *x.sparse());
        }
      }
    }
  }
}

void axpy_prod(const SiconosMatrix& A, const SiconosVector& x, SiconosVector& y, bool init)
{
  // To compute y = A * x ( init = true) or y += A * x (init = false) using ublas::axpy_prod
  assert(!(A.isPLUFactorized()) && "A is PLUFactorized in prod !!");

  if (A.size(1) != x.size())
    SiconosMatrixException::selfThrow("prod(A,x,y) error: inconsistent sizes between A and x.");

  if (A.size(0) != y.size())
    SiconosMatrixException::selfThrow("prod(A,x,y) error: inconsistent sizes between A and y.");

  unsigned int numA = A.num();
  unsigned int numX = x.num();
  unsigned int numY = y.num();

  if (numA == 0) // If A is Block
    SiconosMatrixException::selfThrow("axpy_prod(A,x,y) error: not yet implemented for block matrices.");

  if (numA == 6) // A = 0
  {
    if (init) y.zero(); // else nothing ...
  }

  else if (numA == 7) // A = identity
  {
    if (!init) y += x;
    else
    {
      if (&x != &y)
        y = x ; // if x and y do not share memory (ie are different objects)
    }
    // else nothing
  }

  else // A is not 0 or identity
  {
    {
      {
        if (&x != &y) // if no common memory between x and y.
        {
          if (numX == 1)
          {
            if (numY != 1)
              SiconosMatrixException::selfThrow("prod(A,x,y) error: y (output) must be a dense vector.");

            if (numA == 1)
              ublas::axpy_prod(*A.dense(), *x.dense(), *y.dense(), init);
            else if (numA == 2)
              ublas::axpy_prod(*A.triang(), *x.dense(), *y.dense(), init);
            else if (numA == 3)
              ublas::axpy_prod(*A.sym(), *x.dense(), *y.dense(), init);
            else if (numA == 4)
              ublas::axpy_prod(*A.sparse(), *x.dense(), *y.dense(), init);
            else //if(numA==5)
              ublas::axpy_prod(*A.banded(), *x.dense(), *y.dense(), init);
          }
          else //if(numX == 4)
          {
            if (numY != 1 && numA != 4)
              SiconosMatrixException::selfThrow("axpy_prod(A,x,y) error: y (output) must be a dense vector.");

            if (numA == 1)
              ublas::axpy_prod(*A.dense(), *x.sparse(), *y.dense(), init);
            else if (numA == 2)
              ublas::axpy_prod(*A.triang(), *x.sparse(), *y.dense(), init);
            else if (numA == 3)
              ublas::axpy_prod(*A.sym(), *x.sparse(), *y.dense(), init);
            else if (numA == 4)
            {
              if (numY == 1)
                ublas::axpy_prod(*A.sparse(), *x.sparse(), *y.dense(), init);
              else
                ublas::axpy_prod(*A.sparse(), *x.sparse(), *y.sparse(), init);
            }
            else //if(numA==5)
              ublas::axpy_prod(*A.banded(), *x.sparse(), *y.dense(), init);
          }
        }
        else // if x and y are the same object => alias
        {
          if (numX == 1)
          {
            if (numA == 1)
              ublas::axpy_prod(*A.dense(), *x.dense(), *x.dense(), init);
            else if (numA == 2)
              ublas::axpy_prod(*A.triang(), *x.dense(), *x.dense(), init);
            else if (numA == 3)
              ublas::axpy_prod(*A.sym(), *x.dense(), *x.dense(), init);
            else if (numA == 4)
              ublas::axpy_prod(*A.sparse(), *x.dense(), *x.dense(), init);
            else //if(numA==5)
              ublas::axpy_prod(*A.banded(), *x.dense(), *x.dense(), init);
          }
          else //if(numX == 4)
          {
            if (numA == 1)
              ublas::axpy_prod(*A.dense(), *x.sparse(), *x.sparse(), init);
            else if (numA == 2)
              ublas::axpy_prod(*A.triang(), *x.sparse(), *x.sparse(), init);
            else if (numA == 3)
              ublas::axpy_prod(*A.sym(), *x.sparse(), *x.sparse(), init);
            else if (numA == 4)
              ublas::axpy_prod(*A.sparse(), *x.sparse(), *x.sparse(), init);
            else //if(numA==5)
              ublas::axpy_prod(*A.banded(), *x.sparse(), *x.sparse(), init);
          }
        }
      }
    }
  }
}

void gemvtranspose(double a, const SiconosMatrix& A, const SiconosVector& x, double b, SiconosVector& y)
{
  if (A.isBlock())
    SiconosMatrixException::selfThrow("gemv(...) not yet implemented for block vectors or matrices.");
  assert(!(A.isPLUFactorized()) && "A is PLUFactorized in prod !!");

  unsigned int numA = A.num();
  unsigned int numX = x.num();
  unsigned int numY = y.num();
  if (numA != 1 || numX != 1 || numY != 1)
    SiconosMatrixException::selfThrow("gemv(...) failed: reserved to dense matrices or vectors.");

  siconosBindings::blas::gemv(a, siconosBindings::trans(*A.dense()), *x.dense(), b, *y.dense());
}

void gemv(double a, const SiconosMatrix& A, const SiconosVector& x, double b, SiconosVector& y)
{
  if (A.isBlock())
    SiconosMatrixException::selfThrow("gemv(...) not yet implemented for block vectors or matrices.");
  assert(!(A.isPLUFactorized()) && "A is PLUFactorized in prod !!");

  unsigned int numA = A.num();
  unsigned int numX = x.num();
  unsigned int numY = y.num();
  if (numA != 1 || numX != 1 || numY != 1)
    SiconosMatrixException::selfThrow("gemv(...) failed: reserved to dense matrices or vectors.");

  siconosBindings::blas::gemv(a, *A.dense(), *x.dense(), b, *y.dense());
}


#pragma once
/**
 * @file kdtree.hpp
 * @author Thomas Grandits
 * @brief Header of the KD-Tree implementation (CPU & GPU)
 * @version 0.1
 * @date 2020-12-16
 *
 * @copyright Copyright (c) 2020
 *
 */

#include <iostream>     // std::cout
#include <iterator>     // std::back_inserter
#include <vector>       // std::vector
#include <algorithm>    // std::copy
#include <array>
#include <assert.h>     /* assert */
#include <cmath>
#include <tuple>
#include <numeric>
#include <functional>
#include "nndistance.hpp"

#include <Eigen/Dense>

//https://stackoverflow.com/questions/32014839/how-to-use-a-cuda-class-header-file-in-both-cpp-and-cuda-modules
#ifdef __CUDACC__
#define CUDA_HOSTDEV __host__ __device__
#else
#define CUDA_HOSTDEV
#endif

typedef uint8_t dim_t;
typedef uint32_t point_i_t;
typedef uint32_t point_i_knn_t;
typedef point_i_t tree_ind_t;
typedef int level_t;
typedef uint8_t slot_t;

template <typename T, dim_t dims>
using Vec = Eigen::Matrix<T, dims, 1>;

/**
 * @brief PingPongBuffer consisting of two actual buffers that can be easily accessed and swapped
 *
 * @tparam T Type of the PingPongBuffer
 */
template <typename T>
struct PingPongBuffer
{
    //std::array<T*, 2> buffers;
    T* buffers[2];
    slot_t current_slot = 0;

    CUDA_HOSTDEV PingPongBuffer(){}
    CUDA_HOSTDEV PingPongBuffer(T* ping, T* pong){buffers[0] = ping; buffers[1] = pong;}

    CUDA_HOSTDEV inline T* getCurrentSlot() { return buffers[current_slot]; }
    CUDA_HOSTDEV inline slot_t getCurrentSlotInd() const { return current_slot; }
    CUDA_HOSTDEV inline T* getPongSlot() { return buffers[(current_slot + 1) % 2]; }
    CUDA_HOSTDEV inline void increment()  { current_slot = (current_slot + 1) % 2; }
    CUDA_HOSTDEV inline T* getCurrentSlotAndIncrement() { T* buf = getCurrentSlot();  increment(); return buf;}
};

/**
 * @brief Computes the total number of nodes necessary for a KD-Tree with the given levels
 *
 * @tparam T Type of the levels
 * @param levels Number of levels of the tree
 * @return T Number of total nodes of the tree
 */
template <typename T>
inline T compTotalNrNodes(const T levels)
{
    std::vector<T> range(levels);
    std::iota(range.begin(), range.end(), 0);

    return std::accumulate(range.begin(), range.end(), 0, [](const T accum, const T cur_level){return accum + std::pow<T>(2, cur_level);});
}

/**
 * @brief Computes the total number of leaves necessary for a KD-Tree with the given levels
 *
 * @tparam T Type of the levels
 * @param levels Number of levels of the tree
 * @return CUDA_HOSTDEV T Number of total leaves of the tree
 */
template <typename T>
CUDA_HOSTDEV inline T compTotalNrLeaves(const T levels)
{
    return std::pow<T>(2, levels);
}

/**
 * @brief Enum to tag which nodes have been visited when traversing the tree
 *
 */
enum NodeTag : unsigned char
{
    uncharted = 0,
    //visited = 1,
    left_visited = 1,
    right_visited = 2,
    left_right_visited = left_visited | right_visited
};

/**
 * @brief Enum to specify in which direction an algorithm should continue traversing the tree
 *
 */
enum NodeDirection : unsigned char
{
    up = 0,
    left = 1,
    right = 2,
    finished = 3
};

/**
 * @brief Specifies data associated to each leaf of the tree
 *
 * @tparam T Type of the leaf
 * @tparam dims Dimensionality of the contained points
 */
template <typename T, dim_t dims>
struct PartitionLeaf
{
    std::array<T, dims>* data;
    point_i_t nr_points;
    point_i_t offset;

    PartitionLeaf(std::array<T, dims>* data_, const point_i_t nr_points_, const point_i_t offset_)
    : data(data_), nr_points(nr_points_), offset(offset_){}

    PartitionLeaf(){}
};

/**
 * @brief Partition info present at each node, denoting the split-axis and the median used for the partitioning
 *
 * @tparam T Type of the median
 */
template <typename T>
struct Partition //: PartitionAligned<T>
{
    dim_t axis_split;
    T median;

    Partition(const dim_t axis_split_, const T median_) : axis_split(axis_split_), median(median_){}
    Partition(){}
};

template <typename T>
using tree_visit_f = std::function<NodeDirection(const Partition<T>&, const NodeTag, const tree_ind_t, const level_t)>;

/**
 * @brief Contains all partition info associated with a KD-Tree:
 *        All partitions, all leaves, levels of the tree and the underyling structured points, as well as the shuffled indices are all managed here.
 *
 * @tparam T Type of the structured points and all underlying data
 * @tparam dims Dimensionality of the points
 * @tparam delete_partitions If true, deletes all associated partition data once the tree is destructed
 */
template <typename T, dim_t dims, bool delete_partitions = true>
struct PartitionInfo
{
    Partition<T>* partitions;
    PartitionLeaf<T, dims>* leaves;
    level_t levels;
    std::array<T, dims>* structured_points = NULL;
    point_i_t* shuffled_inds = NULL;
    const point_i_t nr_points = 0;
    const tree_ind_t nr_partitions, nr_leaves;

    PartitionInfo(std::vector<Partition<T>>&& parts, std::vector<PartitionLeaf<T, dims>>&& leaves_, point_i_t* shuffled_inds_, const point_i_t nr_points_);

    PartitionInfo(PartitionInfo&& to_move) : partitions(std::move(to_move.partitions)), leaves(std::move(to_move.leaves)), levels(to_move.levels),
        structured_points(to_move.structured_points), shuffled_inds(to_move.shuffled_inds), nr_points(to_move.nr_points),
        nr_partitions(to_move.nr_partitions), nr_leaves(to_move.nr_leaves)
    {
        to_move.partitions = NULL;
        to_move.leaves = NULL;
        to_move.structured_points = NULL;
        to_move.shuffled_inds = NULL;
    }

    PartitionInfo(const PartitionInfo<T, dims>& to_copy) : partitions(to_copy.partitions), leaves(to_copy.leaves), levels(to_copy.levels),
        structured_points(to_copy.structured_points), shuffled_inds(to_copy.shuffled_inds), nr_points(to_copy.nr_points),
        nr_partitions(to_copy.nr_partitions), nr_leaves(to_copy.nr_leaves)
    {
        if(delete_partitions)
            throw std::runtime_error("Copies are not meant to do this here");
    }

    CUDA_HOSTDEV ~PartitionInfo()
    {
        if(delete_partitions)
        {
            delete partitions;
            delete leaves;
        }
    }
};

template <typename T, dim_t dims>
using PartitionInfoDevice = PartitionInfo<T, dims, false>;

/**
 * @brief Helper class to effectively traverse the KD-tree without any recursions.
 *
 * @tparam T Type of the points
 * @tparam dims Dimensionality of the points
 */
template <typename T, dim_t dims>
class TreeTraversal
{
//protected:
public:
    PartitionInfo<T, dims>* partition_info = NULL;

    NodeTag* visited_info;
    tree_ind_t current_lin_ind;
    level_t current_level;
    point_i_t nr_nodes;

public:

    /**
     * @brief Resets all positions and tags, effectively reinitializing the object without the need to destruct it.
     *
     * @return CUDA_HOSTDEV
     */
    CUDA_HOSTDEV void resetPositionAndTags()
    {
        std::fill(visited_info, visited_info + nr_nodes, NodeTag::uncharted);
        current_lin_ind = current_level = 0;
    }

    TreeTraversal(PartitionInfo<T, dims>* partition_info_) : partition_info(partition_info_), nr_nodes(partition_info_->nr_partitions)
    {
        //visited_info = std::move(std::vector<NodeTag>(compTotalNrNodes(partition_info->levels), NodeTag::uncharted));
        visited_info = new NodeTag[nr_nodes];
        std::fill(visited_info, visited_info + nr_nodes, NodeTag::uncharted);
        resetPositionAndTags();
    }

    CUDA_HOSTDEV TreeTraversal(){}

    CUDA_HOSTDEV TreeTraversal(PartitionInfoDevice<T, dims>* partition_info_) :
    //TODO: This could cause serious problems down the road!!!
    partition_info(reinterpret_cast<PartitionInfo<T, dims>*>(partition_info_)),
    //partition_info_d(partition_info_),
    nr_nodes(partition_info_->nr_partitions)
    {
        //visited_info = std::move(std::vector<NodeTag>(compTotalNrNodes(partition_info->levels), NodeTag::uncharted));
        visited_info = new NodeTag[nr_nodes];
        std::fill(visited_info, visited_info + nr_nodes, NodeTag::uncharted);
        current_lin_ind = current_level = 0;
    }

    CUDA_HOSTDEV ~TreeTraversal()
    {
        delete visited_info;
    }

    template <typename T_ind>
    static T_ind compParentInd(const T_ind lin_ind) { return (lin_ind - 1)/2; }

    template <typename T_ind>
    static T_ind compLeftChildInd(const T_ind lin_ind) { return lin_ind * 2 + 1; }

    template <typename T_ind>
    static T_ind compRightChildInd(const T_ind lin_ind) { return lin_ind * 2 + 2; }

    template <typename T_ind>
    static T_ind compLevel(T_ind lin_ind)
    {
        if(lin_ind == 0)
            return 1;

        T_ind level = 0;
        while((lin_ind = compParentInd(lin_ind)) != 0)
            level++;

        return level + 2;
    }

    template <typename T_ind, typename T_ind2>
    static inline T_ind compLeftLeafInd(const T_ind lin_ind, const T_ind2 nr_partitions) { return compLeftChildInd(lin_ind) - nr_partitions; }

    template <typename T_ind, typename T_ind2>
    static inline T_ind compRightLeafInd(const T_ind lin_ind, const T_ind2 nr_partitions) { return compLeftLeafInd(lin_ind, nr_partitions) + 1; }

    template <typename T_ind>
    static inline T_ind compLeftLeafInd(const T_ind lin_ind) { return compLeftLeafInd(lin_ind, compTotalNrNodes(compLevel(lin_ind))); }

    template <typename T_ind>
    static inline T_ind compRightLeafInd(const T_ind lin_ind) { return compLeftLeafInd(lin_ind) + 1; }

    CUDA_HOSTDEV inline tree_ind_t compLeftLeafInd(){ return compLeftLeafInd(current_lin_ind); }
    CUDA_HOSTDEV inline tree_ind_t compRighLeafInd(){ return compRightLeafInd(current_lin_ind); }
    CUDA_HOSTDEV inline level_t getTotalLevels() const {return partition_info->levels;}
    CUDA_HOSTDEV inline tree_ind_t getCurrentLinearIndex() const {return current_lin_ind;}
    CUDA_HOSTDEV inline level_t getCurrentLevel() const {return current_level;}
    CUDA_HOSTDEV inline NodeTag getCurrentTag() const {return visited_info[current_lin_ind];}
    CUDA_HOSTDEV //Partition<T>& getCurrentPartition() {return partition_info->partitions[current_lin_ind];}
    CUDA_HOSTDEV inline const Partition<T>& getCurrentConstPartition() const {return partition_info->partitions[current_lin_ind];}
    CUDA_HOSTDEV inline void moveToParent(){ assert(current_level != 0 && current_lin_ind != 0); current_lin_ind = compParentInd(current_lin_ind); current_level -= 1; }
    CUDA_HOSTDEV inline void moveToLeftChild(){ assert(current_level < partition_info->levels);  current_lin_ind = compLeftChildInd(current_lin_ind); current_level += 1; }
    CUDA_HOSTDEV inline void moveToRightChild(){ assert(current_level < partition_info->levels);  current_lin_ind = compRightChildInd(current_lin_ind); current_level += 1; }
    CUDA_HOSTDEV inline bool isLeafParent() const { return current_level == partition_info->levels-1; }
    CUDA_HOSTDEV inline const PartitionLeaf<T, dims>& getLeftLeaf() const{ assert(isLeafParent()); return partition_info->leaves[compLeftLeafInd(current_lin_ind, partition_info->nr_partitions)]; }
    CUDA_HOSTDEV inline const PartitionLeaf<T, dims>& getRightLeaf() const{ assert(isLeafParent()); return partition_info->leaves[compRightLeafInd(current_lin_ind, partition_info->nr_partitions)]; }
    CUDA_HOSTDEV inline void setCurrentNodeTag(const NodeTag new_tag){ visited_info[current_lin_ind] = new_tag; }

    static_assert(NodeTag::left_right_visited == 3 && NodeTag::uncharted == 0, "The binary computation does not work without these values");
    CUDA_HOSTDEV inline NodeTag moveTagToBinaryPosition(const NodeTag tag) const
    {
        const unsigned char bit_idx  = current_lin_ind % 4;
        return static_cast<NodeTag>(tag << (bit_idx*2));
    }
    CUDA_HOSTDEV inline NodeTag getBinaryIndMask() const { return moveTagToBinaryPosition(NodeTag::left_right_visited);}
    CUDA_HOSTDEV inline NodeTag getCurrentTagBinary() const
    {
        const auto byte_idx = current_lin_ind / 4;
        const unsigned char bit_idx  = current_lin_ind % 4;
        const unsigned char byte_data = visited_info[byte_idx];
        const unsigned char binary_bitmask = NodeTag::left_right_visited << (bit_idx*2);
        const unsigned char data = ( (byte_data & binary_bitmask) >> (bit_idx*2)) & NodeTag::left_right_visited;

        return static_cast<NodeTag>(data);
    }
    CUDA_HOSTDEV inline void setCurrentNodeTagBinary(const NodeTag new_bits)
    {
        const auto byte_idx = current_lin_ind / 4;
        const NodeTag old_byte = visited_info[byte_idx];
        const NodeTag modified_mask = moveTagToBinaryPosition(NodeTag::left_right_visited); //getBinaryIndMask();
        const NodeTag unmodified_mask = static_cast<NodeTag>(~modified_mask);
        const unsigned char bit_idx  = current_lin_ind % 4;
        visited_info[byte_idx] = static_cast<NodeTag>((unmodified_mask & old_byte) | (moveTagToBinaryPosition(new_bits) & modified_mask));
    }

    /**
     * @brief Traverses the tree with which the object was constructed and calls a function at each node.
     *
     * @param direction_f Function to call at each node. This function will receive the current node along with some positional information.
     *                    The function has to return a direction in which the tree should be further traversed. Finished when the function
     *                    returns NodeDirection::finished.
     * @return CUDA_HOSTDEV
     */
    CUDA_HOSTDEV void traverseTree(const tree_visit_f<T>& direction_f)
    {
        NodeDirection next_dir;
        while((next_dir = direction_f(getCurrentConstPartition(), getCurrentTag(), current_lin_ind, current_level))
                        != NodeDirection::finished)
        {
            switch(next_dir)
            {
                case NodeDirection::up:
                    moveToParent(); break;
                case NodeDirection::left:
                    moveToLeftChild(); break;
                case NodeDirection::right:
                    moveToRightChild(); break;
            }
        }
    }
};

/**
 * @brief Computes the median along given axis for the ordered points
 *
 * @tparam T Type of the points
 * @tparam dims Dimensionality of the points
 * @param points_ordered Points for which the median should be computed. Note that the points are required to be ordered
 *                       in the respective dimension. This should be ensured after the tree was already constructed.
 * @param nr_points Number of points present in the array
 * @param current_axis Axis for which the median should be computed
 * @return T The median
 */
template <typename T, dim_t dims>
inline T compMedian(const std::array<T, dims>* points_ordered, const point_i_t nr_points, const dim_t current_axis)
{
    assert(nr_points > 1);
    const point_i_t half_ind = nr_points / 2;
    if(nr_points % 2 == 1)
        return points_ordered[half_ind][current_axis];
    else
        return 0.5 * (points_ordered[half_ind][current_axis] + points_ordered[half_ind-1][current_axis]);
}

/**
 * @brief Creates the KD-Tree's partitions using a recursive algorithms
 *
 * @tparam T Type of the points
 * @tparam dims Dimensionality of the points
 * @param partitions Vector of already generated (possibly empty) partitions, which will be initialized
 * @param leaves Empty vector of leaves that will be generated when creating the KD-tree
 * @param lin_ind Current node index (initially 0)
 * @param structured_points Pointer to the array of underlying ordered points
 * @param shuffled_inds Pointer to the array of underlying indices of the ordered points
 * @param nr_points Number of total points
 * @param levels Levels of the final KD-Tree
 * @param current_axis Current dimension in which we will split (initially 0)
 * @param arr_offset Current recursive offset of in the structured points array (initially 0)
 */
template <typename T, dim_t dims>
void createPartitionRecursive(
                                std::vector<Partition<T>>& partitions,
                                std::vector<PartitionLeaf<T, dims>>& leaves,
                                const tree_ind_t lin_ind,
                                std::array<T, dims>* structured_points,
                                point_i_t* shuffled_inds, const point_i_t nr_points,
                                const int levels, const dim_t current_axis, const point_i_t arr_offset)
{
    typedef TreeTraversal<T, dims> tree_t;
    if(nr_points == 0)
        throw std::runtime_error("Error: Ran out of points while building KD-Tree. Either you required too many levels, or you used a lot of coplanar points");

    //Track global indices
    std::vector<point_i_t> idx(nr_points);
    std::iota(idx.begin(), idx.end(), 0);
    std::sort(idx.begin(), idx.end(),
    [current_axis, &structured_points](const point_i_t lhs, const point_i_t rhs)
    { return structured_points[lhs][current_axis] < structured_points[rhs][current_axis]; });

    //Reorder in the original
    //Work on copy and then copy back
    std::vector<std::array<T, dims>> structured_points_copy;
    std::vector<point_i_t> shuffle_copy(shuffled_inds, shuffled_inds + nr_points);
    for(size_t elem_i = 0; elem_i < nr_points; elem_i++)
        shuffled_inds[elem_i] = shuffle_copy[idx[elem_i]];

    std::sort(structured_points, structured_points + nr_points,
    [current_axis](const std::array<T, dims>& lhs, const std::array<T, dims>& rhs){return lhs[current_axis] < rhs[current_axis];});

    const T median = compMedian<T, dims>(structured_points, nr_points, current_axis);

    std::array<T, dims>* points_lower = structured_points;
    std::array<T, dims>* points_higher = structured_points + nr_points/2;
    const point_i_t nr_points_lower = nr_points/2;
    const point_i_t nr_points_higher = nr_points - nr_points_lower; //TODO: For duplicates, this may cause alterations...

    assert((nr_points - nr_points_lower) > 0);

    if(levels == 0)
    {
        assert(tree_t::compLeftLeafInd(lin_ind) == leaves.size());
        leaves.push_back(std::move(PartitionLeaf<T, dims>(points_lower, nr_points_lower, arr_offset)));

        assert(tree_t::compRightLeafInd(lin_ind) == leaves.size());
        leaves.push_back(std::move(PartitionLeaf<T, dims>(points_higher, nr_points_higher, arr_offset + nr_points_lower)));
        partitions[lin_ind] = std::move(Partition<T>(current_axis, median));
    }
    else
    {
        createPartitionRecursive<T, dims>(
            partitions, leaves, TreeTraversal<T, dims>::compLeftChildInd(lin_ind),
            points_lower, shuffled_inds, nr_points_lower, levels-1,
                (current_axis + 1) % dims, arr_offset);
        createPartitionRecursive<T, dims>(
            partitions, leaves, TreeTraversal<T, dims>::compRightChildInd(lin_ind),
            points_higher, shuffled_inds + nr_points_lower, nr_points_higher, levels-1,
                (current_axis + 1) % dims, arr_offset + nr_points_lower);

        Partition<T> partition(current_axis, median);
        partitions[lin_ind] = std::move(partition);
    }
}

/**
 * @brief Recursively generates a KD-Tree using \ref createPartitionRecursive 'createPartitionRecursive'
 *
 * @tparam T Type of the points
 * @tparam dims Dimensionality of the points
 * @param points_flat Array of points
 * @param nr_points Number of points
 * @param levels KD-Tree levels
 * @return PartitionInfo<T, dims> The PartitionInfo containing the information regarding the KD-Tree
 */
template <typename T, dim_t dims>
PartitionInfo<T, dims> createKDTree(const T* points_flat,
                                        const point_i_t nr_points, const int levels)
{
    assert(levels >= 1);
    std::vector<PartitionLeaf<T, dims>> leaves; //(compTotalNrLeaves(levels));
    std::vector<Partition<T>> partitions(compTotalNrNodes(levels), Partition<T>(-1, -1));

    point_i_t* orig_inds = new point_i_t[nr_points];
    std::iota(orig_inds, orig_inds + nr_points, 0); //Initializing
    //std::cout << "Creating main partition: " << nr_points << ", " << levels << std::endl;
    const std::array<T, dims>* points = reinterpret_cast<const std::array<T, dims>*>(points_flat);
    //std::vector<std::array<T, dims>> structured_points(points, points + nr_points);
    std::array<T, dims>* structured_points = new std::array<T, dims>[nr_points];
    std::copy(points, points + nr_points, structured_points);
    createPartitionRecursive<T, dims>(partitions, leaves, 0,
                                        structured_points,
                                        orig_inds, nr_points, levels-1, 0, 0);
    //partitions[0] = new MainPartition<T, dims>(std::move(partitions[0]));
    /*MainPartition<T, dims> main_partition = MainPartition<T, dims>(std::move(
                        ),
                                        structured_points, orig_inds, nr_points, levels);*/

    //Sanity checks
    #ifndef NDEBUG
    point_i_t offset = 0;
    for(int i = 0; i < leaves.size(); i++)
    {
        assert(leaves[i].offset == offset);
        offset += leaves[i].nr_points;
    }
    for(int i = 1; i < leaves.size(); i++)
    {
        assert((leaves[i-1].data + leaves[i-1].nr_points) == leaves[i].data);
    }
    #endif

    assert(leaves.size() == compTotalNrLeaves(levels));

    return std::move(PartitionInfo<T, dims>(std::move(partitions), std::move(leaves), /*structured_points,*/ orig_inds, nr_points)); //std::move(main_partition);
}

/**
 * @brief Computes the difference between two points (lhs - rhs)
 *
 * @tparam T Type of the points
 * @tparam T_calc Type of the resulting difference
 * @tparam dims Dimensionality of the points
 * @param lhs Array of points
 * @param rhs Array of points
 * @return std::array<T_calc, dims> The difference between the lhs and rhs
 */
template
<typename T, typename T_calc, dim_t dims>
inline CUDA_HOSTDEV std::array<T_calc, dims> compDiff(const std::array<T, dims>& lhs, const std::array<T, dims>& rhs)
{
    std::array<T_calc, dims> diffs;

    for(dim_t dim_i = 0; dim_i < dims; dim_i++)
    {
        diffs[dim_i] = static_cast<T_calc>(lhs[dim_i]) - static_cast<T_calc>(rhs[dim_i]);
    }

    return std::move(diffs);
}

/**
 * @brief Computes the sum of squared elem
 *
 * @tparam T Type of the array
 * @tparam dims Number of points
 * @param x Array of values
 * @return T sum(x^2)
 */
template
<typename T, dim_t dims>
inline CUDA_HOSTDEV T compQuadrSum(const std::array<T, dims>& x)
{
    return std::accumulate(x.begin(), x.end(), 0., [](const T accum, const T elem){return accum + elem*elem;});
}

/**
 * @brief Computes the element-wise quadratic euclidean distance between lhs and rhs, i.e. ||lhs_i - rhs_i||
 *
 * @tparam T Type of the array
 * @tparam T_calc Type of the calculation
 * @tparam dims Dimensionality of the points
 * @param lhs Left-hand-side of the calculation
 * @param lhs Right-hand-side of the calculation
 * @return T_calc ||lhs_i - rhs_i||
 */
template
<typename T, typename T_calc, dim_t dims>
inline CUDA_HOSTDEV T_calc compQuadrDist(const std::array<T, dims>& lhs, const std::array<T, dims>& rhs)
{
    std::array<T_calc, dims> diffs = compDiff<T, T_calc, dims>(lhs, rhs);
    return compQuadrSum<T_calc, dims>(diffs);
}

/**
 * @brief Computes the distance to the projected point onto the median of the current dimension. This is used to see if we need to further
 *        traverse the tree or already found the KNNs, which is true in case the projection is farther away than
 *        the distance to the current KNN most far away.
 *
 * @tparam T Type of the array
 * @tparam dims Dimensionality of the points
 * @param point Original point for which to search the KNN
 * @param point_proj The original point, already projected on possible multiple previous dimensions. Note that in the simplest case
 *                   this is equal to point.
 * @param Partition<T> Node information of the current axis/dimension split that we want to compute the projection in
 * @return T Resulting distance of the current projection
 */
template
<typename T, dim_t dims>
inline CUDA_HOSTDEV T projectionDist(const std::array<T, dims>& point, const std::array<T, dims>& point_proj, const Partition<T>& partition)
{
    std::array<T, dims> proj_vec = compDiff<T, T, dims>(point_proj, point);
    const dim_t current_axis = partition.axis_split;
    proj_vec[current_axis] += partition.median - point_proj[current_axis];
    return compQuadrSum<T, dims>(proj_vec);
}

/**
 * @brief See \ref projectionDist 'projectionDist'
 */
template
<typename T, dim_t dims>
inline CUDA_HOSTDEV T projectionDist(const Vec<T, dims>& point, const Vec<T, dims>& point_proj, const Partition<T>& partition)
{
    Vec<T, dims> proj_vec = (point_proj - point);
    const dim_t current_axis = partition.axis_split;
    proj_vec[current_axis] += partition.median - point_proj[current_axis];
    return proj_vec.squaredNorm();
}

/**
 * @brief Checks if we need to compute the distances to any of the points in the partition. Achieved by using \ref projectionDist 'projectionDist'.
 */
template
<typename T, typename T_calc, dim_t dims>
inline CUDA_HOSTDEV bool partitionNecessary(const std::array<T, dims>& point, const std::array<T, dims>& point_proj,
                                const Partition<T>& partition, const T current_worst_dist)
{
    const auto proj_dist = projectionDist<T, dims>(point, point_proj, partition);
    return proj_dist < current_worst_dist;
}

/**
 * @brief See \ref partitionNecessary 'partitionNecessary'.
 */
template
<typename T, typename T_calc, dim_t dims>
inline CUDA_HOSTDEV bool partitionNecessary(const Vec<T, dims>& point, const Vec<T, dims>& point_proj,
                                const Partition<T>& partition, const T current_worst_dist)
{
    const auto proj_dist = projectionDist<T, dims>(point, point_proj, partition);
    return proj_dist < current_worst_dist;
}

/**
 * @brief Computes the quadratic distances between point to all points in partition_leaf and updates the current knn, marked by best_dists and best_idx.
 *
 * @param point Query point for which the KNN should be calculated
 * @param partition_leaf Partition for which all quadratic distances should be computed
 * @param best_dists Quadr. distances to the currently assumed KNNs (starts with infinity). Has a length of K
 * @param best_idx Indices of the currently assumed KNNs (starts with -1). Has a length of K
 * @param nr_nns_searches How many nearest neighbors will be searched (=K)
 *
 */
template
<typename T, typename T_calc, dim_t dims>
void compQuadrDistLeafPartition(const std::array<T, dims>& point, const PartitionLeaf<T, dims>& partition_leaf,
                                    T* best_dists, point_i_knn_t* best_idx,
									const point_i_knn_t nr_nns_searches)
{
	const std::array<T, dims>* partition_data = partition_leaf.data;
    const point_i_t partition_size = partition_leaf.nr_points;
	const point_i_t partition_offset = partition_leaf.offset;
    for(point_i_t ref_i = 0; ref_i < partition_size; ref_i++)
    {
        const T_calc dist = compQuadrDist<T, T_calc, dims>(point, partition_data[ref_i]);
        const auto insertion_idx = knnInsertionDynamic<T_calc>(dist, best_dists, nr_nns_searches);
        if(insertion_idx < nr_nns_searches)
        {
            const auto best_dists_end = best_dists + nr_nns_searches;
            const auto best_idx_end = best_idx + nr_nns_searches;
            assert(best_dists + insertion_idx < best_dists_end);
            assert(best_idx + insertion_idx < best_idx_end);
            //Shift elements to the right
            //std::cout << k << ", " << d << "insert into " << insertion_idx << std::endl;
            //std::move_backward(best_dists + insertion_idx, best_dists +  nr_nns_searches - 1, best_dists.end());
            //std::move_backward(best_idx.begin() + insertion_idx, best_idx.end() - 1, best_idx.end());
            moveBackward(best_dists + insertion_idx, best_dists_end - 1, best_dists_end);
            moveBackward(best_idx + insertion_idx, best_idx_end - 1, best_idx_end);
            best_dists[insertion_idx] = dist;
            best_idx[insertion_idx] = ref_i + partition_offset;
        }
    }
}

/**
 * @brief See \ref compQuadrDistLeafPartition 'compQuadrDistLeafPartition'.
 */
template
<typename T, typename T_calc, dim_t dims>
CUDA_HOSTDEV void compQuadrDistLeafPartition(const Vec<T, dims>& point, const PartitionLeaf<T, dims>& partition_leaf,
                                    T* best_dists, point_i_knn_t* best_idx,
									const point_i_knn_t nr_nns_searches)
{
    //printf("compQuadrDistLeafPartition: %x, ", partition_leaf.data);
    //printf("%d, ", partition_leaf.nr_points);
    //printf("%d\n", partition_leaf.offset);
	const Vec<T, dims>* partition_data = reinterpret_cast<Vec<T, dims>*>(partition_leaf.data);
    const point_i_t partition_size = partition_leaf.nr_points;
	const point_i_t partition_offset = partition_leaf.offset;
    for(point_i_t ref_i = 0; ref_i < partition_size; ref_i++)
    {
        const T_calc dist = (point - partition_data[ref_i]).squaredNorm();
        const auto insertion_idx = knnInsertionDynamic<T_calc>(dist, best_dists, nr_nns_searches);
        if(insertion_idx < nr_nns_searches)
        {
            const auto best_dists_end = best_dists + nr_nns_searches;
            const auto best_idx_end = best_idx + nr_nns_searches;
            assert(best_dists + insertion_idx < best_dists_end);
            assert(best_idx + insertion_idx < best_idx_end);
            //Shift elements to the right
            //std::cout << k << ", " << d << "insert into " << insertion_idx << std::endl;
            //std::move_backward(best_dists + insertion_idx, best_dists +  nr_nns_searches - 1, best_dists.end());
            //std::move_backward(best_idx.begin() + insertion_idx, best_idx.end() - 1, best_idx.end());
            moveBackward(best_dists + insertion_idx, best_dists_end - 1, best_dists_end);
            moveBackward(best_idx + insertion_idx, best_idx_end - 1, best_idx_end);
            best_dists[insertion_idx] = dist;
            best_idx[insertion_idx] = ref_i + partition_offset;
        }
    }
}

/**
 * @brief See \ref compQuadrDistLeafPartition 'compQuadrDistLeafPartition'.
 */
template
<typename T, typename T_calc, dim_t dims>
void compQuadrDistPartition(const std::array<T, dims>& point, const PartitionLeaf<T, dims>& partition_leaf,
                                    std::vector<T>& best_dists, std::vector<point_i_knn_t>& best_idx,
									const point_i_knn_t nr_nns_searches);

/**
 * @brief Recursively traverses the node and calls compQuadrDistLeafPartition on the leaves.
 */
template
<typename T, typename T_calc, dim_t dims>
void compQuadrDistNode(const std::array<T, dims>& point, const Partition<T>& partition,
                                    std::vector<T>& best_dists, std::vector<point_i_knn_t>& best_idx,
                                    const int current_level, const std::array<T, dims>& point_proj, const point_i_knn_t nr_nns_searches);

/**
 * @brief Computes the KNN on the KD-Tree defined by partition_info. Iteratively calls compQuadrDistNode on the main node for each point in points_query.
 *        The resulting quadr. distances and indices of the KNNs will be stored in dist [M, K] in a C-ordering fashion.
 */
template
<typename T, typename T_calc, dim_t dims>
void KDTreeKNNSearch(PartitionInfo<T, dims>& partition_info,
                    const point_i_t nr_query,
                    const std::array<T, dims>* points_query, T * dist, point_i_knn_t* idx, const point_i_knn_t nr_nns_searches);


/* main.cpp - a "main" file, just a debugging tool
 *
 * Copyright (C) 2021, LWE-PVSS
 *
 * 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.
 **/

// This file is just a convenience, a handy tool that lets us run
// small porgrams without having to use the awkward ctest syntax.
#include <cassert>
#include <cmath>
#include <random>
#include <chrono>
#include <string>
#include <sys/time.h>
#include <sys/resource.h>
using namespace std;

#include <NTL/version.h>
#include "regevEnc.hpp"
#include "regevProofs.hpp"

using namespace ALGEBRA;
using namespace REGEVENC;

int main(int argc, char** argv) {
    // std::cout << "- Found GMP version "<<__GNU_MP__ <<std::endl;
    // std::cout << "- Found NTL version "<<NTL_VERSION <<std::endl;
    // std::cout << "- Found Sodium version "<<SODIUM_VERSION_STRING<<std::endl;

    int nParties = 512;
    if (argc > 1) {
        nParties = std::stoi(argv[1]);
    }
    if (nParties < 32 || nParties > 4096)
        nParties = 512;
    std::cout << "nParties="<<nParties << std::endl;

    // The dimensions of the the CRX is k-by-m, but note that this is a
    // matrix over GF(p^2) so the lattice dimensions we get is twice that
    KeyParams kp(nParties);
    // kp.k=64;
    GlobalKey gpk("testContext", kp);
    std::cout <<"{ kay:"<<gpk.kay<<", enn:"<<gpk.enn
      <<", sigmaEnc1:"<<gpk.sigmaEnc1<<", sigmaEnc2:"<<gpk.sigmaEnc2<<" }\n";

    TernaryEMatrix::init();
    MerlinRegev mer;
    PedersenContext ped;
    SharingParams ssp(interval(1,gpk.enn+1), gpk.tee);
    VerifierData vd(gpk, ped, mer, ssp);
    ProverData pd(vd);

    // Generate/verify the proofs by the second party (idx=1)
    int partyIdx = 1;

    // Key generation for all the parties
    std::vector<ALGEBRA::EVector> kgNoise(gpk.enn);
    std::vector<ALGEBRA::EVector> sk(gpk.enn);
    std::vector<ALGEBRA::EVector> pk(gpk.enn);
    auto start = chrono::steady_clock::now();
    crsTicks = 0;
    for (int i=0; i<gpk.enn; i++) {
        std::tie(sk[i],pk[i]) = gpk.genKeys(&kgNoise[i]);
        gpk.addPK(pk[i]);
    }
    gpk.setKeyHash();
    auto end = chrono::steady_clock::now();
    auto ticks = chrono::duration_cast<chrono::milliseconds>(end - start).count();
    std::cout <<gpk.enn<<" keyGens in "<<ticks<<" milliseconds, avg="<<(ticks/double(gpk.enn))
        << " ("<< (crsTicks/double(gpk.enn)) << " for s x A)\n";

    // encryption
    std::vector<ALGEBRA::SVector> ptxt1(gpk.enn);
    std::vector<GlobalKey::CtxtPair> ctxt1(gpk.enn);
    // secret sharing of a random value , the secret itself is sshr[0]
    ALGEBRA::SVector sshr;
    ssp.randomSharing(sshr);
    for (int i=0; i<gpk.enn; i++) {
        resize(ptxt1[i], gpk.enn);
        for (int j=0; j<gpk.enn; j++) ptxt1[i][j] = sshr[i+1];
    }
    start = chrono::steady_clock::now();
    crsTicks = 0;
    for (int i=0; i<gpk.enn; i++) {
        ctxt1[i] = gpk.encrypt(ptxt1[i]);
    }
    end = chrono::steady_clock::now();
    ticks = chrono::duration_cast<chrono::milliseconds>(end - start).count();
    std::cout <<gpk.enn<<" encryptions in "<<ticks<<" milliseconds, avg="<<(ticks/double(gpk.enn))
        << " ("<< (crsTicks/double(gpk.enn)) << " for A x r)\n";

    // decryption at party #1
    ALGEBRA::SVector ptxt2;    resize(ptxt2, gpk.tee);
    ALGEBRA::EVector decNoise; resize(decNoise, gpk.tee);
    start = chrono::steady_clock::now();
    for (int i=0; i<gpk.tee; i++) { // decrypt 2nd entry in i'th ctxt
        ptxt2[i] = gpk.decrypt(sk[partyIdx], partyIdx, ctxt1[i], &(decNoise[i]));
    }
    end = chrono::steady_clock::now();
    ticks = chrono::duration_cast<chrono::milliseconds>(end - start).count();
    std::cout << gpk.tee << " decryptions in "<<ticks<< " milliseconds, avg="
        << (ticks/double(gpk.tee)) << std::endl;

    for (int i=0; i<gpk.tee; i++) { // decrypt 2nd entry in i'th ctxt
        if (ptxt2[i] != ptxt1[i][partyIdx])
            std::cout << "decryption error in "<<i<<"th ciphertext\n";
    }

    // re-encryption at party #1
    ALGEBRA::SVector ptxt3;
    resize(ptxt3, gpk.enn);
    for (int j=0; j<gpk.enn; j++) ptxt3[j] = sshr[j+1];
    ALGEBRA::EVector encRnd;
    REGEVENC::GlobalKey::CtxtPair eNoise;
    auto ctxt2 = gpk.encrypt(ptxt3, encRnd, eNoise);

    // Copy the first t ciphertexts into a k x t matrix and another t-vector
    EMatrix ctxtMat;
    resize(ctxtMat, gpk.kay, gpk.tee);
    EVector ctxtVec;
    resize(ctxtVec, gpk.tee);
    for (int i=0; i<gpk.tee; i++) {
        for (int j=0; j<gpk.kay; j++)
            ctxtMat[j][i] = ctxt1[i].first[j];
        ctxtVec[i] = ctxt1[i].second[partyIdx];
    }

    // prepare for proof, commit to the secret key
    DLPROOFS::Point::counter = 0;
    DLPROOFS::Point::timer = 0;
    int origSize = sk[partyIdx].length();

    vd.sk1Com = commit(sk[partyIdx], vd.sk1Idx, vd.Gs, pd.sk1Rnd);

    start = chrono::steady_clock::now();
    SVector lagrange = vd.sp->lagrangeCoeffs(interval(1,gpk.tee+1));

    crsTicks = 0;
    proveDecryption(pd, ctxtMat, ctxtVec, ptxt2, sk[partyIdx], decNoise);
    proveEncryption(pd, ctxt2.first, ctxt2.second, ptxt3, encRnd, eNoise.first, eNoise.second);
    proveKeyGen(pd, partyIdx, sk[partyIdx], kgNoise[partyIdx]);
    proveReShare(pd, lagrange, ptxt2, ptxt3);
    proveSmallness(pd);

    end = chrono::steady_clock::now();
    ticks = chrono::duration_cast<chrono::milliseconds>(end - start).count();
    std::cout << "preparing to prove and committing in "<<ticks<< " milliseconds, "
        << DLPROOFS::Point::counter << " exponentiations in "
        << ((500+DLPROOFS::Point::timer)/1000) << " milliseconds"
        << " ("<< crsTicks << " for xR x A)\n";

    // aggregate the constraints and flatten everything before proving
    DLPROOFS::Point::counter = 0;
    DLPROOFS::Point::timer = 0;
    std::cout<<"aggregting constraints\n";
    start = chrono::steady_clock::now();
    ReadyToProve rtp;
    rtp.aggregateProver(pd);

    // Make copies of the Merlin transcripts and specialize them
    // for the final constraints before proving/verifying them
    auto merLin = *vd.mer;
    merLin.processConstraint("linear", rtp.linCnstr);

    auto merQuad = *vd.mer;
    merQuad.processConstraint("quadratic", rtp.quadCnstr);

    // Flatten the statements, this relases the memory of the constraints
    // (hence the Merlin processing above must be done before doing this).
    std::cout<<"flatenning constraints\n";
    rtp.flattenLinPrv(pd);
    rtp.flattenQuadPrv(pd);

    ReadyToVerify rtv = rtp; // a copy without the secret variables

    // prove and verify the linear statement
    auto merLinVer = merLin; // another copy for verification
    DLPROOFS::LinPfTranscript pfL("Linear");
    pfL.C = rtp.linCom;

    end = chrono::steady_clock::now();
    ticks = chrono::duration_cast<chrono::milliseconds>(end - start).count();
    std::cout << "aggregating constaints in "<<ticks<< " milliseconds, "
        << DLPROOFS::Point::counter << " exponentiations in "
        << ((500+DLPROOFS::Point::timer)/1000) << " milliseconds\n";

    // The actual proof
    DLPROOFS::Point::counter = 0;
    DLPROOFS::Point::timer = 0;
    start = chrono::steady_clock::now();
    DLPROOFS::proveLinear(pfL, rtp.lComRnd, merLin, rtp.linWtns.data(),
            rtp.linStmnt.data(), rtp.linGs.data(), rtp.linGs.size());
    end = chrono::steady_clock::now();
    ticks = chrono::duration_cast<chrono::milliseconds>(end - start).count();
    std::cout << "proving linear in "<<ticks<< " milliseconds, "
        << DLPROOFS::Point::counter << " exponentiations in "
        << ((500+DLPROOFS::Point::timer)/1000) << " milliseconds\n";

    DLPROOFS::Point::counter = 0;
    DLPROOFS::Point::timer = 0;
    start = chrono::steady_clock::now();
    if (!DLPROOFS::verifyLinear(pfL, rtv.linStmnt.data(), rtv.linGs.data(),
                      rtv.linGs.size(), rtv.linCnstr.equalsTo, merLinVer))
        std::cout << "failed linear verification\n";
    end = chrono::steady_clock::now();
    ticks = chrono::duration_cast<chrono::milliseconds>(end - start).count();
    std::cout << "verifying linear in "<<ticks<< " milliseconds, "
        << DLPROOFS::Point::counter << " exponentiations in "
        << ((500+DLPROOFS::Point::timer)/1000) << " milliseconds\n";

    // prove and verify the quadratic statement
    DLPROOFS::Point::counter = 0;
    DLPROOFS::Point::timer = 0;
    auto merQuadVer = merQuad; // another copy for verification
    DLPROOFS::QuadPfTranscript pfQ("Quadratic");
    pfQ.C = rtp.quadCom;
     // The actual proof
    start = chrono::steady_clock::now();
    DLPROOFS::proveQuadratic(pfQ, rtp.qComRnd, merQuad, rtp.quadGs.data(),
                rtp.quadWtnsG.data(), rtp.quadHs.data(), rtp.quadWtnsH.data(),
                rtp.quadGs.size());
    end = chrono::steady_clock::now();
    ticks = chrono::duration_cast<chrono::milliseconds>(end - start).count();
    std::cout << "proving quadratic in "<<ticks<< " milliseconds, "
        << DLPROOFS::Point::counter << " exponentiations in "
        << ((500+DLPROOFS::Point::timer)/1000) << " milliseconds\n";

    // The actual verification
    DLPROOFS::Point::counter = 0;
    DLPROOFS::Point::timer = 0;
    start = chrono::steady_clock::now();
    if (!DLPROOFS::verifyQuadratic(pfQ, rtv.quadGs.data(), rtv.quadHs.data(),
                        rtp.quadGs.size(), rtv.quadCnstr.equalsTo, merQuadVer,
                        rtv.offstG.data(), rtv.offstH.data()))
        std::cout << "failed quadratic verification\n";
    end = chrono::steady_clock::now();
    ticks = chrono::duration_cast<chrono::milliseconds>(end - start).count();
    std::cout << "verifying quadratic in "<<ticks<< " milliseconds, "
        << DLPROOFS::Point::counter << " exponentiations in "
        << ((500+DLPROOFS::Point::timer)/1000) << " milliseconds\n";

    struct rusage ru;
    getrusage(RUSAGE_SELF, &ru);
    std::cout << " max mem: " << ru.ru_maxrss << " kilobytes\n";

    return 0;
}

#if 0
#include <vector>
#include <iostream>
#include "bulletproof.hpp"

int main(int, char**) {
    constexpr size_t pfSize = 13;

    // build a constraint: sum_i ai*bi = b = \sum_bi^2
    DLPROOFS::LinConstraint cnstrL;
    for (size_t i=0; i<pfSize; i++) {
        CRV25519::Scalar& x = cnstrL.terms[i+1].setInteger(i+1);
        cnstrL.equalsTo += x * x;
    }
    DLPROOFS::PtxtVec& xes = cnstrL.terms;
    DLPROOFS::LinPfTranscript pfL = proveLinear("blah", cnstrL, xes);
    std::cout << "linear: "<<verifyLinear(cnstrL, pfL) <<std::endl;

    DLPROOFS::QuadConstraint cnstrQ;
    for (auto& x : xes) {
        cnstrQ.indexes.insert(cnstrQ.indexes.end(), x.first);
        cnstrQ.equalsTo += x.second * x.second;
    }
    DLPROOFS::PtxtVec& ys = cnstrL.terms;
    DLPROOFS::QuadPfTranscript pfQ = proveQuadratic("blah", cnstrQ, xes, ys);
    std::cout << "quadratic: "<<verifyQuadratic(cnstrQ, pfQ) <<std::endl;

    std::set<size_t> indexes;
    for (auto& elem: xes) // elem = {idx:scalar}
        indexes.insert(indexes.end(), elem.first);

    auto [normSq, prNS] = DLPROOFS::proveNormSquared("blah", xes);
    std::cout << "norm: "<<verifyNormSquared(indexes,normSq,prNS)<<std::endl;

    return 0;
}
#endif


// Copyright András Vukics 2006–2020. Distributed under the Boost Software License, Version 1.0. (See accompanying file LICENSE.txt)
#include "ModeCorrelations.h"

#include "LazyDensityOperator.h"
#include "Structure.h"

#include <boost/range/algorithm/copy.hpp>

using namespace boost;

ModeCorrelations::ModeCorrelations()
  : EA_Base("ModeCorrelations",{"<a number operator>","VAR(a number operator)","real(<a>)","imag(\")","<X^2>-<X>^2","<Y^2>-<Y>^2","<(XY+YX)/2>-<X><Y>",
                                "<b number operator>","VAR(b number operator)","real(<b>)","imag(\")","<Q^2>-<Q>^2","<P^2>-<P>^2","<(QP+PQ)/2>-<Q><P>",
                                "<XQ>-<X><Q>","<XP>-<X><P>","<YQ>-<Y><Q>","<YP>-<Y><P>"}
           ),
    averagedMode_()
{
}

namespace {

#include "details_BinaryHelper.h"

}

const ModeCorrelations::Averages
ModeCorrelations::average_v(structure::NoTime, const LazyDensityOperator& matrix) const
{
  using quantumdata::partialTrace;

  auto averages(initializedAverages());

  {
    auto fun=[&](const auto& psiS){return averagedMode_.average(0,psiS);};

    const Averages
      a0{partialTrace<V0>(matrix,fun)},
      a1{partialTrace<V1>(matrix,fun)};

    copy(a1,copy(a0,averages.begin()));
  }

  for (int n=0; n<int(matrix.getDimension(0)); n++) for (int m=1; m<int(matrix.getDimension(1)); m++) {
      if(n<int(matrix.getDimension(0))-1) {
        dcomp temp=sqrt(m*(n+1))*matrix(n,m)(n+1,m-1);
        averages(14)+=real(temp);
        averages(15)+=imag(temp);
      }
      if(n>0) {
        dcomp temp=sqrt(m*n)*matrix(n,m)(n-1,m-1);
        averages(16)+=real(temp);
        averages(17)+=imag(temp);
      }
    }

  return averages;

}

void
ModeCorrelations::process_v(Averages& averages) const
{
  {
    Averages ranged(averages(blitz::Range(0, 6)));
    averagedMode_.process(ranged);
  }
  {
    Averages ranged(averages(blitz::Range(7,13)));
    averagedMode_.process(ranged);
  }

  double
    xAvr=sqrt(2)*averages( 2),
    yAvr=sqrt(2)*averages( 3),
    qAvr=sqrt(2)*averages( 9),
    pAvr=sqrt(2)*averages(10);

  double
    xq= averages(14)+averages(16)-xAvr*qAvr,
    xp= averages(15)+averages(17)-xAvr*pAvr,
    yq=-averages(15)+averages(17)-yAvr*qAvr,
    yp= averages(14)-averages(16)-yAvr*pAvr;

  averages(14)=xq;
  averages(15)=xp;
  averages(16)=yq;
  averages(17)=yp;

}


/* Copyright (c) 2021, the adamantine authors.
 *
 * This file is subject to the Modified BSD License and may not be distributed
 * without copyright and license information. Please refer to the file LICENSE
 * for the text and further information on this license.
 */

#include <DataAssimilator.hh>

#include <deal.II/dofs/dof_tools.h>
#include <deal.II/lac/linear_operator_tools.h>

#include <execution>

namespace adamantine
{

DataAssimilator::DataAssimilator(boost::property_tree::ptree const &database)
{
  // Set the solver parameters from the input database
  // PropertyTreeInput data_assimilation.solver.max_number_of_temp_vectors
  if (boost::optional<unsigned int> max_num_temp_vectors =
          database.get_optional<unsigned int>(
              "solver.max_number_of_temp_vectors"))
    _additional_data.max_n_tmp_vectors = *max_num_temp_vectors;

  // PropertyTreeInput data_assimilation.solver.max_iterations
  if (boost::optional<unsigned int> max_iterations =
          database.get_optional<unsigned int>("solver.max_iterations"))
    _solver_control.set_max_steps(*max_iterations);

  // PropertyTreeInput data_assimilation.solver.convergence_tolerance
  if (boost::optional<double> tolerance =
          database.get_optional<double>("solver.convergence_tolerance"))
    _solver_control.set_tolerance(*tolerance);
}

void DataAssimilator::update_ensemble(
    std::vector<dealii::LA::distributed::Vector<double>> &sim_data,
    std::vector<double> const &expt_data, dealii::SparseMatrix<double> &R)
{
  // Set some constants
  _num_ensemble_members = sim_data.size();
  if (sim_data.size() > 0)
  {
    _sim_size = sim_data[0].size();
  }
  else
  {
    _sim_size = 0;
  }
  _expt_size = expt_data.size();

  // Get the perturbed innovation, ( y+u - Hx )
  std::vector<dealii::Vector<double>> perturbed_innovation(
      _num_ensemble_members);
  for (unsigned int member = 0; member < _num_ensemble_members; ++member)
  {
    perturbed_innovation[member].reinit(_expt_size);
    fill_noise_vector(perturbed_innovation[member], R);
    dealii::Vector<double> temporary = calc_Hx(sim_data[member]);
    for (unsigned int i = 0; i < _expt_size; ++i)
    {
      perturbed_innovation[member][i] += expt_data[i] - temporary[i];
    }
  }

  // Apply the Kalman filter to the perturbed innovation, K ( y+u - Hx )
  std::vector<dealii::LA::distributed::Vector<double>> forecast_shift =
      apply_kalman_gain(sim_data, R, perturbed_innovation);

  // Update the ensemble, x = x + K ( y+u - Hx )
  for (unsigned int member = 0; member < _num_ensemble_members; ++member)
  {
    sim_data[member] += forecast_shift[member];
  }
}

std::vector<dealii::LA::distributed::Vector<double>>
DataAssimilator::apply_kalman_gain(
    std::vector<dealii::LA::distributed::Vector<double>> &vec_ensemble,
    dealii::SparseMatrix<double> &R,
    std::vector<dealii::Vector<double>> &perturbed_innovation)
{
  /*
   * Currently this function uses GMRES to apply the inverse of HPH^T+R in the
   * Kalman gain calculation for each ensemble member individually. Depending on
   * the size of the datasets, the number of ensembles, and other factors doing
   * a direct solve of (HPH^T+R)^-1 once and then applying to the perturbed
   * innovation from each ensemble member might be more efficient.
   */

  dealii::SparsityPattern pattern_H(_expt_size, _sim_size, _expt_size);
  dealii::SparseMatrix<double> H = calc_H(pattern_H);

  dealii::FullMatrix<double> P = calc_sample_covariance_dense(vec_ensemble);

  const auto op_H = dealii::linear_operator(H);
  const auto op_P = dealii::linear_operator(P);
  const auto op_R = dealii::linear_operator(R);

  const auto op_HPH_plus_R =
      op_H * op_P * dealii::transpose_operator(op_H) + op_R;

  std::vector<dealii::LA::distributed::Vector<double>> output(
      _num_ensemble_members,
      dealii::LA::distributed::Vector<double>(_sim_size));

  // Create non-member versions of these for use in the lambda function
  auto solver_control = _solver_control;
  auto additional_data = _additional_data;
  auto sim_size = _sim_size;

  // Apply the Kalman gain to the perturbed innovation for the ensemble members
  // in parallel
  std::transform(
      std::execution::par, perturbed_innovation.begin(),
      perturbed_innovation.end(), output.begin(),
      [&](dealii::Vector<double> entry) {
        dealii::SolverGMRES<dealii::Vector<double>> HPH_plus_R_inv_solver(
            solver_control, additional_data);

        auto op_HPH_plus_R_inv =
            dealii::inverse_operator(op_HPH_plus_R, HPH_plus_R_inv_solver);

        const auto op_K =
            op_P * dealii::transpose_operator(op_H) * op_HPH_plus_R_inv;

        // Apply the Kalman gain to each innovation vector
        dealii::Vector<double> temporary = op_K * entry;

        // Copy into a distributed vector, this is the only place where the
        // mismatch matters, using dealii::Vector for the experimental data
        // and dealii::LA::distributed::Vector for the simulation data.
        dealii::LA::distributed::Vector<double> output_member(sim_size);
        for (unsigned int i = 0; i < sim_size; ++i)
        {
          output_member(i) = temporary(i);
        }

        return output_member;
      });

  return output;
}

dealii::SparseMatrix<double>
DataAssimilator::calc_H(dealii::SparsityPattern &pattern) const
{
  int num_expt_dof_map_entries = _expt_to_dof_mapping.first.size();

  for (auto i = 0; i < num_expt_dof_map_entries; ++i)
  {
    auto sim_index = _expt_to_dof_mapping.second[i];
    auto expt_index = _expt_to_dof_mapping.first[i];
    pattern.add(expt_index, sim_index);
  }

  pattern.compress();

  dealii::SparseMatrix<double> H(pattern);

  for (auto i = 0; i < num_expt_dof_map_entries; ++i)
  {
    auto sim_index = _expt_to_dof_mapping.second[i];
    auto expt_index = _expt_to_dof_mapping.first[i];
    H.add(expt_index, sim_index, 1.0);
  }

  return H;
}

template <int dim>
void DataAssimilator::update_dof_mapping(
    dealii::DoFHandler<dim> const &dof_handler,
    std::pair<std::vector<int>, std::vector<int>> const &indices_and_offsets)
{
  std::map<dealii::types::global_dof_index, dealii::Point<dim>> indices_points;
  dealii::DoFTools::map_dofs_to_support_points(
      dealii::StaticMappingQ1<dim>::mapping, dof_handler, indices_points);
  // Change the format to the one used by ArborX
  std::vector<dealii::types::global_dof_index> dof_indices(
      indices_points.size());
  unsigned int pos = 0;
  for (auto map_it = indices_points.begin(); map_it != indices_points.end();
       ++map_it, ++pos)
  {
    dof_indices[pos] = map_it->first;
  }

  _expt_to_dof_mapping.first.resize(indices_and_offsets.first.size());
  _expt_to_dof_mapping.second.resize(indices_and_offsets.first.size());

  for (unsigned int i = 0; i < _expt_size; ++i)
  {
    for (int j = indices_and_offsets.second[i];
         j < indices_and_offsets.second[i + 1]; ++j)
    {
      _expt_to_dof_mapping.first[j] = i;
      _expt_to_dof_mapping.second[j] =
          dof_indices[indices_and_offsets.first[j]];
    }
  }
}

dealii::Vector<double> DataAssimilator::calc_Hx(
    const dealii::LA::distributed::Vector<double> &sim_ensemble_member) const
{
  int num_expt_dof_map_entries = _expt_to_dof_mapping.first.size();

  dealii::Vector<double> out_vec(_expt_size);

  // Loop through the observation map to get the observation indices
  for (auto i = 0; i < num_expt_dof_map_entries; ++i)
  {
    auto sim_index = _expt_to_dof_mapping.second[i];
    auto expt_index = _expt_to_dof_mapping.first[i];
    out_vec(expt_index) = sim_ensemble_member(sim_index);
  }

  return out_vec;
}

void DataAssimilator::fill_noise_vector(dealii::Vector<double> &vec,
                                        dealii::SparseMatrix<double> &R)
{
  auto vector_size = vec.size();

  // Do Cholesky decomposition
  dealii::FullMatrix<double> L(vector_size);
  dealii::FullMatrix<double> R_full(vector_size);
  R_full.copy_from(R);
  L.cholesky(R_full);

  // Get a vector of normally distributed values
  dealii::Vector<double> uncorrelated_noise_vector(vector_size);

  for (unsigned int i = 0; i < vector_size; ++i)
  {
    uncorrelated_noise_vector(i) = _normal_dist_generator(_prng);
  }

  L.vmult(vec, uncorrelated_noise_vector);
}

template <typename VectorType>
dealii::FullMatrix<double> DataAssimilator::calc_sample_covariance_dense(
    std::vector<VectorType> vec_ensemble) const
{
  unsigned int num_ensemble_members = vec_ensemble.size();
  unsigned int vec_size = 0;
  if (vec_ensemble.size() > 0)
  {
    vec_size = vec_ensemble[0].size();
  }

  // Calculate the mean
  dealii::Vector<double> mean(vec_size);
  for (unsigned int i = 0; i < vec_size; ++i)
  {
    double sum = 0.0;
    for (unsigned int sample = 0; sample < num_ensemble_members; ++sample)
    {
      sum += vec_ensemble[sample][i];
    }
    mean[i] = sum / num_ensemble_members;
  }

  // Calculate the anomaly
  dealii::FullMatrix<double> anomaly(vec_size, num_ensemble_members);
  for (unsigned int member = 0; member < num_ensemble_members; ++member)
  {
    for (unsigned int i = 0; i < vec_size; ++i)
    {
      anomaly(i, member) = (vec_ensemble[member][i] - mean[i]) /
                           std::sqrt(num_ensemble_members - 1.0);
    }
  }

  dealii::FullMatrix<double> cov(vec_size);
  anomaly.mTmult(cov, anomaly);

  return cov;
}

// Explicit instantiation
template void DataAssimilator::update_dof_mapping<2>(
    dealii::DoFHandler<2> const &dof_handler,
    std::pair<std::vector<int>, std::vector<int>> const &indices_and_offsets);
template void DataAssimilator::update_dof_mapping<3>(
    dealii::DoFHandler<3> const &dof_handler,
    std::pair<std::vector<int>, std::vector<int>> const &indices_and_offsets);
template dealii::FullMatrix<double>
DataAssimilator::calc_sample_covariance_dense<dealii::Vector<double>>(
    std::vector<dealii::Vector<double>> vec_ensemble) const;
template dealii::FullMatrix<double>
DataAssimilator::calc_sample_covariance_dense<
    dealii::LA::distributed::Vector<double>>(
    std::vector<dealii::LA::distributed::Vector<double>> vec_ensemble) const;

} // namespace adamantine