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

#ifndef BOOST_MATH_ELLINT_HL_HPP
#define BOOST_MATH_ELLINT_HL_HPP

#ifdef _MSC_VER
#pragma once
#endif

#include <boost/math/special_functions/math_fwd.hpp>
#include <boost/math/special_functions/ellint_rj.hpp>
#include <boost/math/special_functions/ellint_rj.hpp>
#include <boost/math/special_functions/ellint_1.hpp>
#include <boost/math/special_functions/jacobi_zeta.hpp>
#include <boost/math/constants/constants.hpp>
#include <boost/math/policies/error_handling.hpp>
#include <boost/math/tools/workaround.hpp>

// Elliptic integral the Jacobi Zeta function.

namespace boost { namespace math {

namespace detail{

// Elliptic integral - Jacobi Zeta
template <typename T, typename Policy>
T heuman_lambda_imp(T phi, T k, const Policy& pol)
{
    BOOST_MATH_STD_USING
    using namespace boost::math::tools;
    using namespace boost::math::constants;

    const char* function = "boost::math::heuman_lambda<%1%>(%1%, %1%)";

    if(fabs(k) > 1)
       return policies::raise_domain_error<T>(function, "We require |k| <= 1 but got k = %1%", k, pol);

    T result;
    T sinp = sin(phi);
    T cosp = cos(phi);
    T s2 = sinp * sinp;
    T k2 = k * k;
    T kp = 1 - k2;
    T delta = sqrt(1 - (kp * s2));
    if(fabs(phi) <= constants::half_pi<T>())
    {
       result = kp * sinp * cosp / (delta * constants::half_pi<T>());
       result *= ellint_rf_imp(T(0), kp, T(1), pol) + k2 * ellint_rj(T(0), kp, T(1), T(1 - k2 / (delta * delta)), pol) / (3 * delta * delta);
    }
    else
    {
       T rkp = sqrt(kp);
       T ratio;
       if(rkp == 1)
       {
          return policies::raise_domain_error<T>(function, "When 1-k^2 == 1 then phi must be < Pi/2, but got phi = %1%", phi, pol);
       }
       else
          ratio = ellint_f_imp(phi, rkp, pol) / ellint_k_imp(rkp, pol);
       result = ratio + ellint_k_imp(k, pol) * jacobi_zeta_imp(phi, rkp, pol) / constants::half_pi<T>();
    }
    return result;
}

} // detail

template <class T1, class T2, class Policy>
inline typename tools::promote_args<T1, T2>::type heuman_lambda(T1 k, T2 phi, const Policy& pol)
{
   typedef typename tools::promote_args<T1, T2>::type result_type;
   typedef typename policies::evaluation<result_type, Policy>::type value_type;
   return policies::checked_narrowing_cast<result_type, Policy>(detail::heuman_lambda_imp(static_cast<value_type>(phi), static_cast<value_type>(k), pol), "boost::math::heuman_lambda<%1%>(%1%,%1%)");
}

template <class T1, class T2>
inline typename tools::promote_args<T1, T2>::type heuman_lambda(T1 k, T2 phi)
{
   return boost::math::heuman_lambda(k, phi, policies::policy<>());
}

}} // namespaces

#endif // BOOST_MATH_ELLINT_D_HPP


#include <iostream>
#include <Eigen/LU>

#include "irls.h"
#include "input_validation.h"
#include "disclap_family.h"
#include "beta.h"

using namespace Eigen;

void check_input(const std::vector<MatrixXi>& d, const std::vector<MatrixXd>& w, VectorXd& beta, size_t maxit, bool use_deviance, double eps_deviance, bool use_beta, double eps_beta) {
  check_common_input(maxit, use_deviance, eps_deviance, use_beta, eps_beta);
  check_d_w(d, w);
  check_beta(beta, d, w);
}

void update_lin_pred(VectorXd& lin_pred, const VectorXd& beta_ext, size_t individuals, size_t clusters, size_t loci) {
  size_t idx = 0;

  for (size_t k = 0; k < loci; k++) {
    double b_k = beta_ext[clusters + k];

    for (size_t j = 0; j < clusters; j++) {
      double b_j = beta_ext[j];

      double bj_bk = b_j + b_k;

      for (size_t i = 0; i < individuals; i++) {
        lin_pred[idx] = bj_bk;
        idx += 1;
      }
    }
  }
}

bool deviance_converged(const std::vector<MatrixXi>& d, const std::vector<MatrixXd>& w,
  bool& deviance_calculated,
  VectorXd& lin_pred, const VectorXd& beta, double& deviance, double& old_deviance, double eps_deviance,
  size_t individuals, size_t clusters, size_t loci, bool verbose) {

  VectorXd beta_ext = get_beta_extended(beta);

  update_lin_pred(lin_pred, beta_ext, individuals, clusters, loci);

  double dev = 0;

  for (size_t i = 0; i < individuals; i++) {
    MatrixXi d_i = d[i];
    MatrixXd w_i = w[i];

    for (size_t j = 0; j < clusters; j++) {
      double b_j = beta_ext[j];

      for (size_t k = 0; k < loci; k++) {
        double b_k = beta_ext[clusters + k];

        double eta = b_j + b_k;
        double mu = linkinv(eta);
        int d_ele = d_i(j, k);

        double dev_contrib = 0.0;

        if (d_ele == 0) {
          // y == 0: dev = 2*log((1+p)/(1-p))
          double p = (mu < 1e-6) ? 0.5 * mu : (sqrt(1.0 + mu * mu) - 1.0) / mu;
          dev_contrib = 2 * log((1.0 + p) / (1.0 - p));
        } else {
          // y != 0
          double d_ele_double = (double)d_ele;
          dev_contrib = 2 * (loglikeh(d_ele_double, d_ele_double) - loglikeh(mu, d_ele_double));
        }

        dev_contrib *= w_i(j, k);
        dev += dev_contrib;
      }
    }
  }

  old_deviance = deviance;
  deviance = dev;
  deviance_calculated = true;

  double c = std::abs(dev - old_deviance)/(0.1 + std::abs(dev));

  if (verbose) {
    OUTPUT << "      Deviance convergence investigation:" << std::endl;
    OUTPUT << "        deviance     = " << dev << std::endl;
  }

  if (c < eps_deviance) {
    if (verbose) {
      OUTPUT << "        criteria     = " << c << " < " << eps_deviance << " [CONVERGENCE]" << std::endl;
    }

    return true;
  }

  if (verbose) {
    OUTPUT << "        criteria     = " << c << " >= " << eps_deviance << std::endl;
  }

  return false;
}

bool beta_converged(const VectorXd& beta, const VectorXd& old_beta, double eps_beta, bool verbose) {
  VectorXd den = 0.1 + beta.cwiseAbs().array();
  VectorXd diff = beta - old_beta;

  double c = diff.cwiseAbs().cwiseQuotient(den).maxCoeff();

  if (verbose) {
    OUTPUT << "      Beta convergence investigation:" << std::endl;
    OUTPUT << "        beta         = " << beta.transpose() << std::endl;
  }

  if (c < eps_beta) {
    if (verbose) {
      OUTPUT << "        criteria     = " << c << " < " << eps_beta << " [CONVERGENCE]" << std::endl;
    }

    return true;
  }

  if (verbose) {
    OUTPUT << "        criteria     = " << c << " >= " << eps_beta << std::endl;
  }

  return false;
}

void fill_generics(size_t clusters, size_t loci, VectorXd& beta_ext,
  MatrixXd& lin_pred_generic, MatrixXd& mu_generic) {

  for (size_t j = 0; j < clusters; j++) {
    double beta_j = beta_ext[j];

    for (size_t k = 0; k < loci; k++) {
      double eta = beta_j + beta_ext[clusters + k];
      lin_pred_generic(j, k) = eta;
      mu_generic(j, k) = linkinv(eta);
    }
  }
}

void fill_H_ab(size_t individuals, size_t clusters, size_t loci,
  const std::vector<MatrixXi>& d, const std::vector<MatrixXd>& w,
  const MatrixXd& mu_generic,
  MatrixXd& H, MatrixXd& ab) {

  for (size_t i = 0; i < individuals; i++) {
    MatrixXi d_i = d[i];
    MatrixXd w_i = w[i];

    for (size_t j = 0; j < clusters; j++) {
      for (size_t k = 0; k < loci; k++) {
        double mu_generic_jk = mu_generic(j, k);
        double w_ijk = w_i(j, k);

        // H
        double psi_jk = w_ijk * varfunc(mu_generic_jk);
        H(j, k) += psi_jk;

        // ab
        double d_mu_res = d_i(j, k) - mu_generic_jk;
        ab(j, k) += w_ijk * d_mu_res;
      }
    }
  }
}

/*
void fill_H_ab(size_t individuals, size_t clusters, size_t loci,
  const std::vector<MatrixXi>& d, const std::vector<MatrixXd>& w,
  const MatrixXd& mu_generic,
  MatrixXd& H, MatrixXd& ab) {

  // varfunc(double mu) = mu * sqrt(1.0 + mu*mu)
  MatrixXd mu_1_sq = mu_generic.array().square() + 1;
  MatrixXd mu_sqrt_1_sq = mu_1_sq.cwiseSqrt();
  MatrixXd mu_varfunc = mu_generic.cwiseProduct(mu_sqrt_1_sq);

  for (size_t i = 0; i < individuals; i++) {
    MatrixXi d_i = d[i];
    MatrixXd w_i = w[i];
    MatrixXd d_mu = d_i.cast<double>().array() - mu_generic.array();

    H = H.array() + w_i.cwiseProduct(mu_varfunc).array();
    ab = ab.array() + w_i.cwiseProduct(d_mu).array();
  }
}
*/

void update_beta(const std::vector<MatrixXi>& d, const std::vector<MatrixXd>& w, VectorXd& beta,
  MatrixXd& vcov,
  size_t individuals, size_t clusters, size_t loci, bool verbose) {

  VectorXd beta_ext = get_beta_extended(beta);
  MatrixXd lin_pred_generic(clusters, loci);
  MatrixXd mu_generic(clusters, loci);
  fill_generics(clusters, loci, beta_ext, lin_pred_generic, mu_generic);

  MatrixXd H = MatrixXd::Zero(clusters, loci);
  MatrixXd ab = MatrixXd::Zero(clusters, loci);
  fill_H_ab(individuals, clusters, loci, d, w, mu_generic, H, ab);

  VectorXd Drvec_raw = H.colwise().sum();
  VectorXd Drvec(loci - 1);
  for (size_t k = 0; k < (loci - 1); k++) {
    Drvec[k] = Drvec_raw[k];
  }
  MatrixXd Dr = Drvec.asDiagonal();

  VectorXd Dcinvvec(clusters);
  for (size_t j = 0; j < clusters; j++) {
    Dcinvvec[j] = 1 / H.row(j).sum();
  }
  MatrixXd Dcinv = Dcinvvec.asDiagonal();

  H.conservativeResize(NoChange, loci - 1);

  MatrixXd E = Dr - H.transpose() * Dcinv * H;
  MatrixXd Einv = E.lu().solve(MatrixXd::Identity(loci - 1, loci - 1));
  MatrixXd F = Dcinv * H;
  MatrixXd Ft = F.transpose();

  /*
  OUTPUT << "H:" << std::endl << H << std::endl;
  OUTPUT << "Dr:" << std::endl << Dr << std::endl;
  OUTPUT << "Dcinv:" << std::endl << Dcinv << std::endl;
  OUTPUT << "E:" << std::endl << E << std::endl;
  OUTPUT << "Einv:" << std::endl << Einv << std::endl;
  OUTPUT << "F:" << std::endl << F << std::endl;
  OUTPUT << "Ft:" << std::endl << Ft << std::endl;
  */

  MatrixXd P = MatrixXd::Zero(clusters + loci - 1, clusters + loci - 1);
  P.topLeftCorner(clusters, clusters) = Dcinv + (F * Einv) * Ft;
  P.topRightCorner(clusters, loci - 1) = -F * Einv;
  P.bottomLeftCorner(loci - 1, clusters) = -Einv * Ft;
  P.bottomRightCorner(loci - 1, loci - 1) = Einv;

  vcov = P;

  VectorXd a = ab.rowwise().sum();
  VectorXd b = ab.colwise().sum();
  VectorXd gamma(clusters + loci - 1);

  for (size_t j = 0; j < clusters; j++) {
    gamma[j] = a[j];
  }

  for (size_t k = 0; k < (loci - 1); k++) {
    gamma[clusters + k] = b[k];
  }

  VectorXd beta_correction = P * gamma;

  beta = beta + beta_correction;
}

/*
// iter >= ensures old_deviance has been calculated
if (deviance_calculated && iter >= 2 && !isinf(old_deviance) && old_deviance < dev) {
  OUTPUT << "old_deviance = " << old_deviance << std::endl;
  OUTPUT << "dev          = " << dev << std::endl;
  throw "Deviance increased! Normally step-halving would be tried, but it has yet to be implemented...";
}
*/

