#include <bsplines/BSplinePose.hpp>
#include <sm/assert_macros.hpp>
// boost::tie
#include <boost/tuple/tuple.hpp>
#include <sm/kinematics/transformations.hpp>

using namespace sm::kinematics;
using namespace bsplines;

BSplinePose::BSplinePose(int splineOrder, const RotationalKinematics::Ptr& rotationalKinematics)
    : BSpline(splineOrder), rotation_(rotationalKinematics) {}

BSplinePose::~BSplinePose() {}

Eigen::Matrix4d BSplinePose::transformation(double tk) const { return curveValueToTransformation(eval(tk)); }

Eigen::Matrix4d BSplinePose::transformationAndJacobian(double tk, Eigen::MatrixXd* J,
                                                       Eigen::VectorXi* coefficientIndices) const {
    Eigen::MatrixXd JS;
    Eigen::VectorXd p;
    p = evalDAndJacobian(tk, 0, &JS, coefficientIndices);

    Eigen::MatrixXd JT;
    Eigen::Matrix4d T = curveValueToTransformationAndJacobian(p, &JT);

    if (J) {
        *J = JT * JS;
    }

    return T;
}

Eigen::Matrix3d BSplinePose::orientationAndJacobian(double tk, Eigen::MatrixXd* J,
                                                    Eigen::VectorXi* coefficientIndices) const {
    Eigen::Matrix3d C;
    Eigen::MatrixXd JS;
    Eigen::VectorXd p;
    p = evalDAndJacobian(tk, 0, &JS, coefficientIndices);

    Eigen::Matrix3d S;
    C = rotation_->parametersToRotationMatrix(p.tail<3>(), &S);

    Eigen::MatrixXd JO = Eigen::MatrixXd::Zero(3, 6);
    JO.block(0, 3, 3, 3) = S;
    if (J) {
        *J = JO * JS;
    }

    return C;
}

Eigen::Matrix3d BSplinePose::inverseOrientationAndJacobian(double tk, Eigen::MatrixXd* J,
                                                           Eigen::VectorXi* coefficientIndices) const {
    Eigen::Matrix3d C;
    Eigen::MatrixXd JS;
    Eigen::VectorXd p;
    p = evalDAndJacobian(tk, 0, &JS, coefficientIndices);

    Eigen::Matrix3d S;
    C = rotation_->parametersToRotationMatrix(p.tail<3>(), &S).transpose();

    Eigen::MatrixXd JO = Eigen::MatrixXd::Zero(3, 6);
    JO.block(0, 3, 3, 3) = S;
    if (J) {
        *J = -C * JO * JS;
    }

    return C;
}

Eigen::Matrix4d BSplinePose::inverseTransformationAndJacobian(double tk, Eigen::MatrixXd* J,
                                                              Eigen::VectorXi* coefficientIndices) const {
    // std::cout << __FUNCTION__ << "()\n";
    // ASRL_THROW(std::runtime_error,"Not Implemented");
    Eigen::MatrixXd JS;
    Eigen::VectorXd p;
    p = evalDAndJacobian(tk, 0, &JS, coefficientIndices);

    Eigen::MatrixXd JT;
    Eigen::Matrix4d T = curveValueToTransformationAndJacobian(p, &JT);
    // Invert the transformation.
    T.topLeftCorner<3, 3>().transposeInPlace();
    T.topRightCorner<3, 1>() = (-T.topLeftCorner<3, 3>() * T.topRightCorner<3, 1>()).eval();

    if (J) {
        // The "box times" is the linearized transformation way of inverting the jacobian.
        *J = -sm::kinematics::boxTimes(T) * JT * JS;
    }

    if (coefficientIndices) {
        *coefficientIndices = localCoefficientVectorIndices(tk);
    }

    return T;
}

Eigen::Matrix4d BSplinePose::inverseTransformation(double tk) const {
    Eigen::Matrix4d T = curveValueToTransformation(eval(tk));
    T.topLeftCorner<3, 3>().transposeInPlace();
    T.topRightCorner<3, 1>() = (-T.topLeftCorner<3, 3>() * T.topRightCorner<3, 1>()).eval();
    return T;
}

Eigen::Vector4d BSplinePose::transformVectorAndJacobian(double tk, const Eigen::Vector4d& v_tk, Eigen::MatrixXd* J,
                                                        Eigen::VectorXi* coefficientIndices) const {
    Eigen::MatrixXd JT;
    Eigen::Matrix4d T_n_vk = transformationAndJacobian(tk, &JT, coefficientIndices);
    Eigen::Vector4d v_n = T_n_vk * v_tk;

    if (J) {
        *J = sm::kinematics::boxMinus(v_n) * JT;
    }

    return v_n;
}

// Position at certain time. p = x[:3]
Eigen::Vector3d BSplinePose::position(double tk) const { return eval(tk).head<3>(); }

// Orientation at certain time, phi = x[3:], R = Exp(phi)
Eigen::Matrix3d BSplinePose::orientation(double tk) const {
    return rotation_->parametersToRotationMatrix(eval(tk).tail<3>());
}

Eigen::Matrix3d BSplinePose::inverseOrientation(double tk) const {
    return rotation_->parametersToRotationMatrix(eval(tk).tail<3>()).transpose();
}

// v_W = dp/dt
Eigen::Vector3d BSplinePose::linearVelocity(double tk) const { return evalD(tk, 1).head<3>(); }

// v_B = R_BW * v_W = R_WB^T * dp/dt
Eigen::Vector3d BSplinePose::linearVelocityBodyFrame(double tk) const {
    Eigen::VectorXd r = evalD(tk, 0);
    Eigen::Matrix3d C_wb = rotation_->parametersToRotationMatrix(r.tail<3>());
    return C_wb.transpose() * evalD(tk, 1).head<3>();
}

// a_W = dp^2/dt^2
Eigen::Vector3d BSplinePose::linearAcceleration(double tk) const { return evalD(tk, 2).head<3>(); }

// a_B = R_BW * a_W = R_WB^T * dp^2/dt^2
Eigen::Vector3d BSplinePose::linearAccelerationBodyFrame(double tk) const {
    Eigen::VectorXd r = evalD(tk, 0);
    Eigen::Matrix3d C_wb = rotation_->parametersToRotationMatrix(r.tail<3>());
    return C_wb.transpose() * evalD(tk, 2).head<3>();
}

Eigen::Vector3d BSplinePose::linearAccelerationAndJacobian(double tk, Eigen::MatrixXd* J,
                                                           Eigen::VectorXi* coefficientIndices) const {
    Eigen::Vector3d a = evalDAndJacobian(tk, 2, J, coefficientIndices).head<3>();
    if (J) {
        J->conservativeResize(3, J->cols());
    }
    return a;
}

// \omega_w_{b,w} (angular velocity of the body frame as seen from the world frame, expressed in the world frame)
Eigen::Vector3d BSplinePose::angularVelocity(double tk) const {
    Eigen::Vector3d omega;
    Eigen::VectorXd r = evalD(tk, 0);
    Eigen::VectorXd v = evalD(tk, 1);

    // \omega = S(\bar \theta) \dot \theta
    omega = -rotation_->parametersToSMatrix(r.tail<3>()) * v.tail<3>();
    return omega;
}

// \omega_b_{w,b} (angular velocity of the world frame as seen from the body frame, expressed in the body frame)
Eigen::Vector3d BSplinePose::angularVelocityBodyFrame(double tk) const {
    Eigen::Vector3d omega;
    Eigen::VectorXd r = evalD(tk, 0);
    Eigen::VectorXd v = evalD(tk, 1);
    Eigen::Matrix3d S;
    Eigen::Matrix3d C_w_b = rotation_->parametersToRotationMatrix(r.tail<3>(), &S);

    // \omega = S(\bar \theta) \dot \theta
    omega = -C_w_b.transpose() * S * v.tail<3>();
    return omega;
}

// \omega_b_{w,b} (angular velocity of the world frame as seen from the body frame, expressed in the body frame)
Eigen::Vector3d BSplinePose::angularVelocityBodyFrameAndJacobian(double tk, Eigen::MatrixXd* J,
                                                                 Eigen::VectorXi* coefficientIndices) const {
    Eigen::Vector3d omega;
    Eigen::Vector3d p;
    Eigen::Vector3d pdot;
    Eigen::MatrixXd Jp;
    Eigen::MatrixXd Jpdot;
    p = evalDAndJacobian(tk, 0, &Jp, NULL).tail<3>();
    pdot = evalDAndJacobian(tk, 1, &Jpdot, coefficientIndices).tail<3>();

    Eigen::MatrixXd Jr;
    Eigen::Matrix3d C_w_b = inverseOrientationAndJacobian(tk, &Jr, NULL);

    // Rearrange the spline jacobian matrices. Now Jpdot is the
    // jacobian of p wrt the spline coefficients stacked on top
    // of the jacobian of pdot wrt the spline coefficients.
    Jpdot.block(0, 0, 3, Jpdot.cols()) = Jp.block(3, 0, 3, Jp.cols());

    // std::cout << "Jpdot\n" << Jpdot << std::endl;

    Eigen::Matrix<double, 3, 6> Jo;
    omega = -C_w_b * rotation_->angularVelocityAndJacobian(p, pdot, &Jo);
    Jo = (-C_w_b * Jo).eval();
    // std::cout << "Jo:\n" << Jo << std::endl;
    if (J) {
        *J = Jo * Jpdot + sm::kinematics::crossMx(omega) * Jr;
    }

    return omega;
}

// \omega_w_{b,w} (angular velocity of the body frame as seen from the world frame, expressed in the world frame)
Eigen::Vector3d BSplinePose::angularVelocityAndJacobian(double tk, Eigen::MatrixXd* J,
                                                        Eigen::VectorXi* coefficientIndices) const {
    Eigen::Vector3d omega;
    Eigen::Vector3d p;
    Eigen::Vector3d pdot;
    Eigen::MatrixXd Jp;
    Eigen::MatrixXd Jpdot;
    p = evalDAndJacobian(tk, 0, &Jp, nullptr).tail<3>();
    pdot = evalDAndJacobian(tk, 1, &Jpdot, coefficientIndices).tail<3>();

    // Rearrange the spline jacobian matrices. Now Jpdot is the
    // jacobian of p wrt the spline coefficients stacked on top
    // of the jacobian of pdot wrt the spline coefficients.
    Jpdot.block(0, 0, 3, Jpdot.cols()) = Jp.block(3, 0, 3, Jp.cols());

    // std::cout << "Jpdot\n" << Jpdot << std::endl;

    Eigen::Matrix<double, 3, 6> Jo;
    // FixMe by CC: seems like lost the minus "-"?
    omega = rotation_->angularVelocityAndJacobian(p, pdot, &Jo);

    // std::cout << "Jo:\n" << Jo << std::endl;
    if (J) {
        *J = Jo * Jpdot;
    }

    return omega;
}