/*
Returns true if converged, false otherwise
*/
bool irls(const std::vector<MatrixXi>& d, const std::vector<MatrixXd>& w,
  VectorXd& beta, VectorXd& lin_pred,
  MatrixXd& vcov,
  double& deviance, size_t& iterations,
  bool verbose = true, size_t maxit = 25,
  bool use_deviance = true, double eps_deviance = 1e-6,
  bool use_beta = true, double eps_beta = 1e-6,
  bool force_calculate_deviance = false) {

  check_input(d, w, beta, maxit, use_deviance, eps_deviance, use_beta, eps_beta);

  assert(maxit > 0);
  assert(!use_deviance || (use_deviance && eps_deviance > 0));
  assert(!use_beta || (use_beta && eps_beta > 0));

  size_t individuals = d.size();
  assert(individuals > 0);
  assert(w.size() == individuals);

  size_t clusters = d.at(0).rows();
  size_t loci = d.at(0).cols();

  bool converged = false;

  bool deviance_calculated = false;
  double dev = std::numeric_limits<double>::infinity();
  double old_deviance = std::numeric_limits<double>::infinity();

  VectorXd start_beta = VectorXd(beta);

  VectorXd old_beta(beta.size());
  for (size_t idx = 0; idx < old_beta.size(); idx++) {
    old_beta[idx] = std::numeric_limits<double>::infinity();
  }

  if (verbose) {
    OUTPUT << "    Initial coefficients = " << beta.transpose() << std::endl;
  }

  iterations = 0;

  for (size_t iter = 0; iter < maxit; iter++) {
    iterations = iterations + 1;

    if (verbose) {
      OUTPUT << "    IRLS iteration " << (iter + 1) << std::endl;
    }

    if (use_deviance && use_beta) {
      bool beta_conv = beta_converged(beta, old_beta, eps_beta, verbose);

      if (beta_conv) {
        bool dev_conv = deviance_converged(d, w, deviance_calculated, lin_pred, beta, dev, old_deviance, eps_deviance, individuals, clusters, loci, verbose);

        if (dev_conv) {
          converged = true;
          break;
        }
      }
    } else if (use_deviance) {
      bool dev_conv = deviance_converged(d, w, deviance_calculated, lin_pred, beta, dev, old_deviance, eps_deviance, individuals, clusters, loci, verbose);

      if (dev_conv) {
        converged = true;
        break;
      }
    } else if (use_beta) {
      bool beta_conv = beta_converged(beta, old_beta, eps_beta, verbose);

      if (beta_conv) {
        converged = true;
        break;
      }
    } else {
      throw "Unexpected error (!use_deviance && !use_beta)";
    }

    old_beta = beta;
    update_beta(d, w, beta, vcov, individuals, clusters, loci, verbose);

    #if defined(DISCLAPMIX_USED_IN_R)
    R_CheckUserInterrupt();
    #endif
  }

  if (!deviance_calculated) {
    if (force_calculate_deviance) {
      deviance_converged(d, w, deviance_calculated, lin_pred, beta, dev, old_deviance, eps_deviance, individuals, clusters, loci, false);
    }
  }

  deviance = dev;

  return converged;
}


// Copyright  (C)  2007  Ruben Smits <ruben dot smits at mech dot kuleuven dot be>

// Version: 1.0
// Author: Ruben Smits <ruben dot smits at mech dot kuleuven dot be>
// Maintainer: Ruben Smits <ruben dot smits at mech dot kuleuven dot be>
// URL: http://www.orocos.org/kdl

// This library is free software; you can redistribute it and/or
// modify it under the terms of the GNU Lesser General Public
// License as published by the Free Software Foundation; either
// version 2.1 of the License, or (at your option) any later version.

// This library is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
// Lesser General Public License for more details.

// You should have received a copy of the GNU Lesser General Public
// License along with this library; if not, write to the Free Software
// Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA  02110-1301  USA

//Based on the svd of the KDL-0.2 library by Erwin Aertbelien
#ifndef SVD_EIGEN_HH_HPP
#define SVD_EIGEN_HH_HPP

#include <Eigen/Array>
#include <algorithm>

namespace KDL
{
    template<typename Scalar> inline Scalar PYTHAG(Scalar a,Scalar b) {
        double at,bt,ct;
        at = fabs(a);
        bt = fabs(b);
        if (at > bt ) {
            ct=bt/at;
            return Scalar(at*sqrt(1.0+ct*ct));
        } else {
            if (bt==0)
                return Scalar(0.0);
            else {
                ct=at/bt;
                return Scalar(bt*sqrt(1.0+ct*ct));
            }
        }
    }

    template<typename Scalar> inline Scalar SIGN(Scalar a,Scalar b) {
        return ((b) >= Scalar(0.0) ? fabs(a) : -fabs(a));
    }

    /**
     * svd calculation of boost ublas matrices
     *
     * @param A matrix<double>(mxn)
     * @param U matrix<double>(mxn)
     * @param S vector<double> n
     * @param V matrix<double>(nxn)
     * @param tmp vector<double> n
     * @param maxiter defaults to 150
     *
     * @return -2 if maxiter exceeded, 0 otherwise
     */
	template<typename MatrixA, typename MatrixUV, typename VectorS>
	int svd_eigen_HH(
		const Eigen::MatrixBase<MatrixA>&		A,
		Eigen::MatrixBase<MatrixUV>&			U,
		Eigen::MatrixBase<VectorS>&				S,
		Eigen::MatrixBase<MatrixUV>&			V,
		Eigen::MatrixBase<VectorS>&				tmp,
		int maxiter=150)
	{
        //get the rows/columns of the matrix
        const int rows = A.rows();
        const int cols = A.cols();

        U = A;

        int i(-1),its(-1),j(-1),jj(-1),k(-1),nm=0;
        int ppi(0);
        bool flag;
        e_scalar maxarg1,maxarg2,anorm(0),c(0),f(0),h(0),s(0),scale(0),x(0),y(0),z(0),g(0);

        g=scale=anorm=e_scalar(0.0);

        /* Householder reduction to bidiagonal form. */
        for (i=0;i<cols;i++) {
            ppi=i+1;
            tmp(i)=scale*g;
            g=s=scale=e_scalar(0.0);
            if (i<rows) {
                // compute the sum of the i-th column, starting from the i-th row
                for (k=i;k<rows;k++) scale += fabs(U(k,i));
                if (scale!=0) {
                    // multiply the i-th column by 1.0/scale, start from the i-th element
                    // sum of squares of column i, start from the i-th element
                    for (k=i;k<rows;k++) {
                        U(k,i) /= scale;
                        s += U(k,i)*U(k,i);
                    }
                    f=U(i,i);  // f is the diag elem
                    g = -SIGN(e_scalar(sqrt(s)),f);
                    h=f*g-s;
                    U(i,i)=f-g;
                    for (j=ppi;j<cols;j++) {
                        // dot product of columns i and j, starting from the i-th row
                        for (s=0.0,k=i;k<rows;k++) s += U(k,i)*U(k,j);
                        f=s/h;
                        // copy the scaled i-th column into the j-th column
                        for (k=i;k<rows;k++) U(k,j) += f*U(k,i);
                    }
                    for (k=i;k<rows;k++) U(k,i) *= scale;
                }
            }
            // save singular value
            S(i)=scale*g;
            g=s=scale=e_scalar(0.0);
            if ((i <rows) && (i+1 != cols)) {
                // sum of row i, start from columns i+1
                for (k=ppi;k<cols;k++) scale += fabs(U(i,k));
                if (scale!=0) {
                    for (k=ppi;k<cols;k++) {
                        U(i,k) /= scale;
                        s += U(i,k)*U(i,k);
                    }
                    f=U(i,ppi);
                    g = -SIGN(e_scalar(sqrt(s)),f);
                    h=f*g-s;
                    U(i,ppi)=f-g;
                    for (k=ppi;k<cols;k++) tmp(k)=U(i,k)/h;
                    for (j=ppi;j<rows;j++) {
                        for (s=0.0,k=ppi;k<cols;k++) s += U(j,k)*U(i,k);
                        for (k=ppi;k<cols;k++) U(j,k) += s*tmp(k);
                    }
                    for (k=ppi;k<cols;k++) U(i,k) *= scale;
                }
            }
            maxarg1=anorm;
            maxarg2=(fabs(S(i))+fabs(tmp(i)));
            anorm = maxarg1 > maxarg2 ?	maxarg1 : maxarg2;
        }
        /* Accumulation of right-hand transformations. */
        for (i=cols-1;i>=0;i--) {
            if (i<cols-1) {
                if (g) {
                    for (j=ppi;j<cols;j++) V(j,i)=(U(i,j)/U(i,ppi))/g;
                    for (j=ppi;j<cols;j++) {
                        for (s=0.0,k=ppi;k<cols;k++) s += U(i,k)*V(k,j);
                        for (k=ppi;k<cols;k++) V(k,j) += s*V(k,i);
                    }
                }
                for (j=ppi;j<cols;j++) V(i,j)=V(j,i)=0.0;
            }
            V(i,i)=1.0;
            g=tmp(i);
            ppi=i;
        }
        /* Accumulation of left-hand transformations. */
        for (i=cols-1<rows-1 ? cols-1:rows-1;i>=0;i--) {
            ppi=i+1;
            g=S(i);
            for (j=ppi;j<cols;j++) U(i,j)=0.0;
            if (g) {
                g=e_scalar(1.0)/g;
                for (j=ppi;j<cols;j++) {
                    for (s=0.0,k=ppi;k<rows;k++) s += U(k,i)*U(k,j);
                    f=(s/U(i,i))*g;
                    for (k=i;k<rows;k++) U(k,j) += f*U(k,i);
                }
                for (j=i;j<rows;j++) U(j,i) *= g;
            } else {
                for (j=i;j<rows;j++) U(j,i)=0.0;
            }
            ++U(i,i);
        }

        /* Diagonalization of the bidiagonal form. */
        for (k=cols-1;k>=0;k--) { /* Loop over singular values. */
            for (its=1;its<=maxiter;its++) {  /* Loop over allowed iterations. */
                flag=true;
                for (ppi=k;ppi>=0;ppi--) {  /* Test for splitting. */
                    nm=ppi-1;             /* Note that tmp(1) is always zero. */
                    if ((fabs(tmp(ppi))+anorm) == anorm) {
                        flag=false;
                        break;
                    }
                    if ((fabs(S(nm)+anorm) == anorm)) break;
                }
                if (flag) {
                    c=e_scalar(0.0);           /* Cancellation of tmp(l), if l>1: */
                    s=e_scalar(1.);
                    for (i=ppi;i<=k;i++) {
                        f=s*tmp(i);
                        tmp(i)=c*tmp(i);
                        if ((fabs(f)+anorm) == anorm) break;
                        g=S(i);
                        h=PYTHAG(f,g);
                        S(i)=h;
                        h=e_scalar(1.0)/h;
                        c=g*h;
                        s=(-f*h);
                        for (j=0;j<rows;j++) {
                            y=U(j,nm);
                            z=U(j,i);
                            U(j,nm)=y*c+z*s;
                            U(j,i)=z*c-y*s;
                        }
                    }
                }
                z=S(k);

                if (ppi == k) {       /* Convergence. */
                    if (z < e_scalar(0.0)) {   /* Singular value is made nonnegative. */
                        S(k) = -z;
                        for (j=0;j<cols;j++) V(j,k)=-V(j,k);
                    }
                    break;
                }

                x=S(ppi);            /* Shift from bottom 2-by-2 minor: */
                nm=k-1;
                y=S(nm);
                g=tmp(nm);
                h=tmp(k);
                f=((y-z)*(y+z)+(g-h)*(g+h))/(e_scalar(2.0)*h*y);

                g=PYTHAG(f,e_scalar(1.0));
                f=((x-z)*(x+z)+h*((y/(f+SIGN(g,f)))-h))/x;

                /* Next QR transformation: */
                c=s=1.0;
                for (j=ppi;j<=nm;j++) {
                    i=j+1;
                    g=tmp(i);
                    y=S(i);
                    h=s*g;
                    g=c*g;
                    z=PYTHAG(f,h);
                    tmp(j)=z;
                    c=f/z;
                    s=h/z;
                    f=x*c+g*s;
                    g=g*c-x*s;
                    h=y*s;
                    y=y*c;
                    for (jj=0;jj<cols;jj++) {
                        x=V(jj,j);
                        z=V(jj,i);
                        V(jj,j)=x*c+z*s;
                        V(jj,i)=z*c-x*s;
                    }
                    z=PYTHAG(f,h);
                    S(j)=z;
                    if (z) {
                        z=e_scalar(1.0)/z;
                        c=f*z;
                        s=h*z;
                    }
                    f=(c*g)+(s*y);
                    x=(c*y)-(s*g);
                    for (jj=0;jj<rows;jj++) {
                        y=U(jj,j);
                        z=U(jj,i);
                        U(jj,j)=y*c+z*s;
                        U(jj,i)=z*c-y*s;
                    }
                }
                tmp(ppi)=0.0;
                tmp(k)=f;
                S(k)=x;
            }
        }

        //Sort eigen values:
        for (i=0; i<cols; i++){

            double S_max = S(i);
            int i_max = i;
            for (j=i+1; j<cols; j++){
                double Sj = S(j);
                if (Sj > S_max){
                    S_max = Sj;
                    i_max = j;
                }
            }
            if (i_max != i){
                /* swap eigenvalues */
                e_scalar tmp = S(i);
                S(i)=S(i_max);
                S(i_max)=tmp;

                /* swap eigenvectors */
                U.col(i).swap(U.col(i_max));
                V.col(i).swap(V.col(i_max));
            }
        }

        if (its == maxiter)
            return (-2);
        else
            return (0);
    }

}
#endif