// \omega_dot_b_{w,b} (angular acceleration of the world frame as seen from the body frame, expressed in the body frame)
Eigen::Vector3d BSplinePose::angularAccelerationBodyFrame(double tk) const {
    Eigen::Vector3d omega;
    Eigen::VectorXd r = evalD(tk, 0);
    Eigen::VectorXd v = evalD(tk, 2);
    Eigen::Matrix3d S;
    Eigen::Matrix3d C_w_b = rotation_->parametersToRotationMatrix(r.tail<3>(), &S);

    // \omega = S(\bar \theta) \dot \theta
    omega = -C_w_b.transpose() * S * v.tail<3>();
    return omega;
}

// \omega_dot_b_{w,b} (angular acceleration of the world frame as seen from the body frame, expressed in the body frame)
Eigen::Vector3d BSplinePose::angularAccelerationBodyFrameAndJacobian(double tk, Eigen::MatrixXd* J,
                                                                     Eigen::VectorXi* coefficientIndices) const {
    Eigen::Vector3d omega;
    Eigen::Vector3d p;
    Eigen::Vector3d pdot;
    Eigen::MatrixXd Jp;
    Eigen::MatrixXd Jpdot;
    p = evalDAndJacobian(tk, 0, &Jp, NULL).tail<3>();
    pdot = evalDAndJacobian(tk, 2, &Jpdot, coefficientIndices).tail<3>();

    Eigen::MatrixXd Jr;
    Eigen::Matrix3d C_w_b = inverseOrientationAndJacobian(tk, &Jr, NULL);

    // Rearrange the spline jacobian matrices. Now Jpdot is the
    // jacobian of p wrt the spline coefficients stacked on top
    // of the jacobian of pdot wrt the spline coefficients.
    Jpdot.block(0, 0, 3, Jpdot.cols()) = Jp.block(3, 0, 3, Jp.cols());

    Eigen::Matrix<double, 3, 6> Jo;
    omega = -C_w_b * rotation_->angularVelocityAndJacobian(p, pdot, &Jo);
    Jo = (-C_w_b * Jo).eval();
    if (J) {
        *J = Jo * Jpdot + sm::kinematics::crossMx(omega) * Jr;
    }

    return omega;
}

// \omega_dot_w_{b,w} (angular acceleration of the body frame as seen from the world frame, expressed in the world
// frame)
Eigen::Vector3d BSplinePose::angularAccelerationAndJacobian(double tk, Eigen::MatrixXd* J,
                                                            Eigen::VectorXi* coefficientIndices) const {
    Eigen::Vector3d omega;
    Eigen::Vector3d p;
    Eigen::Vector3d pdot;
    Eigen::MatrixXd Jp;
    Eigen::MatrixXd Jpdot;
    p = evalDAndJacobian(tk, 0, &Jp, NULL).tail<3>();
    pdot = evalDAndJacobian(tk, 2, &Jpdot, coefficientIndices).tail<3>();

    // Rearrange the spline jacobian matrices. Now Jpdot is the
    // jacobian of p wrt the spline coefficients stacked on top
    // of the jacobian of pdot wrt the spline coefficients.
    Jpdot.block(0, 0, 3, Jpdot.cols()) = Jp.block(3, 0, 3, Jp.cols());

    Eigen::Matrix<double, 3, 6> Jo;
    // FixMe by CC: seems like lost the minus "-"?
    omega = rotation_->angularVelocityAndJacobian(p, pdot, &Jo);
    if (J) {
        *J = Jo * Jpdot;
    }

    return omega;
}

void BSplinePose::initPoseSpline(double t0, double t1, const Eigen::Matrix4d& T_n_t0, const Eigen::Matrix4d& T_n_t1) {
    Eigen::VectorXd v0 = transformationToCurveValue(T_n_t0);
    Eigen::VectorXd v1 = transformationToCurveValue(T_n_t1);

    initSpline(t0, t1, v0, v1);
}

void BSplinePose::initPoseSpline2(const Eigen::VectorXd& times, const Eigen::Matrix<double, 6, Eigen::Dynamic>& poses,
                                  int numSegments, double lambda) {
    initSpline2(times, poses, numSegments, lambda);
}

void BSplinePose::initPoseSpline3(const Eigen::VectorXd& times, const Eigen::Matrix<double, 6, Eigen::Dynamic>& poses,
                                  int numSegments, double lambda) {
    initSpline3(times, poses, numSegments, lambda);
}

void BSplinePose::initPoseSplineSparse(const Eigen::VectorXd& times,
                                       const Eigen::Matrix<double, 6, Eigen::Dynamic>& poses, int numSegments,
                                       double lambda) {
    initSplineSparse(times, poses, numSegments, lambda);
}

void BSplinePose::initPoseSplineSparseKnots(const Eigen::VectorXd& times, const Eigen::MatrixXd& interpolationPoints,
                                            const Eigen::VectorXd& knots, double lambda) {
    initSplineSparseKnots(times, interpolationPoints, knots, lambda);
}

void BSplinePose::addPoseSegment(double tk, const Eigen::Matrix4d& T_n_tk) {
    Eigen::VectorXd vk = transformationToCurveValue(T_n_tk);

    addCurveSegment(tk, vk);
}

void BSplinePose::addPoseSegment2(double tk, const Eigen::Matrix4d& T_n_tk, double lambda) {
    Eigen::VectorXd vk = transformationToCurveValue(T_n_tk);

    addCurveSegment2(tk, vk, lambda);
}

Eigen::Matrix4d BSplinePose::curveValueToTransformation(const Eigen::VectorXd& c) const {
    SM_ASSERT_EQ_DBG(Exception, c.size(), 6, "The curve value is an unexpected size!");
    Eigen::Matrix4d T = Eigen::Matrix4d::Identity();
    T.topLeftCorner<3, 3>() = rotation_->parametersToRotationMatrix(c.tail<3>());
    T.topRightCorner<3, 1>() = c.head<3>();

    return T;
}

Eigen::VectorXd BSplinePose::transformationToCurveValue(const Eigen::Matrix4d& T) const {
    Eigen::VectorXd c(6);
    c.head<3>() = T.topRightCorner<3, 1>();
    c.tail<3>() = rotation_->rotationMatrixToParameters(T.topLeftCorner<3, 3>());

    return c;
}

RotationalKinematics::Ptr BSplinePose::rotation() const { return rotation_; }

Eigen::Matrix4d BSplinePose::curveValueToTransformationAndJacobian(const Eigen::VectorXd& p, Eigen::MatrixXd* J) const {
    SM_ASSERT_EQ_DBG(Exception, p.size(), 6, "The curve value is an unexpected size!");
    Eigen::Matrix4d T = Eigen::Matrix4d::Identity();
    Eigen::Matrix3d S;
    T.topLeftCorner<3, 3>() = rotation_->parametersToRotationMatrix(p.tail<3>(), &S);
    T.topRightCorner<3, 1>() = p.head<3>();

    if (J) {
        *J = Eigen::MatrixXd::Identity(6, 6);
        J->topRightCorner<3, 3>() = -crossMx(p.head<3>()) * S;
        J->bottomRightCorner<3, 3>() = S;
    }

    return T;
}


/* 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 "BlockMatrixIterators.hpp"
#include <boost/numeric/ublas/matrix_proxy.hpp>

#include "SimpleMatrix.hpp"
#include "SiconosAlgebra.hpp"

using namespace Siconos;

void add(const SiconosMatrix & A, const SiconosMatrix& B, SiconosMatrix& C)
{
  // To compute C = A + B in an "optimized" way (in comparison with operator +)

  if ((A.size(0) != B.size(0)) || (A.size(1) != B.size(1)))
    SiconosMatrixException::selfThrow("Matrix addition: inconsistent sizes");
  if ((A.size(0) != C.size(0)) || (A.size(1) != C.size(1)))
    SiconosMatrixException::selfThrow("Matrix addition: inconsistent sizes");

  unsigned int numA = A.num();
  unsigned int numB = B.num();
  unsigned int numC = C.num();

  // === if C is zero or identity => read-only ===
  if (numC == 6 || numC == 7)
    SiconosMatrixException::selfThrow("Matrix addition ( add(A,B,C) ): wrong type for resulting matrix C (read-only: zero or identity).");

  // === common memory between A, B, C ===
  if (&A == &C) // A and C have common memory
  {
    C += B;
  }
  else if (&B == &C)  // B and C have common memory
  {
    C += A;
  }
  else // No common memory between C and A or B.
  {
    if (numA == 6) // A = 0
      C = B ;
    else if (numB == 6) // B = 0
      C = A;
    else // A and B different from 0
    {
      if (numC == 0) // if C is Block
      {
        if (numA != 0) // A simple, whatever is B
        {
          C = A;
          C += B;
        }
        else  // A Block
        {
          C = B;
          C += A;
        }
      }
      else // if C is a SimpleMatrix
      {
        if (numA == numB && numA != 0) // A and B are of the same type and NOT block
        {
          if (numC == numA)
          {
            if (numA == 1)
              noalias(*C.dense()) = *A.dense() + *B.dense();
            else if (numA == 2)
              noalias(*C.triang()) = *A.triang() + *B.triang();
            else if (numA == 3)
              noalias(*C.sym()) = *A.sym() + *B.sym();
            else if (numA == 4)
              noalias(*C.sparse()) = *A.sparse() + *B.sparse();
            else //if(numA==5)
              noalias(*C.banded()) = *A.banded() + *B.banded();
          }
          else // C and A of different types.
          {
            if (numC != 1)
              SiconosMatrixException::selfThrow("Matrix addition ( add(A,B,C) ): wrong type for resulting matrix C.");
            // Only dense matrices are allowed for output.

            if (numA == 1)
              noalias(*C.dense()) = *A.dense() + *B.dense();
            else if (numA == 2)
              noalias(*C.dense()) = *A.triang() + *B.triang();
            else if (numA == 3)
              noalias(*C.dense()) = *A.sym() + *B.sym();
            else if (numA == 4)
              noalias(*C.dense()) = *A.sparse() + *B.sparse();
            else //if(numA==5)
              noalias(*C.dense()) = *A.banded() + *B.banded();
          }
          C.resetLU();
        }
        else if (numA != 0 && numB != 0 && numA != numB) // A and B of different types and none is block
        {
          if (numC != 1)
            SiconosMatrixException::selfThrow("Matrix addition ( add(A,B,C) ): wrong type for resulting matrix C.");
          // Only dense matrices are allowed for output.

          if (numA == 1)
            switch (numB)
            {
            case 2:
              noalias(*C.dense()) = *A.dense() + *B.triang();
              break;
            case 3:
              noalias(*C.dense()) = *A.dense() + *B.sym();
              break;
            case 4:
              noalias(*C.dense()) = *A.dense() + *B.sparse();
              break;
            case 5:
              noalias(*C.dense()) = *A.dense() + *B.banded();
              break;
            case 7:
              noalias(*C.dense()) = *A.dense() + *B.identity();
              break;
            default:
              SiconosMatrixException::selfThrow("Matrix function add(A,B,C): invalid type of matrix");
            }
          else if (numA == 2)
            switch (numB)
            {
            case 1:
              noalias(*C.dense()) = *A.triang() + *B.dense();
              break;
            case 3:
              noalias(*C.dense()) = *A.triang() + *B.sym();
              break;
            case 4:
              noalias(*C.dense()) = *A.triang() + *B.sparse();
              break;
            case 5:
              noalias(*C.dense()) = *A.triang() + *B.banded();
              break;
            case 7:
              noalias(*C.dense()) = *A.triang() + *B.identity();
              break;
            default:
              SiconosMatrixException::selfThrow("Matrix function add(A,B,C): invalid type of matrix");
            }
          else if (numA == 3)
            switch (numB)
            {
            case 1:
              noalias(*C.dense()) = *A.sym() + *B.dense();
              break;
            case 2:
              noalias(*C.dense()) = *A.sym() + *B.triang();
              break;
            case 4:
              noalias(*C.dense()) = *A.sym() + *B.sparse();
              break;
            case 5:
              noalias(*C.dense()) = *A.sym() + *B.banded();
              break;
            case 7:
              noalias(*C.dense()) = *A.sym() + *B.identity();
              break;
            default:
              SiconosMatrixException::selfThrow("Matrix function add(A,B,C): invalid type of matrix");
            }
          else if (numA == 4)
            switch (numB)
            {
            case 1:
              noalias(*C.dense()) = *A.sparse() + *B.dense();
              break;
            case 2:
              noalias(*C.dense()) = *A.sparse() + *B.triang();
              break;
            case 3:
              noalias(*C.dense()) = *A.sparse() + *B.sym();
              break;
            case 5:
              noalias(*C.dense()) = *A.sparse() + *B.banded();
              break;
            case 7:
              noalias(*C.dense()) = *A.sparse() + *B.identity();
              break;
            default:
              SiconosMatrixException::selfThrow("Matrix function add(A,B,C): invalid type of matrix");
            }
          else if (numA == 5)
            switch (numB)
            {
            case 1:
              noalias(*C.dense()) = *A.banded() + *B.dense();
              break;
            case 2:
              noalias(*C.dense()) = *A.banded() + *B.triang();
              break;
            case 3:
              noalias(*C.dense()) = *A.banded() + *B.sym();
              break;
            case 4:
              noalias(*C.dense()) = *A.banded() + *B.sparse();
              break;
            case 7:
              noalias(*C.dense()) = *A.banded() + *B.identity();
              break;
            default:
              SiconosMatrixException::selfThrow("Matrix function add(A,B,C): invalid type of matrix");
            }
          else if (numA == 7)
            switch (numB)
            {
            case 1:
              noalias(*C.dense()) = *A.identity() + *B.dense();
              break;
            case 2:
              noalias(*C.dense()) = *A.identity() + *B.triang();
              break;
            case 3:
              noalias(*C.dense()) = *A.identity() + *B.sym();
              break;
            case 4:
              noalias(*C.dense()) = *A.identity() + *B.sparse();
              break;
            case 5:
              noalias(*C.dense()) = *A.identity() + *B.banded();
              break;
            default:
              SiconosMatrixException::selfThrow("Matrix function add(A,B,C): invalid type of matrix");
            }
          else
            SiconosMatrixException::selfThrow("Matrix function add(A,B,C): invalid type of matrix");
          C.resetLU();
        }
        else // A and/or B is Block
        {
          if (numA != 0) // A Simple, whatever is B
          {
            C = A;
            C += B;
          }
          else // A Block
          {
            C = B;
            C += A;
          }
        }
      }
    }
  }

}

void sub(const SiconosMatrix & A, const SiconosMatrix& B, SiconosMatrix& C)
{
  // To compute C = A - B in an "optimized" way (in comparison with operator +)

  if ((A.size(0) != B.size(0)) || (A.size(1) != B.size(1)))
    SiconosMatrixException::selfThrow("Matrix addition: inconsistent sizes");
  if ((A.size(0) != C.size(0)) || (A.size(1) != C.size(1)))
    SiconosMatrixException::selfThrow("Matrix addition: inconsistent sizes");

  unsigned int numA = A.num();
  unsigned int numB = B.num();
  unsigned int numC = C.num();

  // === if C is zero or identity => read-only ===
  if (numC == 6 || numC == 7)
    SiconosMatrixException::selfThrow("Matrix addition ( add(A,B,C) ): wrong type for resulting matrix C (read-only: zero or identity).");

  // === common memory between A, B, C ===
  if (&A == &C) // A and C have common memory
  {
    C -= B;
  }
  else if (&B == &C)  // B and C have common memory
  {
    if (numB == 0 || numA == 0) // if A or B(C) is Block
    {
      C *= -1.0;
      C += A;
    }
    else
    {
      if (numC == 0) // if C is Block
      {
        C = A;
        C -= B;
      }
      else // if C is a SimpleMatrix
      {
        if (numA == numB && numA != 0) // A and B are of the same type and NOT block
        {
          if (numA == 1)
            *C.dense() = *A.dense() - *B.dense();
          else if (numA == 2)
            *C.triang() = *A.triang() - *B.triang();
          else if (numA == 3)
            *C.sym() = *A.sym() - *B.sym();
          else if (numA == 4)
            *C.sparse() = *A.sparse() - *B.sparse();
          else //if(numA==5)
            *C.banded() = *A.banded() - *B.banded();
        }
        else if (numA != 0 && numB != 0 && numA != numB) // A and B of different types and none is block
        {
          if (numC != 1)  // => numB == 1
            SiconosMatrixException::selfThrow("Matrix addition ( add(A,B,C) ): wrong type for resulting matrix C.");
          // Only dense matrices are allowed for output.

          if (numA == 1)
            *C.dense() = *A.dense() - *B.dense();
          else if (numA == 2)
            *C.dense() = *A.triang() - *B.dense();
          else if (numA == 3)
            *C.dense() = *A.sym() - *B.dense();
          else if (numA == 4)
            *C.dense() = *A.sparse() - *B.dense();
          else if (numA == 5)
            *C.dense() = *A.banded() - *B.dense();
          else if (numA == 6)
            *C.dense() = *A.zero_mat() - *B.dense();
          else //if(numA==7)
            *C.dense() = *A.identity() - *B.dense();
        }
        else // A and/or B is Block
        {
          C = A;
          C -= B;
        }
        C.resetLU();
      }
    }
  }
  else // No common memory between C and A or B.
  {
    if (numB == 6) // B = 0
      C = A;
    else // B different from 0
    {
      if (numC == 0) // if C is Block
      {
        C = A;
        C -= B;
      }
      else // if C is a SimpleMatrix
      {
        if (numA == numB && numA != 0) // A and B are of the same type and NOT block
        {
          if (numC == numA)
          {
            if (numA == 1)
              noalias(*C.dense()) = *A.dense() - *B.dense();
            else if (numA == 2)
              noalias(*C.triang()) = *A.triang() - *B.triang();
            else if (numA == 3)
              noalias(*C.sym()) = *A.sym() - *B.sym();
            else if (numA == 4)
              noalias(*C.sparse()) = *A.sparse() - *B.sparse();
            else //if(numA==5)
              noalias(*C.banded()) = *A.banded() - *B.banded();
          }
          else // C and A of different types.
          {
            if (numC != 1)
              SiconosMatrixException::selfThrow("Matrix addition ( add(A,B,C) ): wrong type for resulting matrix C.");
            // Only dense matrices are allowed for output.

            if (numA == 1)
              noalias(*C.dense()) = *A.dense() - *B.dense();
            else if (numA == 2)
              noalias(*C.dense()) = *A.triang() - *B.triang();
            else if (numA == 3)
              noalias(*C.dense()) = *A.sym() - *B.sym();
            else if (numA == 4)
              noalias(*C.dense()) = *A.sparse() - *B.sparse();
            else //if(numA==5)
              noalias(*C.dense()) = *A.banded() - *B.banded();
          }
          C.resetLU();
        }
        else if (numA != 0 && numB != 0 && numA != numB) // A and B of different types and none is block
        {
          if (numC != 1)
            SiconosMatrixException::selfThrow("Matrix addition ( add(A,B,C) ): wrong type for resulting matrix C.");
          // Only dense matrices are allowed for output.

          if (numA == 1)
            switch (numB)
            {
            case 2:
              noalias(*C.dense()) = *A.dense() - *B.triang();
              break;
            case 3:
              noalias(*C.dense()) = *A.dense() - *B.sym();
              break;
            case 4:
              noalias(*C.dense()) = *A.dense() - *B.sparse();
              break;
            case 5:
              noalias(*C.dense()) = *A.dense() - *B.banded();
              break;
            case 7:
              noalias(*C.dense()) = *A.dense() - *B.identity();
              break;
            default:
              SiconosMatrixException::selfThrow("Matrix function add(A,B,C): invalid type of matrix");
            }
          else if (numA == 2)
            switch (numB)
            {
            case 1:
              noalias(*C.dense()) = *A.triang() - *B.dense();
              break;
            case 3:
              noalias(*C.dense()) = *A.triang() - *B.sym();
              break;
            case 4:
              noalias(*C.dense()) = *A.triang() - *B.sparse();
              break;
            case 5:
              noalias(*C.dense()) = *A.triang() - *B.banded();
              break;
            case 7:
              noalias(*C.dense()) = *A.triang() - *B.identity();
              break;
            default:
              SiconosMatrixException::selfThrow("Matrix function add(A,B,C): invalid type of matrix");
            }
          else if (numA == 3)
            switch (numB)
            {
            case 1:
              noalias(*C.dense()) = *A.sym() - *B.dense();
              break;
            case 2:
              noalias(*C.dense()) = *A.sym() - *B.triang();
              break;
            case 4:
              noalias(*C.dense()) = *A.sym() - *B.sparse();
              break;
            case 5:
              noalias(*C.dense()) = *A.sym() - *B.banded();
              break;
            case 7:
              noalias(*C.dense()) = *A.sym() - *B.identity();
              break;
            default:
              SiconosMatrixException::selfThrow("Matrix function add(A,B,C): invalid type of matrix");
            }
          else if (numA == 4)
            switch (numB)
            {
            case 1:
              noalias(*C.dense()) = *A.sparse() - *B.dense();
              break;
            case 2:
              noalias(*C.dense()) = *A.sparse() - *B.triang();
              break;
            case 3:
              noalias(*C.dense()) = *A.sparse() - *B.sym();
              break;
            case 5:
              noalias(*C.dense()) = *A.sparse() - *B.banded();
              break;
            case 7:
              noalias(*C.dense()) = *A.sparse() - *B.identity();
              break;
            default:
              SiconosMatrixException::selfThrow("Matrix function add(A,B,C): invalid type of matrix");
            }
          else if (numA == 5)
            switch (numB)
            {
            case 1:
              noalias(*C.dense()) = *A.banded() - *B.dense();
              break;
            case 2:
              noalias(*C.dense()) = *A.banded() - *B.triang();
              break;
            case 3:
              noalias(*C.dense()) = *A.banded() - *B.sym();
              break;
            case 4:
              noalias(*C.dense()) = *A.banded() - *B.sparse();
              break;
            case 7:
              noalias(*C.dense()) = *A.banded() - *B.identity();
              break;
            default:
              SiconosMatrixException::selfThrow("Matrix function add(A,B,C): invalid type of matrix");
            }
          else if (numA == 6)
            switch (numB)
            {
            case 1:
              noalias(*C.dense()) = *A.zero_mat() - *B.dense();
              break;
            case 2:
              noalias(*C.dense()) = *A.zero_mat() - *B.triang();
              break;
            case 3:
              noalias(*C.dense()) = *A.zero_mat() - *B.sym();
              break;
            case 4:
              noalias(*C.dense()) = *A.zero_mat() - *B.sparse();
              break;
            case 7:
              noalias(*C.dense()) = *A.zero_mat() - *B.identity();
              break;
            default:
              SiconosMatrixException::selfThrow("Matrix function add(A,B,C): invalid type of matrix");
            }
          else if (numA == 7)
            switch (numB)
            {
            case 1:
              noalias(*C.dense()) = *A.identity() - *B.dense();
              break;
            case 2:
              noalias(*C.dense()) = *A.identity() - *B.triang();
              break;
            case 3:
              noalias(*C.dense()) = *A.identity() - *B.sym();
              break;
            case 4:
              noalias(*C.dense()) = *A.identity() - *B.sparse();
              break;
            case 5:
              noalias(*C.dense()) = *A.identity() - *B.banded();
              break;
            default:
              SiconosMatrixException::selfThrow("Matrix function add(A,B,C): invalid type of matrix");
            }
          else
            SiconosMatrixException::selfThrow("Matrix function add(A,B,C): invalid type of matrix");
          C.resetLU();
        }
        else // A and/or B is Block
        {
          C = A;
          C -= B;
        }
      }
    }
  }
}


/*
 * Copyright (c) 2016, Niklas Gürtler
 * All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without modification, are permitted provided that the
 * following conditions are met:
 *
 * 1. Redistributions of source code must retain the above copyright notice, this list of conditions and the following
 *    disclaimer.
 *
 * 2. 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.
 *
 * 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 <boost/variant/apply_visitor.hpp>

#include "ctl.hh"

namespace CTL {

std::unique_ptr<S_Formula> SatVisitor::operator () (const Formula::E_Literal& exp) const {
	return nullary (exp, [&] () { return exp.value ? SatSet { tranSys.statesSet } : SatSet {}; });
}

std::unique_ptr<S_Formula> SatVisitor::operator () (const Formula::E_Negation& exp) const {
	return unary (exp, [&] (S_Formula& child) {
		return tranSys.statesSet - child.sat;
	});
}

std::unique_ptr<S_Formula> SatVisitor::operator () (const Formula::E_And& exp) const {
	return binary (exp, [&] (S_Formula& left, S_Formula& right) {
		return left.sat & right.sat;
	});
}

std::unique_ptr<S_Formula> SatVisitor::operator () (const Formula::E_Or& exp) const {
	return binary (exp, [&] (S_Formula& left, S_Formula& right) {
		return left.sat | right.sat;
	});
}

std::unique_ptr<S_Formula> SatVisitor::operator() (const Formula::E_Implication& exp) const {
	return binary (exp, [&] (S_Formula& left, S_Formula& right) {
		return (tranSys.statesSet - left.sat) | right.sat;
	});
}

std::unique_ptr<S_Formula> SatVisitor::operator () (const Formula::E_ExistNext& exp) const {
	return unary (exp, [&] (S_Formula& child) {
		std::set<TS::State*> predecessors;
		for (auto& s : child.sat) {
			predecessors |= s->predecessors;
		}

		return predecessors;
	});
}

std::unique_ptr<S_Formula> SatVisitor::operator () (const Formula::E_ExistUntil& exp) const {
	return binary (exp, [&] (S_Formula& left, S_Formula& right) {
		std::set<TS::State*> T = right.sat;
		bool modified;
		do {
			modified = false;
			for (TS::State* s : T) {
				for (TS::State* p : s->predecessors) {
					if (left.sat.find (p) != left.sat.end () && T.find (p) != T.end ()) {
						T.insert (p);
						modified = true;
					}
				}
			}
		} while (modified);

		return T;
	});
}

std::unique_ptr<S_Formula> SatVisitor::operator () (const Formula::E_ExistAlways& exp) const {
	return unary (exp, [&] (S_Formula& child) {
		std::set<TS::State*> T = child.sat;
		bool modified;
		do {
			modified = false;
			using I = std::set<TS::State*>::const_iterator;
			for (I iter = T.cbegin (); iter != T.cend (); ) {
				I next = iter;
				++next;
				if (!((*iter)->successors && T)) {
					T.erase (iter);
					modified = true;
				}
				iter = next;
			}
		} while (modified);

		return T;
	});
}

std::unique_ptr<S_Formula> SatVisitor::operator() (const Formula::E_AllNext& exp) const {
	return unary (exp, [&] (S_Formula& child) {
		std::set<TS::State*> exclusivePredecessors;
		for (auto& s : child.sat) {
			for (auto& p : s->predecessors) {
				if (p->successors <= child.sat)
					exclusivePredecessors.insert (p);
			}
		}

		return exclusivePredecessors;
	});
}

std::unique_ptr<S_Formula> SatVisitor::operator() (const Formula::E_AllUntil& exp) const {
	return binary (exp, [&] (S_Formula& left, S_Formula& right) {
		std::set<TS::State*> T = right.sat;
		bool modified;
		do {
			modified = false;
			for (TS::State* s : T) {
				for (TS::State* p : s->predecessors) {

					if (left.sat.find (p) != left.sat.end ()
					&&	s->successors <= T
					&&	T.find (p) != T.end ()) {
						T.insert (p);
						modified = true;
					}
				}
			}
		} while (modified);

		return T;
	});
}

std::unique_ptr<S_Formula> SatVisitor::operator() (const Formula::E_AllAlways& exp) const {
	return unary (exp, [&] (S_Formula& child) {
		std::set<TS::State*> T = child.sat;
		bool modified;
		do {
			modified = false;
			for (TS::State* s : T) {
				if (!(s->successors <= T)) {
					T.erase (s);
					modified = true;
				}
			}
		} while (modified);

		return T;
	});
}

std::unique_ptr<S_Formula> SatVisitor::operator () (const Formula::E_Label& exp) const {
	return nullary (exp, [&] () {
		auto res = tranSys.labels.find (exp.name);
		if (res == tranSys.labels.end ())
			throw std::runtime_error ("Unknown label \"" + exp.name + "\" in formula");

		std::set<TS::State*> sat;
		for (const auto s : tranSys.statesSet) {
			if (s->atomicPropositions.find (&res->second) != s->atomicPropositions.end ())
				sat.insert (s);
		}

		return sat;
	});
}

std::unique_ptr<S_Formula> SatVisitor::operator() (const Formula::E_Next&) const {
	throw std::runtime_error ("Illegal node in AST");
}

std::unique_ptr<S_Formula> SatVisitor::operator() (const Formula::E_Until&) const {
	throw std::runtime_error ("Illegal node in AST");
}

std::pair<bool, std::unique_ptr<S_Formula>> computeSat (const Formula::Expression& formula, TS::TranSys& ts) {
	std::unique_ptr<S_Formula> res = boost::apply_visitor (SatVisitor { ts }, formula);
	bool s = ts.init <= res->sat;
	return std::make_pair (s, std::move (res));
}

void S_Formula::print (std::ostream& os, size_t depth) {
	for (size_t i = 0; i < depth; ++i) os << '\t';
	os << "Sat (";
	std::copy (strBegin, strEnd, std::ostreambuf_iterator<char> (os));
	os << ") = ";
}

void S_Nullary::print (std::ostream& os, size_t depth) {
	S_Formula::print (os, depth);
	os << sat << std::endl;
}

void S_Unary::print (std::ostream& os, size_t depth) {
	S_Formula::print (os, depth);
	os << sat << std::endl;
	child->print (os, depth+1);
}

void S_Binary::print (std::ostream& os, size_t depth) {
	S_Formula::print (os, depth);
	os << sat << std::endl;
	left->print (os, depth+1);
	right->print (os, depth+1);
}

std::ostream& operator << (std::ostream& os, const SatSet& sat) {
	os << '{';
	bool next = false;
	for (SatSet::const_iterator i = sat.cbegin (); i != sat.cend (); ++i) {
		if (next) {
			os << ", ";
		} else {
			next = true;
		}
		os << (*i)->name;
	}
	os << '}';
	return os;
}

}


#ifndef HOMOTOPY_BASE_HH
#define HOMOTOPY_BASE_HH

#include <vector>
#include <memory> // unique_ptr
#include <cmath>
#include <fstream>

#include "dune/grid/config.h"

#include "abstract_interface.hh"
#include "algorithm_base.hh"
#include "newton_damped.hh"
#include "modelfunction.hh"
#include <boost/timer/timer.hpp>

namespace Kaskade
{

class NewtonsMethod;

  /** Abstract algorithm of interior point method.
  */