// Copyright (c) Dietmar Wolz.
//
// This source code is licensed under the MIT license found in the
// LICENSE file in the root directory.

// Eigen based implementation of active CMA-ES

// Supports parallel fitness function evaluation.
//
// For expensive objective functions (e.g. machine learning parameter optimization) use the workers
// parameter to parallelize objective function evaluation. The workers parameter should be limited
// the population size because otherwize poulation update is delayed.

// Derived from http://cma.gforge.inria.fr/cmaes.m which follows
// https://www.researchgate.net/publication/227050324_The_CMA_Evolution_Strategy_A_Comparing_Review
// Requires Eigen version >= 3.3.90 because new slicing capabilities are used, see
// https://eigen.tuxfamily.org/dox-devel/group__TutorialSlicingIndexing.html
// requires https://github.com/imneme/pcg-cpp

#include <Eigen/Core>
#include <Eigen/Eigenvalues>
#include <iostream>
#include <random>
#include <float.h>
#include <stdint.h>
#include <ctime>
#include "pcg_random.hpp"
#include "evaluator.h"

using namespace std;

namespace acmaes {

static ivec inverse(const ivec &indices) {
    ivec inverse = ivec(indices.size());
    for (int i = 0; i < indices.size(); i++)
        inverse(indices(i)) = i;
    return inverse;
}

static vec sequence(double start, double end, double step) {
    int size = (int) ((end - start) / step + 1);
    vec d(size);
    double value = start;
    for (int r = 0; r < size; r++) {
        d(r) = value;
        value += step;
    }
    return d;
}

class AcmaesOptimizer {

public:

    AcmaesOptimizer(long runid_, Fitness *fitfun_, int popsize_, int mu_,
            const vec &guess_, const vec &inputSigma_, int maxEvaluations_,
            double accuracy_, double stopfitness_,
            int update_gap_, long seed) {
        // runid used for debugging / logging
        runid = runid_;
        // fitness function to minimize
        fitfun = fitfun_;
        // initial guess for the arguments of the fitness function
        guess = guess_;
        // accuracy = 1.0 is default, > 1.0 reduces accuracy
        accuracy = accuracy_;
        // number of objective variables/problem dimension
        dim = guess_.size();
        // population size, offspring number. The primary strategy parameter to play
        // with, which can be increased from its default value. Increasing the
        // population size improves global search properties in exchange to speed.
        // Speed decreases, as a rule, at most linearly with increasing population
        // size. It is advisable to begin with the default small population size.
        if (popsize_ > 0)
            popsize = popsize_;
        else
            popsize = 4 + int(3. * log(dim));
        // individual sigma values - initial search volume. inputSigma determines
        // the initial coordinate wise standard deviations for the search. Setting
        // SIGMA one third of the initial search region is appropriate.
        if (inputSigma_.size() == 1)
            inputSigma = vec::Constant(dim, inputSigma_[0]);
        else
            inputSigma = inputSigma_;
        // overall standard deviation - search volume.
        sigma = inputSigma.maxCoeff();
        // termination criteria
        // maximal number of evaluations allowed.
        maxEvaluations = maxEvaluations_;
        // limit for fitness value.
        stopfitness = stopfitness_;
        // stop if x-changes larger stopTolUpX.
        stopTolUpX = 1e3 * sigma;
        // stop if x-change smaller stopTolX.
        stopTolX = 1e-11 * sigma * accuracy;
        // stop if fun-changes smaller stopTolFun.
        stopTolFun = 1e-12 * accuracy;
        // stop if back fun-changes smaller stopTolHistFun.
        stopTolHistFun = 1e-13 * accuracy;
        // selection strategy parameters
        // number of parents/points for recombination.
        mu = mu_ > 0 ? mu_ : popsize / 2;
        // array for weighted recombination.
        weights = (log(sequence(1, mu, 1).array()) * -1.) + log(mu + 0.5);
        double sumw = weights.sum();
        double sumwq = weights.squaredNorm();
        weights *= 1. / sumw;
        // variance-effectiveness of sum w_i x_i.
        mueff = sumw * sumw / sumwq;

        // dynamic strategy parameters and constants
        // cumulation constant.
        cc = (4. + mueff / dim) / (dim + 4. + 2. * mueff / dim);
        // cumulation constant for step-size.
        cs = (mueff + 2.) / (dim + mueff + 3.);
        // damping for step-size.
        damps = (1. + 2. * std::max(0., sqrt((mueff - 1.) / (dim + 1.)) - 1.))
                        * max(0.3,
                                1. - // modification for short runs
                                dim / (1e-6 + (maxEvaluations/popsize)))
                                + cs; // minor increment
        // learning rate for rank-one update.
        ccov1 = 2. / ((dim + 1.3) * (dim + 1.3) + mueff);
        // learning rate for rank-mu update'
        ccovmu = min(1. - ccov1,
                2. * (mueff - 2. + 1. / mueff)
                / ((dim + 2.) * (dim + 2.) + mueff));
        // expectation of ||N(0,I)|| == norm(randn(N,1)).
        chiN = sqrt(dim) * (1. - 1. / (4. * dim) + 1 / (21. * dim * dim));
        ccov1Sep = min(1., ccov1 * (dim + 1.5) / 3.);
        ccovmuSep = min(1. - ccov1, ccovmu * (dim + 1.5) / 3.);
        // lazy covariance update gap
        lazy_update_gap =
                update_gap_ >= 0 ?
                        update_gap_ :
                        1.0 / (ccov1 + ccovmu + 1e-23) / dim / 10.0;
        // CMA internal values - updated each generation
        // objective variables.
        xmean = fitfun->encode(guess);
        // evolution path.
        pc = zeros(dim);
        // evolution path for sigma.
        ps = zeros(dim);
        // norm of ps, stored for efficiency.
        normps = ps.norm();
        // coordinate system.
        B = Eigen::MatrixXd::Identity(dim, dim);
        // diagonal of sqrt(D), stored for efficiency.
        diagD = inputSigma / sigma;
        diagC = diagD.cwiseProduct(diagD);
        // B*D, stored for efficiency.
        BD = B.cwiseProduct(diagD.transpose().replicate(dim, 1));
        // covariance matrix.
        C = B * (Eigen::MatrixXd::Identity(dim, dim) * B.transpose());
        // number of iterations.
        iterations = 1;
        // size of history queue of best values.
        historySize = 10 + int(3. * 10. * dim / popsize);
        // stop criteria
        stop = 0;
        // best value so far
        bestValue = DBL_MAX;
        // best parameters so far
        bestX = guess;
        // history queue of best values.
        fitnessHistory = vec::Constant(historySize, DBL_MAX);
        fitnessHistory(0) = bestValue;
        rs = new pcg64(seed);
    }

    ~AcmaesOptimizer() {
        delete rs;
    }

    // param zmean weighted row matrix of the gaussian random numbers generating the current offspring
    // param xold xmean matrix of the previous generation
    // return hsig flag indicating a small correction

    bool updateEvolutionPaths(const vec &zmean, const vec &xold) {
        ps = ps * (1. - cs) + ((B * zmean) * sqrt(cs * (2. - cs) * mueff));
        normps = ps.norm();
        bool hsig = normps / sqrt(1. - pow(1. - cs, 2. * iterations)) / chiN
                < 1.4 + 2. / (dim + 1.);
        pc *= (1. - cc);
        if (hsig)
            pc += (xmean - xold) * (sqrt(cc * (2. - cc) * mueff) / sigma);
        return hsig;
    }

    // param hsig flag indicating a small correction
    // param bestArx fitness-sorted matrix of the argument vectors producing the current offspring
    // param arz unsorted matrix containing the gaussian random values of the current offspring
    // param arindex indices indicating the fitness-order of the current offspring
    // param xold xmean matrix of the previous generation

    double updateCovariance(bool hsig, const mat &bestArx, const mat &arz,
            const ivec &arindex, const mat &xold) {
        double negccov = 0;
        if (ccov1 + ccovmu > 0) {
            mat arpos = (bestArx - xold.replicate(1, mu)) * (1. / sigma); // mu difference vectors
            mat roneu = pc * pc.transpose() * ccov1;
            // minor correction if hsig==false
            double oldFac = hsig ? 0 : ccov1 * cc * (2. - cc);
            oldFac += 1. - ccov1 - ccovmu;
            // Adapt covariance matrix C active CMA
            negccov = (1. - ccovmu) * 0.25 * mueff
                    / (pow(dim + 2., 1.5) + 2. * mueff);
            double negminresidualvariance = 0.66;
            // keep at least 0.66 in all directions, small popsize are most critical
            double negalphaold = 0.5; // where to make up for the variance loss,
            // prepare vectors, compute negative updating matrix Cneg
            ivec arReverseIndex = arindex.reverse();
            mat arzneg = arz(Eigen::all, arReverseIndex.head(mu));
            vec arnorms = arzneg.colwise().norm();
            ivec idxnorms = sort_index(arnorms);
            vec arnormsSorted = arnorms(idxnorms);
            ivec idxReverse = idxnorms.reverse();
            vec arnormsReverse = arnorms(idxReverse);
            arnorms = arnormsReverse.cwiseQuotient(arnormsSorted);
            vec arnormsInv = arnorms(inverse(idxnorms));
            mat sqarnw = arnormsInv.cwiseProduct(arnormsInv).transpose()
                            * weights;
            double negcovMax = (1. - negminresidualvariance) / sqarnw(0);
            if (negccov > negcovMax)
                negccov = negcovMax;
            arzneg = arzneg.cwiseProduct(
                    arnormsInv.transpose().replicate(dim, 1));
            mat artmp = BD * arzneg;
            mat Cneg = artmp * weights.asDiagonal() * artmp.transpose();
            oldFac += negalphaold * negccov;
            C = (C * oldFac) + roneu
                    + (arpos * (ccovmu + (1. - negalphaold) * negccov)
                            * weights.replicate(1, dim).cwiseProduct(
                                    arpos.transpose())) - (Cneg * negccov);
        }
        return negccov;
    }

    // Update B and diagD from C
    // param negccov Negative covariance factor.

    void updateBD(double negccov) {

        if (ccov1 + ccovmu + negccov > 0
                && (std::fmod(iterations,
                        1. / (ccov1 + ccovmu + negccov) / dim / 10.)) < 1.) {
            // to achieve O(N^2) enforce symmetry to prevent complex numbers
            mat triC = C.triangularView<Eigen::Upper>();
            mat triC1 = C.triangularView<Eigen::StrictlyUpper>();
            C = triC + triC1.transpose();
            Eigen::SelfAdjointEigenSolver<mat> sades;
            sades.compute(C);
            // diagD defines the scaling
            diagD = sades.eigenvalues();
            B = sades.eigenvectors();
            if (diagD.minCoeff() <= 0) {
                for (int i = 0; i < dim; i++)
                    if (diagD(i, 0) < 0)
                        diagD(i, 0) = 0.;
                double tfac = diagD.maxCoeff() / 1e14;
                C += Eigen::MatrixXd::Identity(dim, dim) * tfac;
                diagD += vec::Constant(dim, 1.0) * tfac;
            }
            if (diagD.maxCoeff() > 1e14 * diagD.minCoeff()) {
                double tfac = diagD.maxCoeff() / 1e14 - diagD.minCoeff();
                C += Eigen::MatrixXd::Identity(dim, dim) * tfac;
                diagD += vec::Constant(dim, 1.0) * tfac;
            }
            diagC = C.diagonal();
            diagD = diagD.cwiseSqrt(); // D contains standard deviations now
            BD = B.cwiseProduct(diagD.transpose().replicate(dim, 1));
        }
    }

    void newArgs() {
        // generate popsize offspring.
        xmean = fitfun->getClosestFeasible(xmean);
        arz = normal(dim, popsize, *rs);
        arx = mat(dim, popsize);
        for (int k = 0; k < popsize; k++) {
            vec delta = (BD * arz.col(k)) * sigma;
            arx.col(k) = fitfun->getClosestFeasible(xmean + delta);
        }
        fitness = vec(popsize);
    }

    vec ask() {
        // ask for one new argument vector.
        vec arz1 = normalVec(dim, *rs);
        vec delta = (BD * arz1) * sigma;
        vec arx1 = fitfun->getClosestFeasible(xmean + delta);
        return fitfun->decode(arx1);
    }

    int tell(double y, const vec &x) {
        //tell function value for a argument list retrieved by ask_one().
        if (told == 0) {
            fitness = vec(popsize);
            arx = mat(dim, popsize);
            arz = mat(dim, popsize);
        }
        fitness[told] = isfinite(y) ? y : DBL_MAX;
        arx.col(told) = fitfun->encode(x);
        told++;

        if (told >= popsize) {
            xmean = fitfun->getClosestFeasible(xmean);
            try {
                arz = (BD.inverse()
                        * ((arx - xmean.replicate(1, popsize)) / sigma));
            } catch (std::exception &e) {
                arz = normal(dim, popsize, *rs);
            }
            updateCMA();
            told = 0;
            iterations += 1;
        }
        return stop;
    }