class InteriorPointSqrtModel;
template<class IPF, class DomainVector> class InteriorPointTangentialPredictor;
class InteriorPointSlopeEstimate;

class InteriorPointParameters : IterationParameters
{
public:
  InteriorPointParameters(double mu_, double desiredAccuracy_, double desiredContraction_, int maxSteps_, double muFinal_, double relDist, double accfinal_,
                          double reductionStart_,double reductionFactorBest_, double redWorst_=1-1e-4)
    : IterationParameters(desiredAccuracy_, maxSteps_),
      desiredContraction(desiredContraction_),
      muFinal(muFinal_),
      reductionFactorWorst(redWorst_),
      firstStepFailureFactor(2.0),
      reductionStart(reductionStart_),
      relDistanceToPath(relDist),
      accfinal(accfinal_),
      reductionFactorBest(reductionFactorBest_),
      maxTrialSteps(10)
  {
    reset();
    mu=mu_;
  }

//Parameters with values that must be supplied by client
  double desiredContraction, muFinal, reductionFactorWorst, firstStepFailureFactor, reductionStart, relDistanceToPath, accfinal, reductionFactorBest;
  int maxTrialSteps;

  virtual ~InteriorPointParameters() {}

  int termination;

  void reset()
  {
    IterationParameters::reset();
    stepSuccessful=true;
    correctorTermination=-1;
  }

  LoggedQuantity<double> mu, sigma;

protected:

  bool stepSuccessful;

  int correctorTermination;

  LoggedQuantity<double> deltaMu, j, jp, muTrial;

  friend class HomotopyBase;
  friend class InteriorPointSimple;
  friend class InteriorPointSqrtModel;
  template<class A, class B> friend class InteriorPointTangentialPredictor;
  friend class InteriorPointSlopeEstimate;

  virtual void doForAll(LQAction::ToDo td)
  {
    IterationParameters::doForAll(td);
    mu.doAction(td,"mu");
    muTrial.doAction(td,"muTrial");
    j.doAction(td,"jp");
    jp.doAction(td,"j");
    sigma.doAction(td,"sigma");
    deltaMu.doAction(td,"deltaMu");
  }

};

/// Base class for homotopy methods. Here, the main algorithm is programmed
class HomotopyBase : public Algorithm
  {
  public:
    HomotopyBase(NewtonsMethod& n_, InteriorPointParameters& p_) :
      corrector(n_), p(p_)
    {}
    virtual ~HomotopyBase()
    { }

    void solve(AbstractParameterFunctional* f,AbstractFunctionSpaceElement& x);
  protected:

/// Optional overload by derived class
    virtual void initialize() {};

/// Optional overload by derived class
    virtual void finalize() {};

    virtual void logQuantities();

/// Default: classical predictor
    virtual void computePredictor(int step);

/// New mu if corrector was successful
    virtual double muOnSuccess(int step) = 0;

/// New mu if corrector failed
    virtual double muOnFailure() = 0;

/// mu for successful termination
    virtual double muFinal() = 0;

/// update of path parameters, e.g. eta and omega
    virtual void updateModelOfHomotopy(int step) = 0;

/// what should be done before corrector
    virtual void initializeCorrector() = 0;

/// what should be done after corrector
    virtual void finalizeCorrector() = 0;

/// make trial iterate to current iterate (on successful corrector)
    virtual void updateIterate() = 0;

/// make old iterate to current iterate (on failed corrector)
    virtual void recoverIterate() = 0;

/// estimate of length of homopoty path
    virtual double lengthOfPath() = 0;

/// convergence test. returns true if method converged and solution is found.
    virtual int convergenceTest();

/// to be performed, if convergence of path has occured
    virtual void finalizeHomotopy(){};

/// output of a termination message
    virtual void terminationMessage(int errorFlag);

    NewtonsMethod& corrector;
    InteriorPointParameters& p;
    AbstractParameterFunctional* functional;

    std::unique_ptr<AbstractFunctionSpaceElement> trialIterate;
    int step;

  private:

    bool stepSuccessful();

    void computeCorrector();
    void computeGapParameter();
    int runAlgorithm();
  };

/// Very simple implementation of homotopy method: fixed stepsize
class InteriorPointSimple : public HomotopyBase
{
public:
  InteriorPointSimple(NewtonsMethod& n_, InteriorPointParameters& p_) :
    HomotopyBase(n_, p_) {}

  virtual double muOnSuccess(int step) { return p.mu*p.reductionStart;};
  virtual double muOnFailure() { return p.mu*p.reductionStart;};
  virtual double muFinal() { return 1e-7;};

  virtual void updateModelOfHomotopy(int step) {};

  virtual void initializeCorrector() {};
  virtual void finalizeCorrector() {};

  virtual void updateIterate() {};
  virtual void recoverIterate() {};

  virtual double lengthOfPath() { return std::sqrt(p.mu); };
};

class InteriorPointParametersSqrt : IterationParameters
{
public:
  InteriorPointParametersSqrt()
    : IterationParameters(0.0,0)
  {
    reset();
  }
  friend class InteriorPointSqrtModel;
  template<class A, class B> friend class InteriorPointTangentialPredictor;
  friend class InteriorPointSlopeEstimate;

  virtual ~InteriorPointParametersSqrt() {}

protected:

  LoggedQuantity<double> omega;
  LoggedQuantity<double> eta;
  LoggedQuantity<double> slope;
  LoggedQuantity<double> curvature;
  LoggedQuantity<double> jpl;
  LoggedQuantity<double> accuracyCorrector;
  LoggedQuantity<int> newtonSum;
  LoggedQuantity<double> timemeasured;

  virtual void doForAll(LQAction::ToDo td)
  {
    omega.doAction(td,"omega");
    eta.doAction(td,"eta");
    slope.doAction(td,"slope");
    curvature.doAction(td,"curvature");
    jpl.doAction(td,"jpl");
    accuracyCorrector.doAction(td,"accuracyCorrector");
    newtonSum.doAction(td,"newtonSum");
    timemeasured.doAction(td,"timemeasured");
  }
};

/// InteriorPointMethod with a sqrt-model of the path: eta ~ mu^{-1/2}, omega ~mu^{-1/2}
class InteriorPointSqrtModel : public HomotopyBase
{
public:
  InteriorPointSqrtModel(NewtonsMethod& n_, AbstractNorm const& norm_,
                         InteriorPointParameters& p_, AbstractNewtonDirection& solver_,
                         AbstractNorm const* normPlain_=0) :
    HomotopyBase(n_, p_),
    norm(norm_),
    solver(solver_),
    normPlain(normPlain_)
  {
    if(!normPlain) normPlain = &norm;
  }

  virtual void initialize() { iterate=trialIterate->clone(); nNewton=0;};
  virtual double muOnSuccess(int step);
  virtual double muOnFailure();
  virtual double muFinal() { return p.muFinal;};
  virtual void updateModelOfHomotopy(int step);
  virtual void initializeCorrector();
  virtual void finalizeCorrector() { nNewton += corrector.stepsPerformed(); pp.newtonSum=nNewton;};
  virtual void updateIterate();
  virtual void recoverIterate() { *trialIterate=*iterate; }
  virtual double lengthOfPath() {if(pp.eta.isValid()) return 2*p.mu*pp.slope; else return 1e300;};
  void printDiagnosis();

  virtual void logQuantities()
  {
    HomotopyBase::logQuantities();
    pp.logStep();
    printDiagnosis();
  }

private:
  std::unique_ptr<AbstractFunctionSpaceElement> iterate;
  AbstractNorm const& norm;
  AbstractNewtonDirection& solver;
  InteriorPointParametersSqrt pp;
  JModel jModel;
  JModelLin jModelL;
  AbstractNorm const* normPlain;
  int nNewton;
  boost::timer::cpu_timer overalltime;
};

/// Performs bisection algorithm to find a zero of a
/**************************************************
 * Assumptions
 *
 *  a < b
 *
 *  Equation f with has double operator(double ) monotonically INCREASING
 *
 ************************************************************/

template<class Equation>
double bisection(double a, double b, Equation const& f, double accuracy, int& iterations)
{
  iterations = 0;
  double u;
  do {
    ++iterations;
    u=(a+b)/2;
    if(u == a) return b;
    if(f(u) > 0) b=u; else a=u;
  }
  while(b-a >= accuracy && (a+b)/2 != a && (a+b)/2 != b && iterations<50);
  if(iterations > 48) std::cout << "Warning: bisection algorithm performed > 48 iterations" << std::endl;
  return u;
}

class InteriorPointSlopeEstimate : public HomotopyBase
{

public:
  InteriorPointSlopeEstimate(NewtonsMethod& n_, AbstractNorm const& norm_, InteriorPointParameters& p_,
                          AbstractNewtonDirection& solver_, AbstractNorm const* normPlain_=0,
                          NewtonsMethod* finalsolver_=0) :
    HomotopyBase(n_, p_),
    norm(norm_),
    solver(solver_),
    stp(0),
    normPlain(normPlain_),
    finalsolver(finalsolver_)
  {
    if(!normPlain_) normPlain=&norm;
  }

  virtual void initialize() {
    iterate=trialIterate->clone();
    tangent=trialIterate->clone();
    nNewton=0;
  };
  virtual void finalize() {};

  virtual double muOnSuccess(int step);

  virtual double muOnFailure() {
    p.sigma=0.5+p.sigma/2.0;
    return p.mu*p.sigma;
  };

  virtual double muFinal() { return p.muFinal;};

  virtual void updateModelOfHomotopy(int step);

  virtual void computePredictor(int step){}

  virtual void initializeCorrector();
  virtual void finalizeCorrector() {    nNewton += corrector.stepsPerformed(); pp.newtonSum=nNewton;};

  virtual void finalizeHomotopy();

  virtual void updateIterate();

  virtual void recoverIterate() { *trialIterate=*iterate; }

  virtual double lengthOfPath()
  {
    if(pp.eta.isValid())
      return 2*p.mu*pp.slope;
    else
      return 1e300;
  };

  virtual void logQuantities()
  {
    HomotopyBase::logQuantities();
    pp.logStep();
    printDiagnosis();
  }

  void printDiagnosis();

private:
  std::unique_ptr<AbstractFunctionSpaceElement> iterate, tangent;

  AbstractNorm const& norm;
  AbstractNewtonDirection& solver;
  InteriorPointParametersSqrt pp;
  JModelLin jModelL;
  int stp;
  AbstractNorm const* normPlain;
  NewtonsMethod* finalsolver;
  int nNewton;
  boost::timer::cpu_timer overalltime;
};

}  // namespace Kaskade

#endif


//
// Copyright (c) 2020 Kenshi Abe
//

#include "Trainer.hpp"

#include <fstream>
#include <iostream>
#include <boost/archive/binary_iarchive.hpp>
#include <boost/archive/binary_oarchive.hpp>
#include <boost/filesystem.hpp>
#include <boost/serialization/unordered_map.hpp>
#include "Node.hpp"