    void updateCMA() {
        // sort by fitness and compute weighted mean into xmean
        ivec arindex = sort_index(fitness);
        // calculate new xmean, this is selection and recombination
        vec xold = xmean; // for speed up of Eq. (2) and (3)
        ivec bestIndex = arindex.head(mu);
        mat bestArx = arx(Eigen::all, bestIndex);
        xmean = bestArx * weights;
        mat bestArz = arz(Eigen::all, bestIndex);
        mat zmean = bestArz * weights;
        bool hsig = updateEvolutionPaths(zmean, xold);
        // adapt step size sigma
        sigma *= exp(min(1.0, (normps / chiN - 1.) * cs / damps));
        double bestFitness = fitness(arindex(0));
        double worstFitness = fitness(arindex(arindex.size() - 1));
        if (bestValue > bestFitness) {
            bestValue = bestFitness;
            bestX = fitfun->decode(bestArx.col(0));
            if (isfinite(stopfitness) && bestFitness < stopfitness) {
                stop = 1;
                return;
            }
        }
        if (iterations >= last_update + lazy_update_gap) {
            last_update = iterations;
            double negccov = updateCovariance(hsig, bestArx, arz, arindex,
                    xold);
            updateBD(negccov);
            // handle termination criteria
            vec sqrtDiagC = diagC.cwiseSqrt();
            vec pcCol = pc;
            for (int i = 0; i < dim; i++) {
                if (sigma * (max(abs(pcCol[i]), sqrtDiagC[i])) > stopTolX)
                    break;
                if (i >= dim - 1)
                    stop = 2;
            }
            if (stop > 0)
                return;
            for (int i = 0; i < dim; i++)
                if (sigma * sqrtDiagC[i] > stopTolUpX)
                    stop = 3;
            if (stop > 0)
                return;
        }
        double historyBest = fitnessHistory.minCoeff();
        double historyWorst = fitnessHistory.maxCoeff();
        if (iterations > 2
                && max(historyWorst, worstFitness)
        - min(historyBest, bestFitness) < stopTolFun) {
            stop = 4;
            return;
        }
        if (iterations > fitnessHistory.size()
                && historyWorst - historyBest < stopTolHistFun) {
            stop = 5;
            return;
        }
        // condition number of the covariance matrix exceeds 1e14
        if (diagD.maxCoeff() / diagD.minCoeff() > 1e7 * 1.0 / sqrt(accuracy)) {
            stop = 6;
            return;
        }
        // adjust step size in case of equal function values (flat fitness)
        if (bestValue == fitness[arindex[(int) (0.1 + popsize / 4.)]]) {
            sigma *= exp(0.2 + cs / damps);
        }
        if (iterations > 2
                && max(historyWorst, bestFitness)
        - std::min(historyBest, bestFitness) == 0) {
            sigma *= ::exp(0.2 + cs / damps);
        }
        // store best in history
        for (int i = 1; i < fitnessHistory.size(); i++)
            fitnessHistory[i] = fitnessHistory[i - 1];
        fitnessHistory[0] = bestFitness;
    }

    void doOptimize() {

        // -------------------- Generation Loop --------------------------------
        for (iterations = 1; fitfun->evaluations() < maxEvaluations && !fitfun->terminate();
                iterations++) {
            // generate and evaluate popsize offspring
            newArgs();
            for (int k = 0; k < popsize; k++) {
                fitness[k] = fitfun->eval(fitfun->decode(arx.col(k)))(0);
                if (!isfinite(fitness[k]))
                    fitness[k] = DBL_MAX;
            }
            updateCMA();
            if (stop != 0)
                return;
        }
    }

    void do_optimize_delayed_update(int workers) {
        iterations = 0;
        fitfun->resetEvaluations();
        evaluator eval(fitfun, 1, workers);
        vec evals_x[workers];
        // fill eval queue with initial population
        for (int i = 0; i < workers; i++) {
            vec x = ask();
            eval.evaluate(x, i);
            evals_x[i] = x;
        }
        while (fitfun->evaluations() < maxEvaluations) {
            vec_id* vid = eval.result();
            vec y = vec(vid->_v);
            int p = vid->_id;
            delete vid;
            vec x = evals_x[p];
            tell(y(0), x); // tell evaluated x
            if (fitfun->evaluations() >= maxEvaluations)
                break;
            x = ask();
            eval.evaluate(x, p);
            evals_x[p] = x;
        }
    }

    vec getBestX() {
        return bestX;
    }

    double getBestValue() {
        return bestValue;
    }

    double getIterations() {
        return iterations;
    }

    int getStop() {
        return stop;
    }

    Fitness* getFitfun() {
        return fitfun;
    }

    int getDim() {
        return dim;
    }

private:
    long runid;
    Fitness *fitfun;
    vec guess;
    double accuracy;
    int popsize; // population size
    vec inputSigma;
    int dim;
    int maxEvaluations;
    double stopfitness;
    double stopTolUpX;
    double stopTolX;
    double stopTolFun;
    double stopTolHistFun;
    int mu; //
    vec weights;
    double mueff; //
    double sigma;
    double cc;
    double cs;
    double damps;
    double ccov1;
    double ccovmu;
    double chiN;
    double ccov1Sep;
    double ccovmuSep;
    double lazy_update_gap = 0;
    vec xmean;
    vec pc;
    vec ps;
    double normps;
    mat B;
    mat BD;
    mat diagD;
    mat C;
    vec diagC;
    mat arz;
    mat arx;
    vec fitness;
    int iterations = 0;
    int last_update = 0;
    vec fitnessHistory;
    int historySize;
    double bestValue;
    vec bestX;
    int stop;
    int told = 0;
    pcg64 *rs;
};
}

using namespace acmaes;

extern "C" {
void optimizeACMA_C(long runid, callback_type func, int dim,
        double *init, double *lower, double *upper, double *sigma,
        int maxEvals, double stopfitness, int mu, int popsize, double accuracy,
        long seed, bool normalize, int update_gap, int workers, double* res) {
    int n = dim;
    vec guess(n), lower_limit(n), upper_limit(n), inputSigma(n);
    bool useLimit = false;
    for (int i = 0; i < n; i++) {
        guess[i] = init[i];
        inputSigma[i] = sigma[i];
        lower_limit[i] = lower[i];
        upper_limit[i] = upper[i];
        useLimit |= (lower[i] != 0);
        useLimit |= (upper[i] != 0);
    }
    if (useLimit == false) {
        lower_limit.resize(0);
        upper_limit.resize(0);
    }
    Fitness fitfun(func, n, 1, lower_limit, upper_limit);
    fitfun.setNormalize(normalize);
    AcmaesOptimizer opt(runid, &fitfun, popsize, mu, guess, inputSigma,
            maxEvals, accuracy, stopfitness, update_gap, seed);
    try {
        if (workers <= 1)
            opt.doOptimize();
        else
            opt.do_optimize_delayed_update(workers);
        vec bestX = opt.getBestX();
        double bestY = opt.getBestValue();
        for (int i = 0; i < n; i++)
            res[i] = bestX[i];
        res[n] = bestY;
        res[n + 1] = fitfun.evaluations();
        res[n + 2] = opt.getIterations();
        res[n + 3] = opt.getStop();
    } catch (std::exception &e) {
        cout << e.what() << endl;
    }
}
}


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

#ifndef MTL_OPERATIONS_INCLUDE
#define MTL_OPERATIONS_INCLUDE

#include <boost/numeric/mtl/operation/adjoint.hpp>
#include <boost/numeric/mtl/operation/clone.hpp>
#include <boost/numeric/mtl/operation/cholesky.hpp>
#include <boost/numeric/mtl/operation/column_in_matrix.hpp>
#include <boost/numeric/mtl/operation/conj.hpp>
#include <boost/numeric/mtl/operation/copysign.hpp>
#include <boost/numeric/mtl/operation/crop.hpp>
#include <boost/numeric/mtl/operation/cross.hpp>
#include <boost/numeric/mtl/operation/cuppen.hpp>
#include <boost/numeric/mtl/operation/diagonal.hpp>
#include <boost/numeric/mtl/operation/dot.hpp>
#include <boost/numeric/mtl/operation/eigenvalue.hpp>
#include <boost/numeric/mtl/operation/eigenvalue_symmetric.hpp>
#include <boost/numeric/mtl/operation/entry1D.hpp>
#include <boost/numeric/mtl/operation/entry_similar.hpp>
#include <boost/numeric/mtl/operation/evaluate_lazy.hpp>
#include <boost/numeric/mtl/operation/extended_complex.hpp>
#include <boost/numeric/mtl/operation/fill.hpp>
#include <boost/numeric/mtl/operation/fuse.hpp>
#include <boost/numeric/mtl/operation/givens.hpp>
#include <boost/numeric/mtl/operation/hermitian.hpp>
#include <boost/numeric/mtl/operation/hessenberg.hpp>
#include <boost/numeric/mtl/operation/householder.hpp>
#include <boost/numeric/mtl/operation/imag.hpp>
#include <boost/numeric/mtl/operation/inv.hpp>
#include <boost/numeric/mtl/operation/invert_diagonal.hpp>
#include <boost/numeric/mtl/operation/is_negative.hpp>
#include <boost/numeric/mtl/operation/lazy.hpp>
#include <boost/numeric/mtl/operation/left_scale_inplace.hpp>
#include <boost/numeric/mtl/operation/lower_trisolve.hpp>
#include <boost/numeric/mtl/operation/lu.hpp>
#include <boost/numeric/mtl/operation/make_sparse.hpp>
#include <boost/numeric/mtl/operation/make_tag_vector.hpp>
#include <boost/numeric/mtl/operation/merge_complex_vector.hpp>
#include <boost/numeric/mtl/operation/minimal_increase.hpp>
#include <boost/numeric/mtl/operation/misc.hpp>
#include <boost/numeric/mtl/operation/mult.hpp>
#include <boost/numeric/mtl/operation/norms.hpp>
#include <boost/numeric/mtl/operation/ones.hpp>
#include <boost/numeric/mtl/operation/operators.hpp>
#include <boost/numeric/mtl/operation/orth.hpp>
#include <boost/numeric/mtl/operation/print.hpp>
#include <boost/numeric/mtl/operation/product.hpp>
#include <boost/numeric/mtl/operation/qr.hpp>
#include <boost/numeric/mtl/operation/random.hpp>
#include <boost/numeric/mtl/operation/rank_one_update.hpp>
#include <boost/numeric/mtl/operation/rank_two_update.hpp>
#include <boost/numeric/mtl/operation/real.hpp>
#include <boost/numeric/mtl/operation/resource.hpp>
#include <boost/numeric/mtl/operation/right_scale_inplace.hpp>
#include <boost/numeric/mtl/operation/scale.hpp>
#include <boost/numeric/mtl/operation/set_to_zero.hpp>
#include <boost/numeric/mtl/operation/secular.hpp>
#include <boost/numeric/mtl/operation/signum.hpp>
#include <boost/numeric/mtl/operation/split_complex_vector.hpp>
#include <boost/numeric/mtl/operation/sub_matrix.hpp>
#include <boost/numeric/mtl/operation/sum.hpp>
#include <boost/numeric/mtl/operation/min.hpp>
#include <boost/numeric/mtl/operation/min_pos.hpp>
#include <boost/numeric/mtl/operation/max.hpp>
#include <boost/numeric/mtl/operation/max_pos.hpp>
#include <boost/numeric/mtl/operation/max_abs_pos.hpp>
#include <boost/numeric/mtl/operation/num_cols.hpp>
#include <boost/numeric/mtl/operation/num_rows.hpp>
#include <boost/numeric/mtl/operation/row_in_matrix.hpp>
#include <boost/numeric/mtl/operation/size.hpp>
#include <boost/numeric/mtl/operation/size1D.hpp>
#include <boost/numeric/mtl/operation/static_num_cols.hpp>
#include <boost/numeric/mtl/operation/static_num_rows.hpp>
#include <boost/numeric/mtl/operation/static_size.hpp>
#include <boost/numeric/mtl/operation/svd.hpp>
#include <boost/numeric/mtl/operation/swap_row.hpp>
#include <boost/numeric/mtl/operation/trace.hpp>
#include <boost/numeric/mtl/operation/trans.hpp>
#include <boost/numeric/mtl/operation/unary_dot.hpp>
#include <boost/numeric/mtl/operation/unroll.hpp>
#include <boost/numeric/mtl/operation/upper_trisolve.hpp>

#include <boost/numeric/mtl/matrix/bands.hpp>
#include <boost/numeric/mtl/matrix/identity.hpp>
#include <boost/numeric/mtl/matrix/lower.hpp>
#include <boost/numeric/mtl/matrix/permutation.hpp>
#include <boost/numeric/mtl/matrix/reorder.hpp>
#include <boost/numeric/mtl/matrix/reorder_ref.hpp>
#include <boost/numeric/mtl/matrix/reorder_matrix_rows.hpp>
#include <boost/numeric/mtl/matrix/strict_upper.hpp>
#include <boost/numeric/mtl/matrix/strict_lower.hpp>
#include <boost/numeric/mtl/matrix/upper.hpp>

#include <boost/numeric/mtl/io/path.hpp>

#include <boost/numeric/mtl/utility/exception.hpp>
#include <boost/numeric/mtl/utility/string_to_enum.hpp>
#include <boost/numeric/mtl/utility/make_copy_or_reference.hpp>

#include <boost/numeric/mtl/interface/umfpack_solve.hpp>

#endif // MTL_OPERATIONS_INCLUDE