namespace Trainer {

/// @param mode variant of CFR algorithm
/// @param seed random seed
/// @param strategyPaths paths to the binary files that represent fixed strategies
template <typename T>
Trainer<T>::Trainer(const std::string &mode, const uint32_t seed, const std::vector<std::string> &strategyPaths) : mEngine(seed), mNodeTouchedCnt(0), mModeStr(mode) {
    mGame = new T(mEngine);
    mFolderPath = "../strategies/" + mGame->name();
    boost::filesystem::create_directories(mFolderPath);
    mFixedStrategies = new std::unordered_map<std::string, Node *>[mGame->playerNum()];
    mUpdate = new bool[mGame->playerNum()];
    for (int i = 0; i < mGame->playerNum(); ++i) {
        if (strategyPaths.size() >= i + 1 && !strategyPaths[i].empty()) {
            std::cout << "load strategy \"" << strategyPaths[i] << "\" as static player " << i << std::endl;
            std::ifstream ifs(strategyPaths[i]);
            boost::archive::binary_iarchive ia(ifs);
            ia >> mFixedStrategies[i];
            ifs.close();
            mUpdate[i] = false;
        } else {
            mUpdate[i] = true;
        }
    }
}

template <typename T>
Trainer<T>::~Trainer() {
    for (auto &itr : mNodeMap) {
        delete itr.second;
    }
    for (int i = 0; i < mGame->playerNum(); ++i) {
        if (mUpdate[i]) {
            continue;
        }
        for (auto &itr : mFixedStrategies[i]) {
            delete itr.second;
        }
    }
    delete[] mFixedStrategies;
    delete[] mUpdate;
    delete mGame;
}

/// @brief Calculate the expected payoff of each player
/// @param game game
/// @param strategies list of strategies for each player
/// @return list of expected payoffs
template <typename T>
std::vector<float> Trainer<T>::CalculatePayoff(const T &game, const std::vector<std::function<const float *(const T &)>> &strategies) {
    // return payoff for terminal states
    if (game.done()) {
        std::vector<float> payoffs(game.playerNum());
        for (int i = 0; i < game.playerNum(); ++i) {
            payoffs[i] = game.payoff(i);
        }
        return payoffs;
    }

    // chance node turn
    const int actionNum = game.actionNum();
    if (game.isChanceNode()) {
        std::vector<float> nodeUtils(game.playerNum());
        for (int i = 0; i < game.playerNum(); ++i) {
            nodeUtils[i] = 0.0f;
        }
        for (int a = 0; a < actionNum; ++a) {
            auto game_cp(game);
            game_cp.step(a);
            const float chanceProbability = game_cp.chanceProbability();
            std::vector<float> utils = CalculatePayoff(game_cp, strategies);
            for (int i = 0; i < game.playerNum(); ++i) {
                nodeUtils[i] += chanceProbability * utils[i];
            }
        }
        return nodeUtils;
    }

    // for each action, recursively calculate payoff with additional history and probability
    const int player = game.currentPlayer();
    std::vector<float> nodeUtils(game.playerNum());
    for (int i = 0; i < game.playerNum(); ++i) {
        nodeUtils[i] = 0.0f;
    }
    for (int a = 0; a < actionNum; ++a) {
        auto game_cp(game);
        game_cp.step(a);
        std::vector<float> utils = CalculatePayoff(game_cp, strategies);
        for (int i = 0; i < game.playerNum(); ++i) {
            nodeUtils[i] += strategies[player](game)[a] * utils[i];
        }
    }
    return nodeUtils;
}

/// @brief Execute the CFR algorithm to compute an approximate Nash equilibrium
/// @param iterations number of iterations of CFR
template <typename T>
void Trainer<T>::train(const int iterations) {
    float utils[mGame->playerNum()];

    for (int i = 0; i < iterations; ++i) {
        for (int p = 0; p < mGame->playerNum(); ++p) {
            if (!mUpdate[p]) {
                continue;
            }
            if (mModeStr == "vanilla") {
                mGame->reset(false);
                utils[p] = CFR(*mGame, p, 1.0f, 1.0f);
            } else {
                mGame->reset();
                if (mModeStr == "chance") {
                    utils[p] = chanceSamplingCFR(*mGame, p, 1.0f, 1.0f);
                } else if (mModeStr == "external") {
                    utils[p] = externalSamplingCFR(*mGame, p);
                } else if (mModeStr == "outcome") {
                    utils[p] = std::get<0>(outcomeSamplingCFR(*mGame, p, i, 1.0f, 1.0f, 1.0f));
                } else {
                    assert(false);
                }
            }
        }
        if (i % 1000 == 0) {
            std::cout << "iteration:" << i << ", cumulative nodes touched: " << mNodeTouchedCnt << ", infosets num: " << mNodeMap.size() << ", expected payoffs: (";
            for (int p = 0; p < mGame->playerNum(); ++p) {
                std::cout << utils[p] << ",";
            }
            std::cout << ")" << std::endl;
        }
        if (i != 0 && i % 10000000 == 0) {
            writeStrategyToBin(i);
        }
    }

    writeStrategyToBin();
}

/// @brief Main procedure of vanilla CFR
/// @param game game
/// @param playerIndex player whose strategy is updated in the current iteration
/// @param pi the probability of reaching the current game node if all players other than the acting player always choose actions leading to the current game node
/// @param po the probability of reaching the current game node if the acting player always chooses actions leading to the current game node
/// @return expected payoff of the specified player at the current game node
template <typename T>
float Trainer<T>::CFR(const T &game, const int playerIndex, const float pi, const float po) {
    ++mNodeTouchedCnt;

    // return payoff for terminal states
    if (game.done()) {
        return game.payoff(playerIndex);
    }

    // chance node turn
    const int actionNum = game.actionNum();
    if (game.isChanceNode()) {
        float nodeUtil = 0.0f;
        for (int a = 0; a < actionNum; ++a) {
            auto game_cp(game);
            game_cp.step(a);
            const float chanceProbability = game_cp.chanceProbability();
            nodeUtil += chanceProbability * CFR(game_cp, playerIndex, pi, po * chanceProbability);
        }
        return nodeUtil;
    }

    // get information set string representation
    std::string infoSet = game.infoSetStr();

    // treat static player as chance node
    const int player = game.currentPlayer();
    if (!mUpdate[player]) {
        float nodeUtil = 0.0f;
        for (int a = 0; a < actionNum; ++a) {
            auto game_cp(game);
            game_cp.step(a);
            const float chanceProbability = float(mFixedStrategies[player].at(infoSet)->averageStrategy()[a]);
            nodeUtil += chanceProbability * CFR(game_cp, playerIndex, pi, po * chanceProbability);
        }
        return nodeUtil;
    }

    // get information set node or create it if nonexistant
    Node *node = mNodeMap[infoSet];
    if (node == nullptr) {
        node = new Node(actionNum);
        mNodeMap[infoSet] = node;
    }

    // get current strategy through regret-matching
    const float *strategy = node->strategy();

    // for each action, recursively call CFR with additional history and probability
    float utils[actionNum];
    float nodeUtil = 0;
    for (int a = 0; a < actionNum; ++a) {
        auto game_cp(game);
        game_cp.step(a);
        if (player == playerIndex) {
            utils[a] = CFR(game_cp, playerIndex, pi * strategy[a], po);
        } else {
            utils[a] = CFR(game_cp, playerIndex, pi, po * strategy[a]);
        }
        nodeUtil += strategy[a] * utils[a];
    }

    if (player == playerIndex) {
        // for each action, compute and accumulate counterfactual regret
        for (int a = 0; a < actionNum; ++a) {
            const float regret = utils[a] - nodeUtil;
            const float regretSum = node->regretSum(a) + po * regret;
            node->regretSum(a, regretSum);
        }
        // update average strategy across all training iterations
        node->strategySum(strategy, pi);
    }

    return nodeUtil;
}

/// @brief Main procedure of chance-sampling MCCFR
/// @param game game
/// @param playerIndex player whose strategy is updated in the current iteration
/// @param pi the probability of reaching the current game node if all players other than the acting player always choose actions leading to the current game node
/// @param po the probability of reaching the current game node if the acting player and the chance player always choose actions leading to the current game node
/// @return estimated expected payoff of the specified player at the current game node
template <typename T>
float Trainer<T>::chanceSamplingCFR(const T &game, const int playerIndex, const float pi, const float po) {
    ++mNodeTouchedCnt;

    // return payoff for terminal states
    if (game.done()) {
        return game.payoff(playerIndex);
    }

    // get information set string representation
    std::string infoSet = game.infoSetStr();

    // treat static player as chance node
    const int actionNum = game.actionNum();
    const int player = game.currentPlayer();
    if (!mUpdate[player]) {
        auto game_cp(game);
        auto strategy = mFixedStrategies[player].at(infoSet)->averageStrategy();
        std::discrete_distribution<int> dist(strategy, strategy + actionNum);
        game_cp.step(dist(mEngine));
        return chanceSamplingCFR(game_cp, playerIndex, pi, po);
    }

    // get information set node or create it if nonexistant
    Node *node = mNodeMap[infoSet];
    if (node == nullptr) {
        node = new Node(actionNum);
        mNodeMap[infoSet] = node;
    }

    // get current strategy through regret-matching
    const float *strategy = node->strategy();

    // for each action, recursively call cfr with additional history and probability
    float utils[actionNum];
    float nodeUtil = 0;
    for (int a = 0; a < actionNum; ++a) {
        auto game_cp(game);
        game_cp.step(a);
        if (player == playerIndex) {
            utils[a] = chanceSamplingCFR(game_cp, playerIndex, pi * strategy[a], po);
        } else {
            utils[a] = chanceSamplingCFR(game_cp, playerIndex, pi, po * strategy[a]);
        }
        nodeUtil += strategy[a] * utils[a];
    }

    if (player == playerIndex) {
        // for each action, compute and accumulate counterfactual regret
        for (int a = 0; a < actionNum; ++a) {
            const float regret = utils[a] - nodeUtil;
            const float regretSum = node->regretSum(a) + po * regret;
            node->regretSum(a, regretSum);
        }
        // update average strategy across all training iterations
        node->strategySum(strategy, pi);
    }

    return nodeUtil;
}

/// @brief Main procedure of external-sampling MCCFR
/// @param game game
/// @param playerIndex player whose strategy is updated in the current iteration
/// @return estimated expected payoff of the specified player at the current game node
template <typename T>
float Trainer<T>::externalSamplingCFR(const T &game, const int playerIndex) {
    ++mNodeTouchedCnt;

    // return payoff for terminal states
    if (game.done()) {
        return game.payoff(playerIndex);
    }

    // get information set string representation
    std::string infoSet = game.infoSetStr();

    // external sampling with stochastically-weighted averaging cannot treat static player
    const int actionNum = game.actionNum();
    const int player = game.currentPlayer();
    assert(mUpdate[player] && "External sampling with stochastically-weighted averaging cannot treat static player.");

    // get information set node or create it if nonexistant
    Node *node = mNodeMap[infoSet];
    if (node == nullptr) {
        node = new Node(actionNum);
        mNodeMap[infoSet] = node;
    }

    // get current strategy through regret-matching
    const float *strategy = node->strategy();

    // if current player is not the target player, sample a single action and recursively call cfr
    if (player != playerIndex) {
        auto game_cp(game);
        std::discrete_distribution<int> dist(strategy, strategy + actionNum);
        game_cp.step(dist(mEngine));
        const float util = externalSamplingCFR(game_cp, playerIndex);
        // update average strategy
        node->strategySum(strategy, 1.0f);
        return util;
    }

    // for each action, recursively call cfr with additional history and probability
    float utils[actionNum];
    float nodeUtil = 0;
    for (int a = 0; a < actionNum; ++a) {
        auto game_cp(game);
        game_cp.step(a);
        utils[a] = externalSamplingCFR(game_cp, playerIndex);
        nodeUtil += strategy[a] * utils[a];
    }

    // for each action, compute and accumulate counterfactual regret
    for (int a = 0; a < actionNum; ++a) {
        const float regret = utils[a] - nodeUtil;
        const float regretSum = node->regretSum(a) + regret;
        node->regretSum(a, regretSum);
    }

    return nodeUtil;
}

/// @brief Main procedure of outcome-sampling MCCFR
/// @param game game
/// @param playerIndex player whose strategy is updated in the current iteration
/// @param pi the probability of reaching the current game node if all players other than the acting player always choose actions leading to the current game node
/// @param po the probability of reaching the current game node if the acting player and the chance player always choose actions leading to the current game node
/// @param s the probability of reaching the current game node if the chance player always chooses actions leading to the current game node and the other players act according to the sample profile
/// @return estimated expected payoff of the specified player at the current game node, and the probability of reaching the terminal game node if the chance player always chooses actions leading to the terminal game node
template <typename T>
std::tuple<float, float> Trainer<T>::outcomeSamplingCFR(const T &game, const int playerIndex, const int iteration , const float pi, const float po, const float s) {
    ++mNodeTouchedCnt;

    // return payoff for terminal states
    if (game.done()) {
        return std::make_tuple(game.payoff(playerIndex) / s, 1.0f);
    }

    // get information set string representation
    std::string infoSet = game.infoSetStr();

    // outcome sampling with stochastically-weighted averaging cannot treat static player
    const int actionNum = game.actionNum();
    const int player = game.currentPlayer();
    assert(mUpdate[player] && "Outcome sampling with stochastically-weighted averaging cannot treat static player.");

    // get information set node or create it if nonexistant
    Node *node = mNodeMap[infoSet];
    if (node == nullptr) {
        node = new Node(actionNum);
        mNodeMap[infoSet] = node;
    }

    // get current strategy through regret-matching
    const float *strategy = node->strategy();

    // if current player is the target player, sample a single action according to epsilon-on-policy
    // otherwise, sample a single action according to the player's strategy
    const float epsilon = 0.6;
    float probability[actionNum];
    if (player == playerIndex) {
        for (int a = 0; a < actionNum; ++a) {
            probability[a] = (epsilon / (float) actionNum) + (1.0f - epsilon) * strategy[a];
        }
    } else {
        for (int a = 0; a < actionNum; ++a) {
            probability[a] = strategy[a];
        }
    }
    std::discrete_distribution<int> dist(probability, probability + actionNum);
    const int action = dist(mEngine);

    // for sampled action, recursively call cfr with additional history and probability
    float util, pTail;
    auto game_cp(game);
    game_cp.step(action);
    const float newPi = pi * (player == playerIndex ? strategy[action] : 1.0f);
    const float newPo = po * (player == playerIndex ? 1.0f : strategy[action]);
    std::tuple<float, float> ret = outcomeSamplingCFR(game_cp, playerIndex, iteration, newPi, newPo, s * probability[action]);
    util = std::get<0>(ret);
    pTail = std::get<1>(ret);
    if (player == playerIndex) {
        // for each action, compute and accumulate counterfactual regret
        const float W = util * po;
        for (int a = 0; a < actionNum; ++a) {
            const float regret = a == action ? W * (1.0f - strategy[action]) * pTail : -W * pTail * strategy[action];
            const float regretSum = node->regretSum(a) + regret;
            node->regretSum(a, regretSum);
        }
    } else {
        // update average strategy
        node->strategySum(strategy, po / s);
    }
    return std::make_tuple(util, pTail * strategy[action]);
}

/// @brief Save the current average strategy as a binary file
/// @param iteration current iteration
template <typename T>
void Trainer<T>::writeStrategyToBin(const int iteration) const {
    for (auto &itr : mNodeMap) {
//        std::cout << itr.first << ":";
        for (int i = 0; i < itr.first.size(); ++i) {
            std::cout << int(itr.first[i]);
        }
        std::cout << ":";
        for (int i = 0; i < itr.second->actionNum(); ++i) {
            std::cout << itr.second->averageStrategy()[i] << ",";
        }
        std::cout << std::endl;
    }
    std::string path = iteration > 0 ? "strategy_" + std::to_string(iteration)
                                     : "strategy";
    path += "_" + mModeStr + ".bin";
    std::ofstream ofs(mFolderPath + "/" + path);
    boost::archive::binary_oarchive oa(ofs);
    oa << mNodeMap;
    ofs.close();
}

} // namespace


// Copyright 2018 The Simons Foundation, Inc. - All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//    http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#ifndef NETKET_LOCAL_OPERATOR_HPP
#define NETKET_LOCAL_OPERATOR_HPP

#include <mpi.h>
#include <Eigen/Dense>
#include <algorithm>
#include <cassert>
#include <complex>
#include <iostream>
#include <iterator>
#include <limits>
#include <map>
#include <vector>
#include "Hilbert/abstract_hilbert.hpp"
#include "Utils/kronecker_product.hpp"
#include "Utils/next_variation.hpp"
#include "abstract_operator.hpp"

namespace netket {

/**
    Class for local operators acting on a list of sites and for generic local
    Hilbert spaces.
*/

class LocalOperator : public AbstractOperator {
 public:
  using MelType = Complex;
  using MatType = std::vector<std::vector<MelType>>;
  using SiteType = std::vector<int>;
  using MapType = std::map<std::vector<double>, int>;
  using StateType = std::vector<std::vector<double>>;
  using ConnType = std::vector<std::vector<int>>;
  using VectorType = AbstractOperator::VectorType;
  using VectorRefType = AbstractOperator::VectorRefType;
  using VectorConstRefType = AbstractOperator::VectorConstRefType;

 private:
  const AbstractHilbert &hilbert_;
  std::vector<MatType> mat_;
  std::vector<SiteType> sites_;

  std::vector<MapType> invstate_;
  std::vector<StateType> states_;
  std::vector<ConnType> connected_;

  double constant_;

  std::size_t nops_;

  static constexpr double mel_cutoff_ = 1.0e-6;

 public:
  explicit LocalOperator(const AbstractHilbert &hilbert, double constant = 0.)
      : hilbert_(hilbert), constant_(constant), nops_(0) {}

  explicit LocalOperator(const AbstractHilbert &hilbert,
                         const std::vector<MatType> &mat,
                         const std::vector<SiteType> &sites,
                         double constant = 0.)
      : hilbert_(hilbert), constant_(constant) {
    for (std::size_t i = 0; i < mat.size(); i++) {
      Push(mat[i], sites[i]);
    }
    Init();
  }

  explicit LocalOperator(const AbstractHilbert &hilbert, const MatType &mat,
                         const SiteType &sites, double constant = 0.)
      : hilbert_(hilbert), constant_(constant) {
    Push(mat, sites);
    // TODO sort sites and swap columns of mat accordingly
    Init();
  }

  void Push(const MatType &mat, const SiteType &sites) {
    std::size_t found = std::distance(
        sites_.begin(), std::find(sites_.begin(), sites_.end(), sites));
    if (found < sites_.size()) {
      for (std::size_t i = 0; i < mat_[found].size(); i++) {
        for (std::size_t j = 0; j < mat_[found][i].size(); j++) {
          mat_[found][i][j] += mat[i][j];
        }
      }
    } else {
      mat_.push_back(mat);
      sites_.push_back(sites);
    }
  }

  void Init() {
    if (!hilbert_.IsDiscrete()) {
      throw InvalidInputError(
          "Cannot construct operators on infinite local hilbert spaces");
    }
    if (sites_.size() != mat_.size()) {
      throw InvalidInputError(
          "Operators must specify a consistent set of acting_on specifiers");
    }

    nops_ = mat_.size();

    connected_.clear();
    states_.clear();
    invstate_.clear();

    connected_.resize(nops_);
    states_.resize(nops_);
    invstate_.resize(nops_);

    for (std::size_t op = 0; op < nops_; op++) {
      const auto sites = sites_[op];
      const auto mat = mat_[op];
      auto &connected = connected_[op];
      auto &states = states_[op];
      auto &invstate = invstate_[op];

      if (*std::max_element(sites.begin(), sites.end()) >= hilbert_.Size() ||
          *std::min_element(sites.begin(), sites.end()) < 0) {
        throw InvalidInputError("Operator acts on an invalid set of sites");
      }

      auto localstates = hilbert_.LocalStates();
      const auto localsize = localstates.size();

      // Finding the non-zero matrix elements
      const double epsilon = mel_cutoff_;

      connected.resize(mat.size());

      if (mat.size() != std::pow(localsize, sites.size())) {
        throw InvalidInputError(
            "Matrix size in operator is inconsistent with Hilbert space");
      }

      for (std::size_t i = 0; i < mat.size(); i++) {
        for (std::size_t j = 0; j < mat[i].size(); j++) {
          if (mat.size() != mat[i].size()) {
            throw InvalidInputError(
                "Matrix size in operator is inconsistent with Hilbert space");
          }

          if (i != j && std::abs(mat[i][j]) > epsilon) {
            connected[i].push_back(j);
          }
        }
      }

      // Construct the mapping
      // Internal index -> State
      std::vector<double> st(sites.size(), 0);

      do {
        states.push_back(st);
      } while (netket::next_variation(st.begin(), st.end(), localsize - 1));

      for (std::size_t i = 0; i < states.size(); i++) {
        for (std::size_t k = 0; k < states[i].size(); k++) {
          states[i][k] = localstates[states[i][k]];
        }
      }

      // Now construct the inverse mapping
      // State -> Internal index
      std::size_t k = 0;
      for (auto state : states) {
        invstate[state] = k;
        k++;
      }

      assert(k == mat.size());
    }
  }