// #include <opencv2/core/core.hpp> // needed for verbosity >= 3, DISVISUAL
// #include <opencv2/highgui/highgui.hpp> // needed for verbosity >= 3, DISVISUAL
// #include <opencv2/imgproc/imgproc.hpp> // needed for verbosity >= 3, DISVISUAL

#include <iostream>
#include <string>
#include <vector>
#include <valarray>

#include <thread>

#include <Eigen/Core>
#include <Eigen/LU>
#include <Eigen/Dense>

#include <stdio.h>

#include "patch.h"
#include "patchgrid.h"

using std::cout;
using std::endl;
using std::vector;

namespace OFC
{

  PatGridClass::PatGridClass(
    const camparam* cpt_in,
    const camparam* cpo_in,
    const optparam* op_in)
  :
    cpt(cpt_in),
    cpo(cpo_in),
    op(op_in)
  {

  // Generate grid on current scale
  steps = op->steps;
  nopw = ceil( (float)cpt->width /  (float)steps );
  noph = ceil( (float)cpt->height / (float)steps );
  const int offsetw = floor((cpt->width - (nopw-1)*steps)/2);
  const int offseth = floor((cpt->height - (noph-1)*steps)/2);

  nopatches = nopw*noph;
  pt_ref.resize(nopatches);
  p_init.resize(nopatches);
  pat.reserve(nopatches);

  im_ao_eg = new Eigen::Map<const Eigen::MatrixXf>(nullptr,cpt->height,cpt->width);
  im_ao_dx_eg = new Eigen::Map<const Eigen::MatrixXf>(nullptr,cpt->height,cpt->width);
  im_ao_dy_eg = new Eigen::Map<const Eigen::MatrixXf>(nullptr,cpt->height,cpt->width);

  im_bo_eg = new Eigen::Map<const Eigen::MatrixXf>(nullptr,cpt->height,cpt->width);
  im_bo_dx_eg = new Eigen::Map<const Eigen::MatrixXf>(nullptr,cpt->height,cpt->width);
  im_bo_dy_eg = new Eigen::Map<const Eigen::MatrixXf>(nullptr,cpt->height,cpt->width);

  int patchid=0;
  for (int x = 0; x < nopw; ++x)
  {
    for (int y = 0; y < noph; ++y)
    {
      int i = x*noph + y;

      pt_ref[i][0] = x * steps + offsetw;
      pt_ref[i][1] = y * steps + offseth;
      p_init[i].setZero();

      pat.push_back(new OFC::PatClass(cpt, cpo, op, patchid));
      patchid++;
    }
  }
}

PatGridClass::~PatGridClass()
{
  delete im_ao_eg;
  delete im_ao_dx_eg;
  delete im_ao_dy_eg;

  delete im_bo_eg;
  delete im_bo_dx_eg;
  delete im_bo_dy_eg;

  for (int i=0; i< nopatches; ++i)
    delete pat[i];
}

void PatGridClass::SetComplGrid(PatGridClass *cg_in)
{
  cg = cg_in;
}

void PatGridClass::InitializeGrid(const float * im_ao_in, const float * im_ao_dx_in, const float * im_ao_dy_in)
{
  im_ao = im_ao_in;
  im_ao_dx = im_ao_dx_in;
  im_ao_dy = im_ao_dy_in;

  new (im_ao_eg) Eigen::Map<const Eigen::MatrixXf>(im_ao,cpt->height,cpt->width); // new placement operator
  new (im_ao_dx_eg) Eigen::Map<const Eigen::MatrixXf>(im_ao_dx,cpt->height,cpt->width);
  new (im_ao_dy_eg) Eigen::Map<const Eigen::MatrixXf>(im_ao_dy,cpt->height,cpt->width);

  #pragma omp parallel for schedule(static)
  for (int i = 0; i < nopatches; ++i)
  {
    pat[i]->InitializePatch(im_ao_eg, im_ao_dx_eg, im_ao_dy_eg, pt_ref[i]);
    p_init[i].setZero();
  }

}

void PatGridClass::SetTargetImage(const float * im_bo_in, const float * im_bo_dx_in, const float * im_bo_dy_in)
{
  im_bo = im_bo_in;
  im_bo_dx = im_bo_dx_in;
  im_bo_dy = im_bo_dy_in;

  new (im_bo_eg) Eigen::Map<const Eigen::MatrixXf>(im_bo,cpt->height,cpt->width); // new placement operator
  new (im_bo_dx_eg) Eigen::Map<const Eigen::MatrixXf>(im_bo_dx,cpt->height,cpt->width); // new placement operator
  new (im_bo_dy_eg) Eigen::Map<const Eigen::MatrixXf>(im_bo_dy,cpt->height,cpt->width); // new placement operator

  #pragma omp parallel for schedule(static)
  for (int i = 0; i < nopatches; ++i)
    pat[i]->SetTargetImage(im_bo_eg, im_bo_dx_eg, im_bo_dy_eg);

}

void PatGridClass::Optimize()
{
    #pragma omp parallel for schedule(dynamic,10)
    for (int i = 0; i < nopatches; ++i)
    {
      pat[i]->OptimizeIter(p_init[i], true); // optimize until convergence
    }
}

// void PatGridClass::OptimizeAndVisualize(const float sc_fct_tmp) // needed for verbosity >= 3, DISVISUAL
// {
//   bool allconverged=0;
//   int cnt = 0;
//   while (!allconverged)
//   {
//     cnt++;
//
//     allconverged=1;
//
//     for (int i = 0; i < nopatches; ++i)
//     {
//       if (pat[i]->isConverged()==0)
//       {
//         pat[i]->OptimizeIter(p_init[i], false); // optimize, only one iterations
//         allconverged=0;
//       }
//     }
//
//
//     // Display original image
//     const cv::Mat src(cpt->height+2*cpt->imgpadding, cpt->width+2*cpt->imgpadding, CV_32FC1, (void*) im_ao);
//     cv::Mat img_ao_mat = src(cv::Rect(cpt->imgpadding,cpt->imgpadding,cpt->width,cpt->height));
//     cv::Mat outimg;
//     img_ao_mat.convertTo(outimg, CV_8UC1);
//     cv::cvtColor(outimg, outimg, CV_GRAY2RGB);
//     cv::resize(outimg, outimg, cv::Size(), sc_fct_tmp, sc_fct_tmp, cv::INTER_NEAREST);
//
//     for (int i = 0; i < nopatches; ++i)
//     {
//       // Show displacement vector
//       const Eigen::Vector2f pt_ret = pat[i]->GetPointPos();
//
//       Eigen::Vector2f pta, ptb;
//
//       cv::line(outimg, cv::Point( (pt_ref[i][0]+.5)*sc_fct_tmp, (pt_ref[i][1]+.5)*sc_fct_tmp ), cv::Point( (pt_ret[0]+.5)*sc_fct_tmp, (pt_ret[1]+.5)*sc_fct_tmp ), cv::Scalar(255*pat[i]->isConverged() ,255*(!pat[i]->isConverged()),0),  2);
//
//       cv::line(outimg, cv::Point( (cpt->cx+.5)*sc_fct_tmp, (cpt->cy+.5)*sc_fct_tmp ), cv::Point( (cpt->cx+.5)*sc_fct_tmp, (cpt->cy+.5)*sc_fct_tmp ), cv::Scalar(0,0, 255),  2);
//
//     }
//
//     char str[200];
//     sprintf(str,"Iter: %i",cnt);
//     cv::putText(outimg, str, cv::Point2f(20,20), cv::FONT_HERSHEY_PLAIN, 1,  cv::Scalar(0,0,255,255), 2);
//
//     cv::namedWindow( "Img_iter", cv::WINDOW_AUTOSIZE );
//     cv::imshow( "Img_iter", outimg);
//
//     cv::waitKey(500);
//   }
// }

void PatGridClass::InitializeFromCoarserOF(const float * flow_prev)
{
  #pragma omp parallel for schedule(dynamic,10)
  for (int ip = 0; ip < nopatches; ++ip)
  {
    int x = floor(pt_ref[ip][0] / 2); // better, but slower: use bil. interpolation here
    int y = floor(pt_ref[ip][1] / 2);
    int i = y*(cpt->width/2) + x;

    #if (SELECTMODE==1)
    p_init[ip](0) = flow_prev[2*i  ]*2;
    p_init[ip](1) = flow_prev[2*i+1]*2;
    #else
    p_init[ip](0) = flow_prev[  i  ]*2;
    #endif
  }
}

void PatGridClass::AggregateFlowDense(float *flowout) const
{
  float* we = new float[cpt->width * cpt->height];

  memset(flowout, 0, sizeof(float) * (op->nop * cpt->width * cpt->height) );
  memset(we,      0, sizeof(float) * (          cpt->width * cpt->height) );

  #ifdef USE_PARALLEL_ON_FLOWAGGR // Using this enables OpenMP on flow aggregation. This can lead to race conditions. Experimentally we found that the result degrades only marginally. However, for our experiments we did not enable this.
    #pragma omp parallel for schedule(static)
  #endif
  for (int ip = 0; ip < nopatches; ++ip)
  {

    if (pat[ip]->IsValid())
    {
      #if (SELECTMODE==1)
      const Eigen::Vector2f*            fl = pat[ip]->GetParam(); // flow displacement of this patch
      Eigen::Vector2f flnew;
      #else
      const Eigen::Matrix<float, 1, 1>* fl = pat[ip]->GetParam(); // horz. displacement of this patch
      Eigen::Matrix<float, 1, 1> flnew;
      #endif

      const float * pweight = pat[ip]->GetpWeightPtr(); // use image error as weight

      int lb = -op->p_samp_s/2;
      int ub = op->p_samp_s/2-1;

      for (int y = lb; y <= ub; ++y)
      {
        for (int x = lb; x <= ub; ++x, ++pweight)
        {
          int yt = (y + pt_ref[ip][1]);
          int xt = (x + pt_ref[ip][0]);

          if (xt >= 0 && yt >= 0 && xt < cpt->width && yt < cpt->height)
          {

            int i = yt*cpt->width + xt;

            #if (SELECTCHANNEL==1 | SELECTCHANNEL==2)  // single channel/gradient image
            float absw = 1.0f /  (float)(std::max(op->minerrval  ,*pweight));
            #else  // RGB image
            float absw = (float)(std::max(op->minerrval  ,*pweight)); ++pweight;
                  absw+= (float)(std::max(op->minerrval  ,*pweight)); ++pweight;
                  absw+= (float)(std::max(op->minerrval  ,*pweight));
            absw = 1.0f / absw;
            #endif

            flnew = (*fl) * absw;
            we[i] += absw;

            #if (SELECTMODE==1)
            flowout[2*i]   += flnew[0];
            flowout[2*i+1] += flnew[1];
            #else
            flowout[i] += flnew[0];
            #endif
          }
        }
      }
    }
  }

  // if complementary (forward-backward merging) is given, integrate negative backward flow as well
  if (cg)
  {
      Eigen::Vector4f wbil; // bilinear weight vector
      Eigen::Vector4i pos;

      #ifdef USE_PARALLEL_ON_FLOWAGGR
        #pragma omp parallel for schedule(static)
      #endif
      for (int ip = 0; ip < cg->nopatches; ++ip)
      {
        if (cg->pat[ip]->IsValid())
        {
          #if (SELECTMODE==1)
          const Eigen::Vector2f*            fl = (cg->pat[ip]->GetParam()); // flow displacement of this patch
          Eigen::Vector2f flnew;
          #else
          const Eigen::Matrix<float, 1, 1>* fl = (cg->pat[ip]->GetParam()); // horz. displacement of this patch
          Eigen::Matrix<float, 1, 1> flnew;
          #endif

          const Eigen::Vector2f rppos = cg->pat[ip]->GetPointPos(); // get patch position after optimization
          const float * pweight = cg->pat[ip]->GetpWeightPtr(); // use image error as weight

          Eigen::Vector2f resid;

          // compute bilinear weight vector
          pos[0] = ceil(rppos[0] +.00001); // make sure they are rounded up to natural number
          pos[1] = ceil(rppos[1] +.00001); // make sure they are rounded up to natural number
          pos[2] = floor(rppos[0]);
          pos[3] = floor(rppos[1]);

          resid[0] = rppos[0] - pos[2];
          resid[1] = rppos[1] - pos[3];
          wbil[0] = resid[0]*resid[1];
          wbil[1] = (1-resid[0])*resid[1];
          wbil[2] = resid[0]*(1-resid[1]);
          wbil[3] = (1-resid[0])*(1-resid[1]);

          int lb = -op->p_samp_s/2;
          int ub = op->p_samp_s/2-1;

          for (int y = lb; y <= ub; ++y)
          {
            for (int x = lb; x <= ub; ++x, ++pweight)
            {

              int yt = y + pos[1];
              int xt = x + pos[0];
              if (xt >= 1 && yt >= 1 && xt < (cpt->width-1) && yt < (cpt->height-1))
              {

                #if (SELECTCHANNEL==1 | SELECTCHANNEL==2)  // single channel/gradient image
                float absw = 1.0f /  (float)(std::max(op->minerrval  ,*pweight));
                #else  // RGB
                float absw = (float)(std::max(op->minerrval  ,*pweight)); ++pweight;
                      absw+= (float)(std::max(op->minerrval  ,*pweight)); ++pweight;
                      absw+= (float)(std::max(op->minerrval  ,*pweight));
                absw = 1.0f / absw;
                #endif

                flnew = (*fl) * absw;

                int idxcc =  xt    +  yt   *cpt->width;
                int idxfc = (xt-1) +  yt   *cpt->width;
                int idxcf =  xt    + (yt-1)*cpt->width;
                int idxff = (xt-1) + (yt-1)*cpt->width;

                we[idxcc] += wbil[0] * absw;
                we[idxfc] += wbil[1] * absw;
                we[idxcf] += wbil[2] * absw;
                we[idxff] += wbil[3] * absw;

                #if (SELECTMODE==1)
                flowout[2*idxcc  ] -= wbil[0] * flnew[0];   // use reversed flow
                flowout[2*idxcc+1] -= wbil[0] * flnew[1];

                flowout[2*idxfc  ] -= wbil[1] * flnew[0];
                flowout[2*idxfc+1] -= wbil[1] * flnew[1];

                flowout[2*idxcf  ] -= wbil[2] * flnew[0];
                flowout[2*idxcf+1] -= wbil[2] * flnew[1];

                flowout[2*idxff  ] -= wbil[3] * flnew[0];
                flowout[2*idxff+1] -= wbil[3] * flnew[1];
                #else
                flowout[idxcc] -= wbil[0] * flnew[0]; // simple averaging of inverse horizontal displacement
                flowout[idxfc] -= wbil[1] * flnew[0];
                flowout[idxcf] -= wbil[2] * flnew[0];
                flowout[idxff] -= wbil[3] * flnew[0];
                #endif
              }
            }
          }
        }
      }
  }

  #pragma omp parallel for schedule(static, 100)
  // normalize each pixel by dividing displacement by aggregated weights from all patches
  for (int yi = 0; yi < cpt->height; ++yi)
  {
    for (int xi = 0; xi < cpt->width; ++xi)
    {
      int i    = yi*cpt->width + xi;
      if (we[i]>0)
      {
        #if (SELECTMODE==1)
        flowout[2*i  ] /= we[i];
        flowout[2*i+1] /= we[i];
        #else
        flowout[i] /= we[i];
        #endif
      }
    }
  }

  delete[] we;
}

}


/*%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%  Solving 2D Poisson + Drift Diffusion semiconductor eqns for a solar cell using
%                      Scharfetter-Gummel discretization
%
%                         Written by Timofey Golubev
%
%     This includes the 2D poisson equation and 2D continuity/drift-diffusion
%     equations using Scharfetter-Gummel discretization. The Poisson equation
%     is solved first, and the solution of potential is used to calculate the
%     Bernoulli functions and solve the continuity eqn's.
%
%   Boundary conditions for Poisson equation are:
%
%     -a fixed voltage at (x,0) and (x, Nz) defined by V_bottomBC
%      and V_topBC which are defining the  electrodes
%
%    -insulating boundary conditions: V(0,z) = V(1,z) and
%     V(0,N+1) = V(1,N) (N is the last INTERIOR mesh point).
%     so the potential at the boundary is assumed to be the same as just inside
%     the boundary. Gradient of potential normal to these boundaries is 0.
%
%   Matrix equations are AV*V = bV, Ap*p = bp, and An*n = bn where AV, Ap, and An are sparse matrices
%   (generated using spdiag), for the Poisson and continuity equations.
%   V is the solution for electric potential, p is the solution for hole
%   density, n is solution for electron density
%   bV is the rhs of Poisson eqn which contains the charge densities and boundary conditions
%   bp is the rhs of hole continuity eqn which contains net generation rate
%   and BCs
%
%     The code as is will calculate data for a JV curve
%     as well as carrier densities, current densities, and electric field
%     distributions of a generic solar cell made of an active layer and electrodes.
%     More equations for carrier recombination can be easily added.f
%
%     Photogeneration rate will be inputed from gen_rate.inp file
%     (i.e. the output of an optical model can be used) or an analytic expression
%     for photogeneration rate can be added to photogeneration.cpp. Generation rate file
%     should contain num_cell-2 number of entries in a single column, corresponding to
%     the the generation rate at each mesh point (except the endpoints).
%
%     The code can also be applied to non-illuminated devices by
%     setting photogeneration rate to 0.
%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%*/

#include <iostream>
#include <vector>
#include <iomanip>
#include <algorithm>   //allows to use fill and min
#include <fstream>
#include <chrono>
#include <string>
#include <time.h>
#include <fstream>
#include <string>

#include <omp.h>

#define EIGEN_USE_MKL_ALL  //is for Intel MKL

//#define EIGEN_NO_DEBUG   //this should turn off eigen asserts, //MAKES NO PERFORMANCE DIFFERENCE!, is I think auto turned off in release mode
#include <Eigen/Sparse>
#include <Eigen/Dense>
#include<Eigen/IterativeLinearSolvers>
#include <Eigen/SparseCholesky>
#include<Eigen/SparseQR>
#include <Eigen/OrderingMethods>
#include<Eigen/SparseLU>
#include <unsupported/Eigen/CXX11/Tensor>  //allows for 3D matrices (Tensors)

#include "constants.h"        //these contain physics constants only
#include "parameters.h"
#include "poisson.h"
#include "continuity_p.h"
//#include "continuity_n.h"
//#include "recombination.h"
//#include "photogeneration.h"
#include "Utilities.h"

#include "mkl.h"