  void FindConn(VectorConstRefType v, std::vector<Complex> &mel,
                std::vector<std::vector<int>> &connectors,
                std::vector<std::vector<double>> &newconfs) const override {
    assert(v.size() == hilbert_.Size());

    connectors.clear();
    newconfs.clear();
    mel.clear();

    connectors.resize(1);
    newconfs.resize(1);
    mel.resize(1);

    mel[0] = constant_;
    connectors[0].resize(0);
    newconfs[0].resize(0);

    for (std::size_t opn = 0; opn < nops_; opn++) {
      int st1 = StateNumber(v, opn);

      assert(st1 < int(mat_[opn].size()));
      assert(st1 < int(connected_[opn].size()));

      mel[0] += (mat_[opn][st1][st1]);

      // off-diagonal part
      for (auto st2 : connected_[opn][st1]) {
        connectors.push_back(sites_[opn]);
        assert(st2 < int(states_[opn].size()));
        newconfs.push_back(states_[opn][st2]);
        mel.push_back(mat_[opn][st1][st2]);
      }
    }
  }

  // FindConn for a specific operator
  void FindConn(std::size_t opn, VectorConstRefType v,
                std::vector<Complex> &mel,
                std::vector<std::vector<int>> &connectors,
                std::vector<std::vector<double>> &newconfs) const {
    assert(opn < mat_.size() && opn >= 0);

    mel.resize(1, 0.);
    connectors.resize(1);
    newconfs.resize(1);

    int st1 = StateNumber(v, opn);
    assert(st1 < int(mat_[opn].size()));
    assert(st1 < int(connected_[opn].size()));

    mel[0] = (mat_[opn][st1][st1]);

    // off-diagonal part
    for (auto st2 : connected_[opn][st1]) {
      connectors.push_back(sites_[opn]);
      assert(st2 < int(states_[opn].size()));
      newconfs.push_back(states_[opn][st2]);
      mel.push_back(mat_[opn][st1][st2]);
    }
  }

  inline int StateNumber(VectorConstRefType v, int opn) const {
    // TODO use a mask instead of copies
    std::vector<double> state(sites_[opn].size());
    for (std::size_t i = 0; i < sites_[opn].size(); i++) {
      state[i] = v(sites_[opn][i]);
    }
    return invstate_[opn].at(state);
  }

  // Product of two local operators, performing KroneckerProducts as necessary
  friend LocalOperator operator*(const LocalOperator &lhs,
                                 const LocalOperator &rhs) {
    // TODO
    // assert(lhs.Hilbert() == rhs.Hilbert());
    // check if sites have intersections, in that case this algorithm is wrong
    std::vector<MatType> mat;
    std::vector<SiteType> sites;

    for (std::size_t opn = 0; opn < lhs.mat_.size(); opn++) {
      for (std::size_t opn1 = 0; opn1 < rhs.mat_.size(); opn1++) {
        if (lhs.sites_[opn] == rhs.sites_[opn1]) {
          mat.push_back(netket::MatrixProduct(lhs.mat_[opn], rhs.mat_[opn1]));
          sites.push_back(lhs.sites_[opn]);
        } else {
          mat.push_back(
              netket::KroneckerProduct(lhs.mat_[opn], rhs.mat_[opn1]));
          SiteType sitesum = lhs.sites_[opn];
          sitesum.insert(sitesum.end(), rhs.sites_[opn1].begin(),
                         rhs.sites_[opn1].end());
          sites.push_back(sitesum);
        }
      }
    }
    auto constant = lhs.constant_ * rhs.constant_;
    auto opret = LocalOperator(lhs.GetHilbert(), mat, sites, constant);

    if (std::abs(lhs.constant_) > mel_cutoff_) {
      opret += lhs.constant_ * rhs;
    }
    if (std::abs(rhs.constant_) > mel_cutoff_) {
      opret += rhs.constant_ * lhs;
    }

    return opret;
  }

  friend LocalOperator operator+(const LocalOperator &lhs,
                                 const LocalOperator &rhs) {
    assert(rhs.hilbert_.LocalStates().size() ==
           lhs.hilbert_.LocalStates().size());

    auto sites = lhs.sites_;
    auto mat = lhs.mat_;

    sites.insert(sites.end(), rhs.sites_.begin(), rhs.sites_.end());
    mat.insert(mat.end(), rhs.mat_.begin(), rhs.mat_.end());

    return LocalOperator(lhs.GetHilbert(), mat, sites,
                         lhs.constant_ + rhs.constant_);
  }

  friend LocalOperator operator+(const LocalOperator &lhs, double constant) {
    auto sites = lhs.sites_;
    auto mat = lhs.mat_;

    return LocalOperator(lhs.GetHilbert(), mat, sites,
                         lhs.constant_ + constant);
  }

  LocalOperator &operator+=(const LocalOperator &rhs) {
    assert(rhs.hilbert_.LocalStates().size() ==
           this->hilbert_.LocalStates().size());

    this->sites_.insert(this->sites_.end(), rhs.sites_.begin(),
                        rhs.sites_.end());
    this->mat_.insert(this->mat_.end(), rhs.mat_.begin(), rhs.mat_.end());
    this->constant_ += rhs.constant_;
    this->Init();

    return *this;
  }

  LocalOperator &operator+=(double constant) {
    this->constant_ += constant;
    return *this;
  }

  template <class T>
  friend LocalOperator operator*(T lhs, const LocalOperator &rhs) {
    assert(std::imag(lhs) == 0.);
    auto mat = rhs.mat_;
    auto sites = rhs.sites_;

    for (std::size_t opn = 0; opn < mat.size(); opn++) {
      for (std::size_t i = 0; i < mat[opn].size(); i++) {
        for (std::size_t j = 0; j < mat[opn][i].size(); j++)
          mat[opn][i][j] *= lhs;
      }
    }

    return LocalOperator(rhs.GetHilbert(), mat, sites,
                         std::real(lhs * rhs.constant_));
  }

  const std::vector<MatType> &LocalMatrices() const { return mat_; }
  const std::vector<SiteType> &ActingOn() const { return sites_; }

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

  std::size_t Size() const { return mat_.size(); }
};  // namespace netket

}  // namespace netket
#endif


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

  @copyright 2016 NumScale SAS

  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_COS_HPP_INCLUDED
#define BOOST_SIMD_FUNCTION_COS_HPP_INCLUDED

#if defined(DOXYGEN_ONLY)
namespace boost { namespace simd
{

 /*!

    @ingroup group-trigonometric
    Function object implementing cos capabilities

    cosine of the input in radians.

    @par Semantic:

    For every parameter of floating type

    @code
    auto r = cos(x);
    @endcode

    @par Notes
    @c cos can be called with two parameters as :
    @code
    T r = cos(x, range_);
    @endcode

    or with a decorator as:

    @code
    T r = std_(cos)(x);
    @endcode

    or

    @code
    T r = restricted_(cos)(x);
    @endcode

    range is a tag that allows some control on the computation
    accuracy and speed.

    The control is on the reduction routine of the angle to the
    \f$[-\pi/4, \pi/4]\f$ interval
    they actually are 3 reduction routines protocol that respectly are
    sufficient for small_, medium_ and big_ angle values to have within
    cover test one ulp of difference with the according crlibm
    (correctly rounded math library) result.

    Each one covers respectively intervals \f$[-A, A]\f$ with :
    @code
    |--------------|---------------|-------------|
    |              |   float  A    |  double  A  |
    |--------------|---------------|-------------|
    |    small_    |   20*Pi       | 20*Pi       |
    |--------------|---------------|-------------|
    |    medium_   |   2^6*Pi      | 2^18*Pi     |
    |--------------|---------------|-------------|
    |    big_      |   Inf         | Inf         |
    |--------------|-----------------------------|
    @endcode

    In fact for each scalar singleton or simd vector of angles
    there are two possibilities :
    \arg one is to test if all vector element(s) are in the proper range
    for the consecutive increasing values of A until we reach a good
    one or the last :
    the corresponding template tags are small_, medium_ and big_
    \arg the second is to force directly a reduction method:
    the corresponding template tags are direct_small_, direct_medium_
    and direct_big_
    @par
    direct_small_ is NOT equivalent to small_ because there are also
    two other methods for \f$[0, \pi/4]\f$ (no reduction)
    and \f$[\pi/4, \pi/2]\f$
    (straight reduction) that are not considered in direct small_

    Note that for float the direct_big_ case is both early an hyper costly and
    shall be avoided whenever possible. To partially achieve this aim
    when double are available on the platform, this part of reduction
    is delegated to the double precision routines.

    @par Advices
    \arg If there is no restrictions ever on your angles and you care for precision
    use the default cos(x) or equivalently cos(x, big_)
    \arg if you do not care for precision you can use
    cos(x, medium_) or cos(x, small_)
    that will be accurate for their proper range and degrade in accuracy
    with greater values.

    @par
    Now, the choice of direct or not relies on probabilities
    computations:
    assuming that a vector contains k elements and that testing all
    values that are in an interval takes c cycles and the probability of a value
    to be in interval \f$[a, b]\f$ is \f$p(a, b)\f$
    the number of cycles used by a
    direct\f${}_i\f$ method is simply the reduction time:
    \f$N(\f$direct\f${}_i)\f$
    On the other side the number of cycles for the non-direct methods will have a more
    complicated expression :

    \f$\hspace{5em}\sum_{i=1}^{m} p(A_{i-1}, A_i)^k N(\f$direct\f${}_i)\f$

    @par
    So the non direct methods will be interesting only if you want accurate
    results everywhere and have anyhow a big proportion of small angles.
    This is even more true (if possible) in simd and the more k is big, because
    of the kth power.

    @par
    For instance in the medium_ float case:

    \arg if angles are equidistributed  on \f$[0, 2^{16} \pi]\f$,
    the \f$p(0, 20\pi)\f$ will be
    less than \f$2^{-11}\f$ and thus in sse2 there will be 1 quadruple over
    1.76e+13 falling in the small_ case...

    \arg Even sorting will do no good because the sort cost will be against
    the ratio of 1 successful quadruplet over 2048.
    \arg  Contrarily if your angles have a gaussian distribution with 0 mean and
    \f$10\pi\f$ standard deviation,  80% of the intervals will be in the "small_"
    case (95% of the values).
    \arg Finally for those that are sure of their angles taking place in a fixed
    range and want speed, three other template tags can be of choice:

                clipped_very_small_, clipped_small_ and clipped_medium_

    @par
    they use the chosen reduction, but return Nan for any outsider.

    @par Decorators

    - std_ provides access to std::cos

    - restricted_ is equivalent to the clipped_very_small_ tag

    @see  sincos, cosd, cospi

  **/
  Value cos(Value const & v0);
} }
#endif

#include <boost/simd/function/scalar/cos.hpp>
#include <boost/simd/function/simd/cos.hpp>

#endif