int main()
{

    //trivial MKL function call for testing
    vcAbs(0, 0, 0);  //MY MKL linking works!!, b/c otherwise it wouldn't recognize this function!

    std::chrono::high_resolution_clock::time_point start = std::chrono::high_resolution_clock::now();  //start clock timer
    Parameters params;    //params is struct storing all parameters
    params.Initialize();  //reads parameters from file

    const int num_cell_x = params.num_cell_x;   //create a local num_cell so don't have to type params.num_cell everywhere
    const int num_cell_y = params.num_cell_y;
    const int num_cell_z = params.num_cell_z;

    const int num_V = static_cast<int>(floor((params.Va_max-params.Va_min)/params.increment))+1;  //floor returns double, explicitely cast to int
    params.tolerance_eq = params.tolerance_i;
    const int Nx = params.Nx;
    const int Ny = params.Ny;
    const int Nz = params.Nz;
    const int num_rows = (Nx+1)*(Ny+1)*(Nz+1);  //number of rows in the solution vectors (V, n, p)
    //NOTE: num_rows is the same as num_elements.
    //NOTE: we include the top BC inside the matrix and solution vectors to allow in future to use mixed BC's there.

    std::ofstream JV;
    JV.open("JV.txt");  //note: file will be created inside the build directory

    //-------------------------------------------------------------------------------------------------------
    //Initialize other vectors
    //WILL INDEX FROM 0, b/c that's what Eigen library does.
    //Note: these are Tensor type, so can use reshape on them, but these all are just a single column (i.e. the soln column from the matrix eqn).
    Eigen::Tensor<double, 3> p(num_rows, 1, 1), oldp(num_rows, 1, 1), newp(num_rows, 1, 1);
    Eigen::Tensor<double, 3> oldV(num_rows, 1, 1), newV(num_rows, 1, 1), V(num_rows, 1, 1);

    //create matrices to hold the V, n, and p values (including those at the boundaries) according to the (x,z) coordinates.
    //allows to write formulas in terms of coordinates
    Eigen::VectorXd soln_V(num_rows), soln_p(num_rows);  //vector for storing solutions to the  sparse solver (indexed from 0)
    Eigen::VectorXd p_Xd(num_rows), V_Xd(num_rows);  //Eigen Vector_Xds for the initial buesses to bicgstab method
    //For the following, only need gen rate on insides, so N+1 size is enough
    std::vector<double> Up(num_rows); //will store generation rate as vector, for easy use in rhs

    //Eigen::Tensor<double, 3> R_Langevin(N+1,N+1,N+1), PhotogenRate(N+1,N+1,N+1);
    Eigen::Tensor<double, 3> J_total_Z(num_cell_x+1, num_cell_y+1, num_cell_z+1), J_total_X(num_cell_x+1, num_cell_y+1, num_cell_z+1), J_total_Y(num_cell_x+1, num_cell_y+1, num_cell_z+1);  //we want the indices of J to correspond to the real x,y,z values..., for convinience                //matrices for spacially dependent current
    Eigen::Tensor<double, 3> V_matrix(num_cell_x, num_cell_y, num_cell_z), temp_permuted; //indexed from 0, so just num_cell....//temp_permuted is needed b/c it can't do in place permutations...., need to save to another tensor!  //Note: is actually a Tensor in Eigen
    Eigen::Tensor<double, 3> fullV(num_cell_x+1, num_cell_y+1, num_cell_z+1), fullp(num_cell_x+1, num_cell_y+1, num_cell_z+1);  //for storing all the values in device, including bndrys
    Eigen::SparseMatrix<double> input; //for feeding input matrix into BiCGSTAB, b/c it crashes if try to call get matrix from the solve call.

    //test if openmp is working
    //omp_set_num_threads(8);   //this can allow to set the number of threads that will be used

    std::cout << Eigen::nbThreads() << std::endl;  //displays the # of threads that will be used by Eigen--> mine displays 8, but doesn't seem like it's using 8.

    //this will only make sense if within a region of code which is multithreaded
//    int nthreads = omp_get_num_threads();  //SAYS ONLY 1 THREAD IS AVAILABLE--> MIGHT HAVE AN ISSUE HERE!
//    std::cout << nthreads << " treads available " << std::endl;
//------------------------------------------------------------------------------------
    //Construct objects
    Poisson poisson(params);  //so it can't construct the poisson object
    //Recombo recombo(params);
    Continuity_p continuity_p(params);  //note this also sets up the constant top and bottom electrode BC's

    //Continuity_n continuity_n(params);  //note this also sets up the constant top and bottom electrode BC's
    //Photogeneration photogen(params, params.Photogen_scaling, params.GenRateFileName);
    Utilities utils;
    Eigen::SimplicialLDLT<Eigen::SparseMatrix<double>, Eigen::UpLoType::Lower, Eigen::AMDOrdering<int>> SCholesky; //Note using NaturalOrdering is much much slower

    Eigen::SparseQR<Eigen::SparseMatrix<double>, Eigen::COLAMDOrdering<int>> SQR;
    Eigen::SparseLU<Eigen::SparseMatrix<double> >  poisson_LU, cont_n_LU, cont_p_LU;
    //Eigen::BiCGSTAB<Eigen::SparseMatrix<double>, Eigen::IncompleteLUT<double>> BiCGStab_solver;  //BiCGStab solver object/ /USING THIS PRECONDITIONER IS WAY TOO SLOW FOR LARGE SYSTEMS!
    Eigen::BiCGSTAB<Eigen::SparseMatrix<double>, Eigen::DiagonalPreconditioner<double>> BiCGStab_solver;   //NOTE: WORKS MUCH FASTER WITH DIAGONAL PRECONDITIONER, THAN the IncompleteLUT preconditioner!!--> probably b/c
    //Eigen::BiCGSTAB<Eigen::SparseMatrix<double>, Eigen::IdentityPreconditioner> BiCGStab_solver;  //try with Identity preconditioner, the simplest trivial one
    BiCGStab_solver.setTolerance(1e-14); //set the tolerance explicitely, so matches Matlab's tolerance

    Eigen::ConjugateGradient<Eigen::SparseMatrix<double>, Eigen::UpLoType::Lower|Eigen::UpLoType::Upper > cg;

//--------------------------------------------------------------------------------------------
    //Define boundary conditions and initial conditions. Note: electrodes are at the top and bottom.
    double Va = params.Va_min;
    poisson.set_V_bottomBC(params, Va);
    poisson.set_V_topBC(params, Va);  //THESE CAN BE MOVED TO BE DONE WITHIN constructor of Poisson object

    //Initial conditions
    //for now assume diff is constant everywhere...
    double diff = (poisson.get_V_topBC(0,0) - poisson.get_V_bottomBC(0,0))/num_cell_z;  //this is  calculated correctly

    //Note: B/C I must fill from 0, and don't want to include bottomBC in V matrix, b/c is not part of V
    for (int k = 1; k <= Nz+1; k++) {
        for (int i = 1; i <= Nx+1; i++) {
            for (int j = 1; j <= Ny+1; j++) {
                V_matrix(i-1, j-1, k-1) = poisson.get_V_bottomBC(i,j) +  diff*(k);  //-1's b/c fill from 0
            }
        }
    }

    //need to permute the matrix, to be consistent with the z,y,x ordering which I use for the matrices when solving.
    //NOTE: for Tensor shuffle to work, NEED TO EXPLICTELY CREATE AN ARRAY--> this isn't clear from the documentation
    Eigen::array<ptrdiff_t, 3> permutation = {{2,1,0}}; //array should HAVE THE SPECIFIC type:  ptrdiff_t  ==> used for pointer arithmetic and array indexing.
    //ptrdiff_t is the signed integer type of the result of subtracting 2 pointers.

    temp_permuted = V_matrix.shuffle(permutation);  //shuffle permuts the tensor. Note: dimensions are indexed from 0. This is supposed to swap x and z values...
    V_matrix = temp_permuted;
    //Returns a copy of the input tensor whose dimensions have been reordered according to the specified permutation. The argument shuffle is an array of Index values. Its size is the rank of the input tensor. It must contain a permutation of 0, 1, ..., rank - 1.
    //works to here
    //NOTE: IT CAN'T DO AN INPLACE PERMUTATION!!!--> need to rename the variable!!!

    //reshape the matrix to a single column. //Note: even though reshaping to a column, V still must be a TENSOR type for this to work!
     Eigen::array<ptrdiff_t, 3> reshape_sizes = {{num_rows, 1, 1}};  //need to only go to num_rows..., b/c it fills from 0 !!
     V = V_matrix.reshape(reshape_sizes);  //Note: even though reshaping to a column, V still must be a TENSOR type for this to work!

     //prepare initial guess, for 1st iteration of bicgstab
     for (int i = 0; i < num_rows; i++) {
         V_Xd(i) = V(i,0,0);
         soln_V(i) = V(i,0,0); // do this since we are using soln_V for the inital guess
     }

    //Fill p with initial conditions (need for error calculation)
    double min_dense = continuity_p.get_p_bottomBC(1,1) < continuity_p.get_p_topBC(1,1) ? continuity_p.get_p_bottomBC(1,1):continuity_p.get_p_topBC(1,1);  //this should be same as std::min  fnc which doesn't work for some reason
    //double min_dense = std::min (continuity_n.get_n_bottomBC(1,1), continuity_p.get_p_topBC(1,1));  //Note: I defined the get fnc to take as arguments the i,j values... //Note: the bc's along bottom and top are currently uniform, so index, doesn't really matter.
    for (int i = 0; i < num_rows; i++) {
        p(i,0,0) = min_dense;  //NOTE: p is tensor now, so need () to access elements
    }

    //prepare initial guess, for 1st iteration of bicgstab
    for (int i = 0; i < num_rows; i++) {
        p_Xd(i) = p(i,0,0);  //this is working correctly
        soln_p(i) = p(i,0,0); // do this since we are using soln_p for the inital guess
    }

    //Convert the p to p_matrix
    Eigen::array<ptrdiff_t, 3> to_matrix_sizes{{Nz+1, Ny+1, Nx+1}};  //use reshape according to the  ordering of the matrix, so: z, y, x
    Eigen::Tensor<double, 3> p_matrix = p.reshape(to_matrix_sizes);
    //continuity_p.set_p_matrix(p_matrix);  //THIS FOR SOME REASON FAILS..., so don't use it...//save p_matrix to continuity_p member variable--> THIS MIGHT BE INEFFIIENT, BUT DO IT FOR NOW

    //////////////////////MAIN LOOP////////////////////////////////////////////////////////////////////////////////////////////////////////

    int iter, not_cnv_cnt, Va_cnt;
    bool not_converged;
    double error_np, old_error;  //this stores max value of the error and the value of max error from previous iteration
    std::vector<double> error_np_vector(num_rows);  //note: since n and p solutions are in vector form, can use vector form here also

    poisson.setup_matrix();  //I VERIFIED that size of sparse matrix is correct

    for (Va_cnt = 1; Va_cnt <= num_V; Va_cnt++) {  //+1 b/c 1st Va is the equil run
        not_converged = false;
        not_cnv_cnt = 0;

        Va = params.Va_min+params.increment*(Va_cnt-1);

        if (Va_cnt == 1) {
            params.use_tolerance_i();  //reset tolerance back
            params.use_w_i();
            //PhotogenRate = photogen.getPhotogenRate();    //otherwise PhotogenRate is pre-initialized to 0 in this main.cpp when declared
        }

        if (params.tolerance > 1e-5)
            std::cerr<<"ERROR: Tolerance has been increased to > 1e-5" <<std::endl;

        std::cout << "Va = " << Va <<std::endl;

        //Reset top and bottom BCs (outside of loop b/c don't change iter to iter)
        poisson.set_V_bottomBC(params, Va);
        poisson.set_V_topBC(params, Va);

       //correct through  here

        //-----------------------------------------------------------
        error_np = 1.0;
        iter = 0;

        //get's through here
        while (error_np > params.tolerance) {
            //std::cout << "Va " << Va <<std::endl;

            //-----------------Solve Poisson Equation------------------------------------------------------------------
            poisson.set_rhs(p);  //this finds netcharge and sets rhs
            //std::cout << poisson.get_sp_matrix() << std::endl;
            oldV = V;

            input = poisson.get_sp_matrix();

            //as expected, LU, is way too slow for a 3D matrix!!

            if (iter == 0) {
                BiCGStab_solver.analyzePattern(input);
                BiCGStab_solver.factorize(input);  //this computes preconditioner, Poisson matrix doesn't change, so can factorize just once
            }
            //BiCGStab_solver.compute(input);  //this computes the preconditioner..compute(input);
            //soln_V = BiCGStab_solver.solve(poisson.get_rhs());
            soln_V = BiCGStab_solver.solveWithGuess(poisson.get_rhs(), soln_V); //note: using soln_V for initial guess is faster than using V_Xd/NOTE: use solve with Guess...., b/c need initial guess
            //std::cout << "#iterations:     " << BiCGStab_solver.iterations() << std::endl;
             //std::cout << BiCGStab_solver.info() << std::endl;
            //std::cout << soln_V << std::endl;

           //CHOLESKY is not accurate!! for 3D solve

            //std::cout << poisson.get_sp_matrix() << std::endl;
             //std::cout << "Poisson solver error " << poisson.get_sp_matrix() * soln_Xd - poisson.get_rhs() << std::endl;

            //RECALL, I am starting my V vector from index of 1, corresponds to interior pts...
            for (int i = 0; i < num_rows; i++) {
                newV(i,0,0) = soln_V(i);   //fill VectorXd  rhs of the equation
                //NOTE: newV is a Tensor now, so use () to access
            }

            //Mix old and new solutions for V
            if (iter > 0)
                V  = utils.linear_mix(params, newV, oldV);
            else
                V = newV;

            //update the Eigen Vector_Xd for the initial guess--> LATER DO ALL THIS MORE EFFICIENTLY
            for (int i = 0; i < num_rows; i++) {
                V_Xd(i) = V(i,0,0);
            }

            //reshape solution to a V_matrix
            V_matrix = V.reshape(to_matrix_sizes); // reshapes using z, y, x ordering
            //need to permute before using for fullV
            temp_permuted = V_matrix.shuffle(permutation);  //need to use a temp variable for this to work

            //fill the fullV
            //SHOULD MOVE THIS TO A FUNCTION!
            for (int k = 1; k <= num_cell_z; k++)
               for (int j = 1; j <= num_cell_y; j++)
                   for (int i = 1; i <= num_cell_x; i++)
                       fullV(i,j,k) = temp_permuted(i-1,j-1,k-1);  //-1 b/c temp_permuted from V_matrix was filled from 0

            for (int j = 0; j <= num_cell_y; j++)
                for (int i = 0; i <= num_cell_x; i++)
                    fullV(i,j,0) = poisson.get_V_bottomBC(i,j);

            for (int k = 1; k <= num_cell_z; k++)
               for (int j = 1; j <= num_cell_y; j++)
                   fullV(0,j,k) = temp_permuted(num_cell_x-1,j-1,k-1);  //x BC's

            for (int k = 1; k <= num_cell_z; k++)
               for (int i = 1; i <= num_cell_x; i++)
                   fullV(i,0,k) = temp_permuted(i-1,num_cell_y-1,k-1);  //y BC's

            //fill edges
            for (int k = 1; k <= num_cell_z; k++)
                fullV(0,0,k) = temp_permuted(num_cell_x-1,0,k-1);

            //------------------------------Calculate Net Generation Rate----------------------------------------------------------

            //R_Langevin = recombo.ComputeR_Langevin(params,n,p);
            //FOR NOW CAN USE 0 FOR R_Langevin

                for (int i = 0; i < num_rows; i++) {
                    Up[i] = 0; //params.Photogen_scaling;  //This is what was used in Matlab version for testing.   photogen.getPhotogenRate()(i,j); //- R_Langevin(i,j);
                }

            //--------------------------------Solve equation for p------------------------------------------------------------

            continuity_p.setup_eqn(fullV, Up, p);  //pass it fullV...
            //std::cout << continuity_p.get_sp_matrix() << std::endl;   //Note: get rhs, returns an Eigen VectorXd

            for (int i = 0; i < num_rows; i++)
                oldp(i,0,0) = p(i,0,0);          //explicitely  copy, just in case

            input = continuity_p.get_sp_matrix();
            if (iter == 0) {
                BiCGStab_solver.analyzePattern(input);
                BiCGStab_solver.factorize(input);  //factorize only for the 1st iteration!!, since matrix doesn't change too much, this still will work!-> means it only  computes preconditioner once per Va!! //this computes preconditioner, if use along with analyzePattern (for 1st iter)
            }
            //BiCGStab_solver.compute(input);  //this computes the preconditioner..compute(input);
            soln_p = BiCGStab_solver.solveWithGuess(continuity_p.get_rhs(), soln_p);  //NOTE: if for initial guess use soln_p, INSTEAD OF p_Xd (which is the linearly mixed solution, then does't blow up!!!!!, even with 0.2 = w.
            //soln_p = BiCGStab_solver.solve(continuity_p.get_rhs());

//         std::cout << soln_p << std::endl;
//         exit(1);

            //save results back into n std::vector. RECALL, I am starting my V vector from index of 1, corresponds to interior pts...
            for (int i = 0; i < num_rows; i++) {
                newp(i,0,0) = soln_p(i);   //newp is now a tensor....
            }

            //------------------------------------------------

            //if get negative p's or n's set them = 0
            for (int i = 0; i < num_rows; i++) {
                if (newp(i,0,0) < 0.0) newp(i,0,0) = 0;
                //if (newn[i] < 0.0) newn[i] = 0;
            }

            //calculate the error
            old_error = error_np;
            int count = 0;  //for counting the error_np_vector_index

            //THIS CAN BE MOVED TO A FUNCTION IN UTILS
            std::fill(error_np_vector.begin(), error_np_vector.end(),0.0);  //refill with 0's so have fresh one
            for (int i = 0; i < num_rows; i++) {
                if (newp(i,0,0)!= 0) {
                    error_np_vector[count] = (std::abs(newp(i,0,0)-oldp(i,0,0)))/std::abs(oldp(i,0,0));
                    count++;
//                   if(iter == 2)
//                       std::cout << error_np_vector[i] << std::endl;
               }
            }

//            if(iter == 2) {
//                for (int i = 0; i < num_rows; i++ ) {
//                    std::cout <<  newp(i,0,0) << " " << oldp(i,0,0) <<  std::endl;   //FOR LARGE SYSTEMS SEEMS NEWP AND OLDP ARE EXACT SAME--> GETTING 0 ERRORS, BUT THEN BLOWS UP!!  //compare newp and oldp //newp are not 0's --> have numbers...
//                }
//                exit(1);
//            }

            error_np = *std::max_element(error_np_vector.begin(),error_np_vector.end());

            std::cout << error_np << std::endl;

            //auto decrease w if not converging
            if (error_np >= old_error)
                not_cnv_cnt = not_cnv_cnt+1;
            if (not_cnv_cnt > 1000) {  //Note: 100 is too small for C++, sometimes w is reduced when not necessary!!
                params.reduce_w();
                params.relax_tolerance();
                not_cnv_cnt = 0;
            }

            p = utils.linear_mix(params, newp, oldp);

            //update the Eig vector for initial guess --> LATER SHOULD DO ALL OF THIS MORE EFFICIENTLY
            for (int i = 0; i < num_rows; i++) {
                p_Xd(i) = p(i,0,0);
            }

            //Convert p to p_matrix
            p_matrix = p.reshape(to_matrix_sizes);  // this reshapes based on z,y,x ordering

            //need to permute before using for fullV
            temp_permuted = p_matrix.shuffle(permutation);  //need to use a temp variable for this to work

            //fill the fullp
            //SHOULD MOVE THIS TO A FUNCTION!
            for (int k = 1; k <= num_cell_z; k++)
               for (int j = 1; j <= num_cell_y; j++)
                   for (int i = 1; i <= num_cell_x; i++)
                       fullp(i,j,k) = temp_permuted(i-1,j-1,k-1);  //-1 b/c temp_permuted from V_matrix was filled from 0

            for (int j = 0; j <= num_cell_y; j++)
                for (int i = 0; i <= num_cell_x; i++)
                    fullp(i,j,0) = continuity_p.get_p_bottomBC(i,j);

            for (int k = 1; k <= num_cell_z; k++)
               for (int j = 1; j <= num_cell_y; j++)
                   fullp(0,j,k) = temp_permuted(num_cell_x-1,j-1,k-1);  //x BC's

            for (int k = 1; k <= num_cell_z; k++)
               for (int i = 1; i <= num_cell_x; i++)
                   fullp(i,0,k) = temp_permuted(i-1,num_cell_y-1,k-1);  //y BC's

            //fill edges
            for (int k = 1; k <= num_cell_z; k++)
                fullp(0,0,k) = temp_permuted(num_cell_x-1,0,k-1);

            //continuity_p.set_p_matrix(p_matrix);  //update member variable  //DON'T USE THIS B/C CAUSES ISSUES, just use the p matrix form main

           // std::cout << error_np << std::endl;
            //std::cout << "weighting factor = " << params.w << std::endl << std::endl;

            iter = iter+1;
        }

        //-------------------Calculate Currents using Scharfetter-Gummel definition--------------------------

        //continuity_n.calculate_currents();
        continuity_p.calculate_currents(fullp); //send it fullp, since we find it here anyway

        J_total_Z = continuity_p.get_Jp_Z();// + continuity_n.get_Jn_Z();
        J_total_X = continuity_p.get_Jp_X();// + continuity_n.get_Jn_X();
        J_total_Y = continuity_p.get_Jp_Y();// + continuity_n.get_Jn_Y();

//        for (int k = 1; k < num_cell_z; k++)
//            std::cout << J_total_Z(2,2,k) << std::endl;
//exit(1);

        //---------------------Write to file----------------------------------------------------------------
        utils.write_details(params, Va, fullV, fullp, J_total_Z, Up);

        if(Va_cnt >0) utils.write_JV(params, JV, iter, Va, J_total_Z);

    }//end of main loop

    JV.close();

    std::chrono::high_resolution_clock::time_point finish = std::chrono::high_resolution_clock::now();
    std::chrono::duration<double> time = std::chrono::duration_cast<std::chrono::duration<double>>(finish-start);
    std::cout << "CPU time = " << time.count() << std::endl;

    return 0;
}


#include "problem_generator.h"
#include <Eigen/Dense>
#include <random>
#include <vector>

namespace pose_lib {

static const double kPI = 3.14159265358979323846;

double CalibPoseValidator::compute_pose_error(const ProblemInstance &instance, const CameraPose &pose) {
  return (instance.pose_gt.R - pose.R).norm() + (instance.pose_gt.t - pose.t).norm() + std::abs(instance.pose_gt.alpha - pose.alpha);
}

bool CalibPoseValidator::is_valid(const ProblemInstance &instance, const CameraPose &pose, double tol) {
  if ((pose.R.transpose() * pose.R - Eigen::Matrix3d::Identity()).norm() > tol)
    return false;

  // Point to point correspondences
  // alpha * p + lambda*x = R*X + t
  for (int i = 0; i < instance.x_point_.size(); ++i) {
    double err = 1.0 - std::abs(instance.x_point_[i].dot((pose.R * instance.X_point_[i] + pose.t - pose.alpha * instance.p_point_[i]).normalized()));
    if (err > tol)
      return false;
  }

  // Point to Line correspondences
  // alpha * p + lambda * x = R*(X + mu*V) + t
  for (int i = 0; i < instance.x_line_.size(); ++i) {
    // lambda * x - mu * R*V = R*X + t - alpha * p
    // x.cross(R*V).dot(R*X+t-alpha.p) = 0
    Eigen::Vector3d X = pose.R * instance.X_line_[i] + pose.t - pose.alpha * instance.p_line_[i];
    double err = instance.x_line_[i].cross(pose.R * instance.V_line_[i]).normalized().dot(X);

    if (err > tol)
      return false;
  }

  // Line to point correspondences
  // l'*(R*X + t - alpha*p) = 0
  for (int i = 0; i < instance.l_line_point_.size(); ++i) {

    Eigen::Vector3d X = pose.R * instance.X_line_point_[i] + pose.t - pose.alpha * instance.p_line_point_[i];

    double err = std::abs(instance.l_line_point_[i].dot(X.normalized()));
    if (err > tol)
      return false;
  }

  // Line to line correspondences
  // l'*(R*(X + mu*V) + t - alpha*p) = 0
  for (int i = 0; i < instance.l_line_line_.size(); ++i) {

    Eigen::Vector3d X = pose.R * instance.X_line_line_[i] + pose.t - pose.alpha * instance.p_line_line_[i];
    Eigen::Vector3d V = pose.R * instance.V_line_line_[i];

    double err = std::abs(instance.l_line_line_[i].dot(X.normalized())) + std::abs(instance.l_line_line_[i].dot(V.normalized()));
    if (err > tol)
      return false;
  }

  return true;
}

double UnknownFocalValidator::compute_pose_error(const ProblemInstance &instance, const CameraPose &pose) {
  return (instance.pose_gt.R - pose.R).norm() + (instance.pose_gt.t - pose.t).norm() + std::abs(instance.pose_gt.alpha - pose.alpha);
}

bool UnknownFocalValidator::is_valid(const ProblemInstance &instance, const CameraPose &pose, double tol) {
  if ((pose.R.transpose() * pose.R - Eigen::Matrix3d::Identity()).norm() > tol)
    return false;

  if (pose.alpha < 0)
    return false;

  Eigen::Matrix3d Kinv;
  Kinv.setIdentity();
  Kinv(2, 2) = pose.alpha;
  // lambda*diag(1,1,alpha)*x = R*X + t
  for (int i = 0; i < instance.x_point_.size(); ++i) {
    double err = 1.0 - std::abs((Kinv * instance.x_point_[i]).normalized().dot((pose.R * instance.X_point_[i] + pose.t).normalized()));
    if (err > tol)
      return false;
  }

  return true;
}

double RadialPoseValidator::compute_pose_error(const ProblemInstance &instance, const CameraPose &pose) {
  // Only compute up to sign for radial cameras

  double err1 = (instance.pose_gt.R.topRows(2) - pose.R.topRows(2)).norm() + (instance.pose_gt.t.topRows(2) - pose.t.topRows(2)).norm();
  double err2 = (instance.pose_gt.R.topRows(2) + pose.R.topRows(2)).norm() + (instance.pose_gt.t.topRows(2) + pose.t.topRows(2)).norm();

  return std::min(err1, err2);
}

bool RadialPoseValidator::is_valid(const ProblemInstance &instance, const CameraPose &pose, double tol) {
  if ((pose.R.transpose() * pose.R - Eigen::Matrix3d::Identity()).norm() > tol)
    return false;

  // Point to point correspondences -- Convert these to line correspondences
  // alpha * p + lambda*x = R*X + t
  for (int i = 0; i < instance.x_point_.size(); ++i) {
    Eigen::Vector3d radial_line{-instance.x_point_[i](1),instance.x_point_[i](0), 0.0};
    Eigen::Vector3d X = pose.R * instance.X_point_[i] + pose.t;
    double err = std::abs(radial_line.dot(X.normalized()));
    if (err > tol)
      return false;
  }

  // Line to point correspondences
  // l'*(R*X + t) = 0
  for (int i = 0; i < instance.l_line_point_.size(); ++i) {
    Eigen::Vector3d X = pose.R * instance.X_line_point_[i] + pose.t;

    double err = std::abs(instance.l_line_point_[i].dot(X.normalized()));
    if (err > tol)
      return false;
  }

  return true;
}

void set_random_pose(CameraPose &pose, bool upright) {
  if (upright) {
    Eigen::Vector2d r;
    r.setRandom().normalize();
    //pose.R << r(0), 0.0, r(1), 0.0, 1.0, 0.0, -r(1), 0.0, r(0); // y-gravity
    pose.R << r(0), r(1), 0.0, -r(1), r(0), 0.0, 0.0, 0.0, 1.0; // z-gravity
  } else {
    pose.R = Eigen::Quaternion<double>::UnitRandom();
  }
  pose.t.setRandom();
}

void generate_problems(int n_problems, std::vector<ProblemInstance> *problem_instances,
                       const ProblemOptions &options) {
  problem_instances->clear();
  problem_instances->reserve(n_problems);

  double fov_scale = std::tan(options.camera_fov_ / 2.0 * kPI / 180.0);

  // Random generators
  std::default_random_engine random_engine;
  std::uniform_real_distribution<double> depth_gen(options.min_depth_, options.max_depth_);
  std::uniform_real_distribution<double> coord_gen(-fov_scale, fov_scale);
  std::uniform_real_distribution<double> scale_gen(options.min_scale_, options.max_scale_);
  std::uniform_real_distribution<double> focal_gen(options.min_focal_, options.max_focal_);
  std::normal_distribution<double> direction_gen(0.0, 1.0);
  std::normal_distribution<double> offset_gen(0.0, 1.0);

  for (int i = 0; i < n_problems; ++i) {
    ProblemInstance instance;
    set_random_pose(instance.pose_gt, options.upright_);

    if (options.unknown_scale_) {
      instance.pose_gt.alpha = scale_gen(random_engine);
    } else if (options.unknown_focal_) {
      instance.pose_gt.alpha = focal_gen(random_engine);
    }

    // Point to point correspondences
    instance.x_point_.reserve(options.n_point_point_);
    instance.X_point_.reserve(options.n_point_point_);
    instance.p_point_.reserve(options.n_point_point_);
    for (int j = 0; j < options.n_point_point_; ++j) {

      Eigen::Vector3d p{0.0, 0.0, 0.0};
      Eigen::Vector3d x{coord_gen(random_engine), coord_gen(random_engine), 1.0};
      x.normalize();
      Eigen::Vector3d X;

      if (options.generalized_) {
        p << offset_gen(random_engine), offset_gen(random_engine), offset_gen(random_engine);
      }

      X = instance.pose_gt.alpha * p + x * depth_gen(random_engine);

      X = instance.pose_gt.R.transpose() * (X - instance.pose_gt.t);

      if (options.unknown_focal_) {
        x.block<2, 1>(0, 0) *= instance.pose_gt.alpha;
        x.normalize();
      }

      instance.x_point_.push_back(x);
      instance.X_point_.push_back(X);
      instance.p_point_.push_back(p);
    }

    // Point to line correspondences
    instance.x_line_.reserve(options.n_point_line_);
    instance.X_line_.reserve(options.n_point_line_);
    instance.V_line_.reserve(options.n_point_line_);
    instance.p_line_.reserve(options.n_point_line_);
    for (int j = 0; j < options.n_point_line_; ++j) {
      Eigen::Vector3d p{0.0, 0.0, 0.0};
      Eigen::Vector3d x{coord_gen(random_engine), coord_gen(random_engine), 1.0};
      x.normalize();
      Eigen::Vector3d X;

      if (options.generalized_) {
        p << offset_gen(random_engine), offset_gen(random_engine), offset_gen(random_engine);
      }
      X = instance.pose_gt.alpha * p + x * depth_gen(random_engine);
      X = instance.pose_gt.R.transpose() * (X - instance.pose_gt.t);

      Eigen::Vector3d V{direction_gen(random_engine), direction_gen(random_engine), direction_gen(random_engine)};
      V.normalize();

      // Translate X such that X.dot(V) = 0
      X = X - V.dot(X) * V;

      if (options.unknown_focal_) {
        // TODO implement this.
      }

      instance.x_line_.push_back(x);
      instance.X_line_.push_back(X);
      instance.V_line_.push_back(V);
      instance.p_line_.push_back(p);
    }

    // Line to point correspondences
    instance.l_line_point_.reserve(options.n_line_point_);
    instance.X_line_point_.reserve(options.n_line_point_);
    instance.p_line_point_.reserve(options.n_line_point_);
    for (int j = 0; j < options.n_line_point_; ++j) {
      Eigen::Vector3d p{0.0, 0.0, 0.0};
      Eigen::Vector3d x{coord_gen(random_engine), coord_gen(random_engine), 1.0};
      x.normalize();
      Eigen::Vector3d X;

      if (options.generalized_) {
        p << offset_gen(random_engine), offset_gen(random_engine), offset_gen(random_engine);
      }
      X = instance.pose_gt.alpha * p + x * depth_gen(random_engine);
      X = instance.pose_gt.R.transpose() * (X - instance.pose_gt.t);

      // Cross product with random vector to generate line
      Eigen::Vector3d l;
      if(options.radial_lines_) {
        // Line passing through image center
        l = x.cross(Eigen::Vector3d{0.0, 0.0, 1.0});
      } else {
        // Random line
        l = x.cross(Eigen::Vector3d(direction_gen(random_engine), direction_gen(random_engine), direction_gen(random_engine)));
      }

      l.normalize();

      if (options.unknown_focal_) {
        // TODO implement this.
      }

      instance.l_line_point_.push_back(l);
      instance.X_line_point_.push_back(X);
      instance.p_line_point_.push_back(p);
    }

    // Line to line correspondences
    instance.l_line_line_.reserve(options.n_line_line_);
    instance.X_line_line_.reserve(options.n_line_line_);
    instance.V_line_line_.reserve(options.n_line_line_);
    instance.p_line_line_.reserve(options.n_line_line_);
    for (int j = 0; j < options.n_line_line_; ++j) {
      Eigen::Vector3d p{0.0, 0.0, 0.0};
      Eigen::Vector3d x{coord_gen(random_engine), coord_gen(random_engine), 1.0};
      x.normalize();
      Eigen::Vector3d X;

      if (options.generalized_) {
        p << offset_gen(random_engine), offset_gen(random_engine), offset_gen(random_engine);
      }
      X = instance.pose_gt.alpha * p + x * depth_gen(random_engine);
      X = instance.pose_gt.R.transpose() * (X - instance.pose_gt.t);

      Eigen::Vector3d V{direction_gen(random_engine), direction_gen(random_engine), direction_gen(random_engine)};
      V.normalize();

      // Translate X such that X.dot(V) = 0
      X = X - V.dot(X) * V;

      Eigen::Vector3d l = x.cross(instance.pose_gt.R * V);
      l.normalize();

      if (options.unknown_focal_) {
        // TODO implement this.
      }

      instance.l_line_line_.push_back(l);
      instance.X_line_line_.push_back(X);
      instance.V_line_line_.push_back(V);
      instance.p_line_line_.push_back(p);
    }

    problem_instances->push_back(instance);
  }
}

}; // namespace pose_lib