/*
 * Copyright (c) 2021 The Foundation for Research on Information Technologies in Society (IT'IS).
 *
 * This file is part of iSEG
 * (see https://github.com/ITISFoundation/osparc-iseg).
 *
 * This software is released under the MIT License.
 *  https://opensource.org/licenses/MIT
 */
#include "AutoTubeWidget.h"

#include "Data/ItkUtils.h"
#include "Data/Logger.h"
#include "Data/SlicesHandlerITKInterface.h"

#include "Thirdparty/IJ/BinaryThinningImageFilter3D/itkBinaryThinningImageFilter3D.h"
#include "Thirdparty/IJ/NonMaxSuppression/itkNonMaxSuppressionImageFilter.h"

#include <itkBinaryThinningImageFilter.h>
#include <itkBinaryThresholdImageFilter.h>
#include <itkConnectedComponentImageFilter.h>
#include <itkCurvesLevelSetImageFilter.h>
#include <itkFastMarchingImageFilter.h>
#include <itkGradientMagnitudeRecursiveGaussianImageFilter.h>
#include <itkHessianToObjectnessMeasureImageFilter.h>
#include <itkImage.h>
#include <itkMinimumMaximumImageCalculator.h>
#include <itkMultiScaleHessianBasedMeasureImageFilter.h>
#include <itkRelabelComponentImageFilter.h>
#include <itkRescaleIntensityImageFilter.h>
#include <itkSigmoidImageFilter.h>
#include <itkSliceBySliceImageFilter.h>
#include <itkThresholdImageFilter.h>

#include <accumulators/percentile.hpp>
#include <boost/accumulators/accumulators.hpp>
#include <boost/accumulators/statistics/max.hpp>
#include <boost/accumulators/statistics/mean.hpp>
#include <boost/accumulators/statistics/min.hpp>
#include <boost/accumulators/statistics/stats.hpp>
#include <boost/accumulators/statistics/variance.hpp>
#include <boost/algorithm/string/classification.hpp>
#include <boost/algorithm/string/join.hpp>
#include <boost/algorithm/string/split.hpp>
#include <boost/algorithm/string/trim.hpp>
#include <boost/range/adaptor/transformed.hpp>

#include <QFormLayout>
#include <QMessageBox>
#include <QScrollArea>

#include <algorithm>
#include <sstream>

namespace acc = boost::accumulators;
using boost::adaptors::transformed;
using boost::algorithm::join;

template<typename TInputImage, typename TOutputImage, unsigned int Dimension>
class BinaryThinningImageFilter : public itk::BinaryThinningImageFilter<TInputImage, TOutputImage>
{
};

template<typename TInputImage, typename TOutputImage>
class BinaryThinningImageFilter<TInputImage, TOutputImage, 3> : public itk::BinaryThinningImageFilter3D<TInputImage, TOutputImage>
{
};

AutoTubeWidget::AutoTubeWidget(iseg::SlicesHandlerInterface* hand3D)
		: m_Handler3D(hand3D)
{
	setToolTip(Format("Sandbox Tool"));

	m_Metric2d = new QCheckBox;
	m_Metric2d->setChecked(false);
	m_Metric2d->setToolTip(Format("Use 2D pseudo vesselness (blob feature per slice), or use full 3D vesselness."));

	m_NonMaxSuppression = new QCheckBox;
	m_NonMaxSuppression->setChecked(true);
	m_NonMaxSuppression->setToolTip(Format("Extract approx. one pixel wide paths based on non-maximum suppression."));

	m_Skeletonize = new QCheckBox;
	m_Skeletonize->setChecked(true);
	m_Skeletonize->setToolTip(Format("Compute 1-pixel wide centerlines (skeleton)."));

	m_SigmaLow = new QLineEdit(QString::number(0.3));
	m_SigmaLow->setValidator(new QDoubleValidator);

	m_SigmaHi = new QLineEdit(QString::number(0.6));
	m_SigmaHi->setValidator(new QDoubleValidator);

	m_NumberSigmaLevels = new QLineEdit(QString::number(2));
	m_NumberSigmaLevels->setValidator(new QIntValidator);

	m_Threshold = new QLineEdit;
	m_Threshold->setValidator(new QDoubleValidator);

	m_MaxRadius = new QLineEdit(QString::number(1));
	m_MaxRadius->setValidator(new QDoubleValidator);

	m_MinObjectSize = new QLineEdit(QString::number(10));
	m_MinObjectSize->setValidator(new QIntValidator);

	m_SelectObjectsButton = new QPushButton("Select Mask");
	m_SelectedObjects = new QLineEdit;
	m_SelectedObjects->setReadOnly(true);

	m_ExecuteButton = new QPushButton("Execute");

	auto layout = new QFormLayout;
	layout->addRow("Sigma Min", m_SigmaLow);
	layout->addRow("Sigma Max", m_SigmaHi);
	layout->addRow("Number of Sigmas", m_NumberSigmaLevels);
	layout->addRow("2D Vesselness", m_Metric2d);
	layout->addRow("Feature Threshold", m_Threshold);
	layout->addRow("Non-maximum Suppression", m_NonMaxSuppression);
	layout->addRow("Centerlines", m_Skeletonize);
	//layout->addRow("Maximum radius", _max_radius);
	layout->addRow("Minimum object size", m_MinObjectSize);
	layout->addRow(m_SelectObjectsButton, m_SelectedObjects);
	layout->addRow(m_ExecuteButton);

	auto big_view = new QWidget;
	big_view->setLayout(layout);

	auto scroll_area = new QScrollArea(this);
	scroll_area->setWidget(big_view);

	auto top_layout = new QGridLayout(1, 1);
	top_layout->addWidget(scroll_area, 0, 0);
	setLayout(top_layout);

	QObject_connect(m_SelectObjectsButton, SIGNAL(clicked()), this, SLOT(SelectObjects()));
	QObject_connect(m_ExecuteButton, SIGNAL(clicked()), this, SLOT(DoWork()));
}

void AutoTubeWidget::Init()
{
	OnSlicenrChanged();
	HideParamsChanged();
}

void AutoTubeWidget::NewLoaded()
{
	OnSlicenrChanged();
}

void AutoTubeWidget::OnSlicenrChanged()
{
}

void AutoTubeWidget::Cleanup()
{
	m_SelectedObjects->setText("");
	m_CachedFeatureImage.img = nullptr;
}

void AutoTubeWidget::OnMouseClicked(iseg::Point p)
{
}

void AutoTubeWidget::SelectObjects()
{
	auto sel = m_Handler3D->TissueSelection();
	std::cout << "sel " << sel[0] << std::endl;
	std::string text = join(sel | transformed([](int d) { return std::to_string(d); }), ", ");
	m_SelectedObjects->setText(QString::fromStdString(text));
}

void AutoTubeWidget::DoWork()
{
	iseg::SlicesHandlerITKInterface itk_handler(m_Handler3D);
	try
	{
		if ((true)) //(all_slices->isChecked())
		{
			using input_type = itk::SliceContiguousImage<float>;
			using tissues_type = itk::SliceContiguousImage<iseg::tissues_size_t>;
			auto source = itk_handler.GetSource(true); // active_slices -> correct seed z-position
			auto target = itk_handler.GetTarget(true);
			auto tissues = itk_handler.GetTissues(true);
			DoWorkNd<input_type, tissues_type, input_type>(source, tissues, target);
		}
		else
		{
			using input_type = itk::Image<float, 2>;
			using tissues_type = itk::Image<iseg::tissues_size_t, 2>;
			auto source = itk_handler.GetSourceSlice();
			auto target = itk_handler.GetTargetSlice();
			auto tissues = itk_handler.GetTissuesSlice();
			DoWorkNd<input_type, tissues_type, input_type>(source, tissues, target);
		}
	}
	catch (const std::exception& e)
	{
		QMessageBox::warning(this, "iSeg", QString("Error: ") + e.what(), QMessageBox::Ok | QMessageBox::Default);
	}
}

template<class TInput, class TImage>
typename TImage::Pointer AutoTubeWidget::ComputeFeatureImage(TInput* source) const
{
	itkStaticConstMacro(ImageDimension, size_t, TInput::ImageDimension);
	using hessian_pixel_type = itk::SymmetricSecondRankTensor<float, ImageDimension>;
	using hessian_image_type = itk::Image<hessian_pixel_type, ImageDimension>;
	using feature_filter_type = itk::HessianToObjectnessMeasureImageFilter<hessian_image_type, TImage>;
	using multiscale_hessian_filter_type = itk::MultiScaleHessianBasedMeasureImageFilter<TInput, hessian_image_type, TImage>;

	double sigm_min = m_SigmaLow->text().toDouble();
	double sigm_max = m_SigmaHi->text().toDouble();
	int num_levels = m_NumberSigmaLevels->text().toInt();

	auto objectness_filter = feature_filter_type::New();
	objectness_filter->SetBrightObject(false);
	objectness_filter->SetObjectDimension(1);
	objectness_filter->SetScaleObjectnessMeasure(true);
	objectness_filter->SetAlpha(0.5);
	objectness_filter->SetBeta(0.5);
	objectness_filter->SetGamma(5.0);

	auto multi_scale_enhancement_filter = multiscale_hessian_filter_type::New();
	multi_scale_enhancement_filter->SetInput(source);
	multi_scale_enhancement_filter->SetHessianToMeasureFilter(objectness_filter);
	multi_scale_enhancement_filter->SetSigmaStepMethodToEquispaced();
	multi_scale_enhancement_filter->SetSigmaMinimum(std::min(sigm_min, sigm_max));
	multi_scale_enhancement_filter->SetSigmaMaximum(std::max(sigm_min, sigm_max));
	multi_scale_enhancement_filter->SetNumberOfSigmaSteps(std::min(1, num_levels));
	multi_scale_enhancement_filter->Update();
	return multi_scale_enhancement_filter->GetOutput();
}

template<class TInput, class TImage>
typename TImage::Pointer AutoTubeWidget::ComputeFeatureImage2d(TInput* source) const
{
	itkStaticConstMacro(ImageDimension, size_t, TInput::ImageDimension);
	using input_type = TInput;
	using output_type = TImage;
	using image_type = itk::Image<float, 2>;

	using hessian_filter_type = itk::HessianRecursiveGaussianImageFilter<image_type>;
	using hessian_pixel_type = itk::SymmetricSecondRankTensor<double, 2>;
	using hessian_image_type = itk::Image<hessian_pixel_type, 2>;

	using feature_filter_type = itk::HessianToObjectnessMeasureImageFilter<hessian_image_type, image_type>;
	using multiscale_hessian_filter_type = itk::MultiScaleHessianBasedMeasureImageFilter<image_type, hessian_image_type, image_type>;
	using slice_by_slice_filter_type = itk::SliceBySliceImageFilter<input_type, output_type, multiscale_hessian_filter_type>;

	double sigm_min = m_SigmaLow->text().toDouble();
	double sigm_max = m_SigmaHi->text().toDouble();
	int num_levels = m_NumberSigmaLevels->text().toInt();

	auto objectness_filter = feature_filter_type::New();
	objectness_filter->SetBrightObject(false);
	objectness_filter->SetObjectDimension(ImageDimension == 2 ? 1 : 0); // for 2D analysis we are looking for lines in a single slice
	objectness_filter->SetScaleObjectnessMeasure(true);
	objectness_filter->SetAlpha(0.5);
	objectness_filter->SetBeta(0.5);
	objectness_filter->SetGamma(5.0);

	auto multi_scale_enhancement_filter = multiscale_hessian_filter_type::New();
	multi_scale_enhancement_filter->SetHessianToMeasureFilter(objectness_filter);
	multi_scale_enhancement_filter->SetSigmaStepMethodToEquispaced();
	multi_scale_enhancement_filter->SetSigmaMinimum(std::min(sigm_min, sigm_max));
	multi_scale_enhancement_filter->SetSigmaMaximum(std::max(sigm_min, sigm_max));
	multi_scale_enhancement_filter->SetNumberOfSigmaSteps(std::min(1, num_levels));

	auto slice_filter = slice_by_slice_filter_type::New();
	slice_filter->SetInput(source);
	slice_filter->SetFilter(multi_scale_enhancement_filter);
	slice_filter->Update();

	return slice_filter->GetOutput();
}

template<class TInput, class TTissue, class TTarget>
void AutoTubeWidget::DoWorkNd(TInput* source, TTissue* tissues, TTarget* target)
{
	itkStaticConstMacro(ImageDimension, size_t, TInput::ImageDimension);
	using input_type = TInput;
	using real_type = itk::Image<float, ImageDimension>;
	using mask_type = itk::Image<unsigned char, ImageDimension>;
	using labelfield_type = itk::Image<unsigned short, ImageDimension>;

	using threshold_filter_type = itk::BinaryThresholdImageFilter<real_type, mask_type>;
	using nonmax_filter_type = itk::NonMaxSuppressionImageFilter<real_type>;
	using thinnning_filter_type = BinaryThinningImageFilter<mask_type, mask_type, ImageDimension>;

	typename real_type::Pointer feature_image;
	std::vector<double> feature_params;
	feature_params.push_back(m_SigmaLow->text().toDouble());
	feature_params.push_back(m_SigmaHi->text().toDouble());
	feature_params.push_back(m_NumberSigmaLevels->text().toInt());
	feature_params.push_back(m_Metric2d->isChecked());
	if (!m_CachedFeatureImage.Get(feature_image, feature_params))
	{
		feature_image = m_Metric2d->isChecked()
												? ComputeFeatureImage2d<input_type, real_type>(source)
												: ComputeFeatureImage<input_type, real_type>(source);
		m_CachedFeatureImage.Store(feature_image, feature_params);
		m_CachedSkeleton.img = nullptr;
	}

	// initialize threshold to reasonable value
	bool ok = false;
	float lower = m_Threshold->text().toFloat(&ok);
	if (!ok)
	{
		auto calculator = itk::MinimumMaximumImageCalculator<real_type>::New();
		calculator->SetImage(feature_image);
		calculator->Compute();

		auto min_gm = calculator->GetMinimum();
		auto max_gm = calculator->GetMaximum();

		lower = min_gm + 0.8 * (max_gm - min_gm); // stupid way to guess threshold
		m_Threshold->setText(QString::number(lower));
	}

	// extract IDs if any were set
	std::vector<int> object_ids;
	if (!m_SelectedObjects->text().isEmpty())
	{
		std::vector<std::string> tokens;
		std::string selected_objects_text = m_SelectedObjects->text().toStdString();
		boost::algorithm::split(tokens, selected_objects_text, boost::algorithm::is_any_of(","));
		std::transform(tokens.begin(), tokens.end(), std::back_inserter(object_ids), [](std::string s) {
			boost::algorithm::trim(s);
			return stoi(s);
		});
	}

	// mask feature image before skeletonization
	if (!object_ids.empty())
	{
		using map_functor_type = iseg::Functor::MapLabels<unsigned short, unsigned char>;
		map_functor_type map;
		map.m_Map.assign(m_Handler3D->TissueNames().size() + 1, 0);
		for (size_t i = 0; i < object_ids.size(); i++)
		{
			map.m_Map.at(object_ids[i]) = 1;
		}

		auto map_filter = itk::UnaryFunctorImageFilter<TTissue, mask_type, map_functor_type>::New();
		map_filter->SetFunctor(map);
		map_filter->SetInput(tissues);

		auto masker = itk::MaskImageFilter<real_type, mask_type>::New();
		masker->SetInput(feature_image);
		masker->SetMaskImage(map_filter->GetOutput());
		SAFE_UPDATE(masker, return );
		feature_image = masker->GetOutput();
	}

	typename mask_type::Pointer skeleton;
	std::vector<double> skeleton_params(object_ids.begin(), object_ids.end());
	skeleton_params.push_back(lower);
	skeleton_params.push_back(m_NonMaxSuppression->isChecked());
	skeleton_params.push_back(m_Skeletonize->isChecked());
	skeleton_params.push_back(m_MinObjectSize->text().toInt());
	if (!m_CachedSkeleton.Get(skeleton, skeleton_params))
	{
		// disconnect bright tubes via non-maxi suppression
		if (m_NonMaxSuppression->isChecked())
		{
			auto masking = itk::ThresholdImageFilter<real_type>::New();
			masking->SetInput(feature_image);
			masking->ThresholdBelow(lower);
			masking->SetOutsideValue(std::min(lower, 0.f));

			// do thinning to get brightest pixels ("one" pixel wide)
			auto nonmax_filter = nonmax_filter_type::New();
			nonmax_filter->SetInput(masking->GetOutput());

			auto threshold = threshold_filter_type::New();
			threshold->SetInput(nonmax_filter->GetOutput());
			threshold->SetLowerThreshold(std::nextafter(std::min(lower, 0.f), 1.f));
			SAFE_UPDATE(threshold, return );
			skeleton = threshold->GetOutput();
		}
		else
		{
			auto threshold = threshold_filter_type::New();
			threshold->SetInput(feature_image);
			threshold->SetLowerThreshold(lower);
			SAFE_UPDATE(threshold, return );
			skeleton = threshold->GetOutput();
		}

		if (m_Skeletonize->isChecked())
		{
			// get centerline: either thresholding or non-max suppression must be done before this
			auto thinning = thinnning_filter_type::New();
			thinning->SetInput(skeleton);

			auto rescale = itk::RescaleIntensityImageFilter<mask_type>::New();
			rescale->SetInput(thinning->GetOutput());
			rescale->SetOutputMinimum(0);
			rescale->SetOutputMaximum(255);
			rescale->InPlaceOn();

			SAFE_UPDATE(rescale, return );
			skeleton = rescale->GetOutput();
		}
		m_CachedSkeleton.Store(skeleton, skeleton_params);
	}

	typename mask_type::Pointer output;
	if (skeleton && m_MinObjectSize->text().toInt() > 1)
	{
		auto connectivity = itk::ConnectedComponentImageFilter<mask_type, labelfield_type>::New();
		connectivity->SetInput(skeleton);
		connectivity->FullyConnectedOn();

		auto relabel = itk::RelabelComponentImageFilter<labelfield_type, labelfield_type>::New();
		relabel->SetInput(connectivity->GetOutput());
		relabel->SetMinimumObjectSize(m_MinObjectSize->text().toInt());

		auto threshold = itk::BinaryThresholdImageFilter<labelfield_type, mask_type>::New();
		threshold->SetInput(relabel->GetOutput());
		threshold->SetLowerThreshold(1);
		SAFE_UPDATE(threshold, return );
		output = threshold->GetOutput();
	}

	iseg::DataSelection data_selection;
	data_selection.allSlices = true; // all_slices->isChecked();
	data_selection.sliceNr = m_Handler3D->ActiveSlice();
	data_selection.work = true;
	emit BeginDatachange(data_selection, this);

	if (output && iseg::Paste<mask_type, input_type>(output, target))
	{
		// good, else maybe output is not defined
	}
	else if (!iseg::Paste<mask_type, input_type>(skeleton, target))
	{
		std::cerr << "Error: could not set output because image regions don't match.\n";
	}

	emit EndDatachange(this);
}


/* * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * */
/*                                                                           */
/*  This file is part of the library KASKADE 7                               */
/*    see http://www.zib.de/projects/kaskade7-finite-element-toolbox         */
/*                                                                           */
/*  Copyright (C) 2002-2012 Zuse Institute Berlin                            */
/*                                                                           */
/*  KASKADE 7 is distributed under the terms of the ZIB Academic License.    */
/*    see $KASKADE/academic.txt                                              */
/*                                                                           */
/* * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * */

#ifndef ILUPRECOND_HH
#define ILUPRECOND_HH

#include "globheads.h"
#include "protospp.h"
#include "ios.h"
extern "C" void set_arms_pars(io_t* io,  int Dscale, int *ipar, double *tolcoef, int *lfil);

#include <boost/timer/timer.hpp>
#include "dune/istl/preconditioners.hh"
#include "linalg/triplet.hh"

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

namespace Kaskade
{
/**
 * \ingroup linalg
 *
 * PrecondType::ILUT Preconditioner from Saads ITSOL library
 */
template <class Op>
class ILUTPreconditioner: public Dune::Preconditioner<typename Op::Range, typename Op::Range>
{
  typedef typename Op::Range Range;
  typedef Range               Domain;
  typedef typename Op::field_type field_type;

public:
  static int const category = Dune::SolverCategory::sequential;

  ILUTPreconditioner(Op& op, int lfil=240, double tol=1.0e-2, int verbosity=2) {
    boost::timer::cpu_timer iluTimer;
    std::cout << "ilut constructor: ";
    MatrixAsTriplet<field_type> A = op.template get<MatrixAsTriplet<field_type> >();
    int n = A.nrows(), nnz = A.nnz(), ierr;
    FILE *flog = stdout;

    csmat = (csptr)Malloc( sizeof(SparMat), "main" );
    ierr = COOcs(n, nnz,  &A.data[0], &A.cidx[0],  &A.ridx[0], csmat);
    if( ierr != 0 )
      {
        printf(" *** ILU error - COOcs code %d csmat=%p\n", ierr, csmat);
        throw ierr;
      }

    lu = (iluptr)Malloc(sizeof(ILUSpar), "main");
    ierr = ilut(csmat, lu, lfil, tol, flog);
    if( ierr != 0 )
      {
        printf(" *** PrecondType::ILUT error - ilut code %d lu=%p\n", ierr, lu);
        throw ierr;
      }

    xx   = (double *)Malloc(n*sizeof(double), "main");
    yy   = (double *)Malloc(n*sizeof(double), "main");

    if ( verbosity>=2 )
    {
	std::cout << "PrecondType::ILUT: n=" << n << ", nnz=" << nnz << " dropTol=" << tol << ", fillfac="
	          << 1.0*nnz_ilu(lu)/(nnz+1.0) << ", lfil=" << lfil << ", time=" << (double)(iluTimer.elapsed().user)/1e9 << "s\n";
	}
  }
  ~ILUTPreconditioner()
    {
      cleanILU(lu);
      cleanCS(csmat);
      free(xx);
      free(yy);
    }

  virtual void pre (Domain&, Range&) {}
  virtual void post (Domain&) {}

  virtual void apply (Domain& x, Range const& y) {
    y.write(yy);
    lusolC(yy, xx, lu);
    x.read(xx);
  }

private:
  csptr csmat;
  iluptr lu;
  double *xx, *yy;
};

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

/**
 * \ingroup linalg
 *
 * PrecondType::ILUT Preconditioner from Saads ITSOL library
 */
template <class Op>
class ILUKPreconditioner: public Dune::Preconditioner<typename Op::Range, typename Op::Range>
{
  typedef typename Op::Range Range;
  typedef Range               Domain;
  typedef typename Op::field_type field_type;

public:
  static int const category = Dune::SolverCategory::sequential;

  ILUKPreconditioner(Op& op, int fill_lev=3, int verbosity=2) {
    boost::timer::cpu_timer iluTimer;
    MatrixAsTriplet<field_type> A = op.template get<MatrixAsTriplet<field_type> >();
    int n = A.nrows(), nnz = A.nnz(), ierr;
    FILE *flog = stdout;

    csmat = (csptr)Malloc( sizeof(SparMat), "main" );
    ierr = COOcs(n, nnz,  &A.data[0], &A.cidx[0],  &A.ridx[0], csmat);
    if( ierr != 0 )
      {
        printf(" *** ILU error - COOcs code %d csmat=%p\n", ierr, csmat);
        throw ierr;
      }

    lu = (iluptr)Malloc(sizeof(ILUSpar), "main");
    ierr = ilukC(fill_lev, csmat, lu, flog);
    if( ierr != 0 )
      {
        printf(" *** PrecondType::ILUK error - ilut code %d lu=%p\n", ierr, lu);
        throw ierr;
      }

    xx   = (double *)Malloc(n*sizeof(double), "main");
    yy   = (double *)Malloc(n*sizeof(double), "main");

    if ( verbosity>=2 )
    {
	std::cout << "PrecondType::ILUK: n=" << n << ", nnz=" << nnz << ", fillfac="
	          << 1.0*nnz_ilu(lu)/(nnz+1.0) << ", fill_lev="
	          << fill_lev << ", time=" << (double)(iluTimer.elapsed().user)/1e9 << "s\n";
    }
  }
  ~ILUKPreconditioner()
    {
      cleanILU(lu);
      cleanCS(csmat);
      free(xx);
      free(yy);
    }

  virtual void pre (Domain&, Range&) {}
  virtual void post (Domain&) {}

  virtual void apply (Domain& x, Range const& y) {
    y.write(yy);
    lusolC(yy, xx, lu);
    x.read(xx);
  }

private:
  csptr csmat;
  iluptr lu;
  double *xx, *yy;
};
//---------------------------------------------------------------------

/**
 * \ingroup linalg
 *
 * PrecondType::ILUT Preconidoner from Saads ITSOL library
 */
template <class Op>
class ARMSPreconditioner: public Dune::Preconditioner<typename Op::Range, typename Op::Range>
{
  typedef typename Op::Range Range;
  typedef Range               Domain;
  typedef typename Op::field_type field_type;

public:
  static int const category = Dune::SolverCategory::sequential;

  ARMSPreconditioner(Op& op, int lfil=4, double tol=1.0e-2, int lev_reord=1, double tolind = 0.2, int verbosity=2) {
    boost::timer::cpu_timer iluTimer;
    MatrixAsTriplet<field_type> A = op.template get<MatrixAsTriplet<field_type> >();
    int n = A.nrows(), nnz = A.nnz(), ierr, diagscal = 1;

    int lfil_arr[7];
    double droptol[7], dropcoef[7];
    int ipar[18];
    io_t io;
    FILE *flog = stdout;

    csmat = (csptr)Malloc( sizeof(SparMat), "main" );
    ierr = COOcs(n, nnz,  &A.data[0], &A.cidx[0],  &A.ridx[0], csmat);
    if( ierr != 0 )
      {
        printf(" *** PrecondType::ARMS error - COOcs code %d csmat=%p\n", ierr, csmat);
        throw ierr;
      }

    memset(&io, 0, sizeof(io) );
    io.perm_type = 0;
    io.Bsize = 400;
    set_arms_pars(&io, diagscal, ipar, dropcoef, lfil_arr);
    for (int j=0; j<7; j++)
      {
         lfil_arr[j] = lfil*((int) nnz/n);
         droptol[j] =  tol*dropcoef[j];
      }
    ipar[1] = lev_reord;
    ArmsSt = (arms) Malloc(sizeof(armsMat),"main:ArmsSt");
    setup_arms(ArmsSt);

    ierr = arms2(csmat, ipar, droptol, lfil_arr, tolind, ArmsSt, flog);
    if( ierr != 0 )
      {
        printf(" *** PrecondType::ARMS error - arms2 code %d ArmsSt=%p\n", ierr, ArmsSt);
        throw ierr;
      }

    yy   = (double *)Malloc(n*sizeof(double), "main");

    if ( verbosity>=2 )
    {
	std::cout << "PrecondType::ARMS: n=" << n << ", nnz=" << nnz << ", dropTol=" << tol
	          << ", lev_reord=" << lev_reord
	          << ", lfil=" << lfil << ", time=" << (double)(iluTimer.elapsed().user)/1e9 << "s\n";
    }
  }
  ~ARMSPreconditioner()
    {
      cleanARMS(ArmsSt);
      cleanCS(csmat);
      free(yy);
    }

  virtual void pre (Domain&, Range&) {}
  virtual void post (Domain&) {}

  virtual void apply (Domain& x, Range const& y) {
    y.write(yy);
    armsol2(yy, ArmsSt);
    x.read(yy);
  }

private:
  csptr csmat;
  arms ArmsSt;
  double *yy;
};

}  // namespace Kaskade
#endif


// ************************  Prize-collecting Steiner Forest (PCSF) *********************//
// This is a C++ implementation of our heuristic algorithm for the PCSF (Akhmedov et al 2017).
// The heuristic is developed for functional analyses of large biological networks in a reasonable time.

// Loading the required libraries
#include <Rcpp.h>
#include <boost/config.hpp>
#include <boost/graph/adjacency_list.hpp>
#include <boost/graph/breadth_first_search.hpp>
#include <boost/graph/visitors.hpp>
#include <boost/graph/graph_utility.hpp>
#include <boost/random.hpp>
#include <boost/random/linear_congruential.hpp>
#include <boost/property_map/property_map.hpp>
#include <boost/graph/erdos_renyi_generator.hpp>
#include <boost/program_options.hpp>
#include <boost/format.hpp>
#include <boost/lexical_cast.hpp>
#include <boost/graph/connected_components.hpp>
#include <boost/graph/dijkstra_shortest_paths.hpp>
#include <boost/graph/prim_minimum_spanning_tree.hpp>
#include <time.h>
#include <stdio.h>
#include <boost/graph/graphviz.hpp>
#include <fstream>
#include <vector>
#include <utility>
#include <string>
#include <math.h>
#include <iomanip>
#include <boost/limits.hpp>
#include <queue>

using namespace boost;
using namespace Rcpp;
using namespace std;

// Node class
class Node
{  public:
  vector <int> children;
  int father;
  int size;
  double price;
  Node(){};
  Node(const Node & other){
    father=other.father;
    children= other.children;}

  Node & operator= (const Node & other){
    father=other.father;
    children= other.children;
    return *this;
  }

};

// Properties of Vertices in the network
struct VertexProperties  {
  VertexProperties() :c(0) {}
  VertexProperties(string const & name) : name(name),c(0){}
  string name;
  double c;
};

// Properties of the graph that is used within the BOOST
typedef adjacency_list<vecS, vecS, undirectedS, VertexProperties, property < edge_weight_t, double > > GraphBase;
typedef graph_traits<GraphBase>::vertex_iterator vertex_iterator;
typedef graph_traits<GraphBase>::out_edge_iterator edge_iterator;
typedef graph_traits<GraphBase>::edge_iterator graph_edge_iterator;
typedef graph_traits<GraphBase>::edge_descriptor Edge;
typedef graph_traits<GraphBase>::vertex_descriptor Vertex;

GraphBase g, g_adjusted; GraphBase G; GraphBase G_pruned;
property_map<GraphBase, edge_weight_t>::type weight_g;
property_map<GraphBase, edge_weight_t>::type weight_g_adjusted;
property_map<GraphBase, edge_weight_t>::type weight_G;
property_map<GraphBase, edge_weight_t>::type weight_G_pruned;

int Root = -1;
static map <string, int> g_map;
static map <string, int> G_map;
static map <string, int> G_pruned_map;

void clear_variables(){
  g.clear();
  g_adjusted.clear();
  G.clear();
  G_pruned.clear();
  g_map.clear();
  G_map.clear();
  G_pruned_map.clear();
}

// Map
int idx_g(String const & id)
{
  map<string, int>::iterator mit = g_map.find(id);
  if (mit == g_map.end())
    return g_map[id] = add_vertex(VertexProperties(id), g);
  return mit->second;
}
int idx_G(string const & id)
{
  map<string, int>::iterator mit = G_map.find(id);
  if (mit == G_map.end())
    return G_map[id] = add_vertex(VertexProperties(id), G);
  return mit->second;
}
int idx_G_pruned(string const & id)
{
  map<string, int>::iterator mit = G_pruned_map.find(id);
  if (mit == G_pruned_map.end())
    return G_pruned_map[id] = add_vertex(VertexProperties(id), G_pruned);
  return mit->second;
}

// Reading the input network
void read_input_graph(CharacterVector from, CharacterVector to,  NumericVector cost, CharacterVector prize, NumericVector prize_v)
{
  for(int i=0; i < from.size(); i++){
    add_edge(vertex(idx_g(from[i]), g), vertex(idx_g(to[i]), g), cost[i], g);
  }
  for(int i=0; i<prize.size(); i++){
    g[idx_g(prize[i])].c = prize_v[i];
  }

  // cerr << num_edges(g) << " edges, " << num_vertices(g) << " vertices" << endl;

  g_adjusted = g;
  weight_g_adjusted = get(edge_weight, g_adjusted);
  graph_edge_iterator ei, ei_end; double penalty; Vertex sour, tar;
  for(tie(ei, ei_end) = edges(g_adjusted); ei != ei_end; ++ei){
    sour = source(*ei,g_adjusted); tar = target(*ei,g_adjusted); penalty = 0;
    if(g_adjusted[sour].c < 0 && g_adjusted[tar].c < 0){
      penalty = g_adjusted[sour].c + g_adjusted[tar].c;
    } else if( g_adjusted[sour].c < 0 ){
      penalty = g_adjusted[sour].c;
    } else if( g_adjusted[tar].c < 0 ){
      penalty = g_adjusted[tar].c;
    }
    weight_g_adjusted[*ei] = weight_g_adjusted[*ei] + abs(penalty);
  }
}

// A function to dynamically remove the leaf node if its prize smaller than connection cost.
// It is used called within the process_leafs() function.
void clear(vector <Node> & predecessor, int & current_node){
  int node=current_node;
  for(unsigned int j=0; j<predecessor[node].children.size(); j++){
    clear(predecessor, predecessor[node].children[j]);
  }
  predecessor[node].father=-1;
}

// A function to dynamically remove the leaf node if its prize smaller than connection cost.
void process_leafs(vector <Node> & predecessor, int & current_node, Edge &e, bool &found){
  for(unsigned int i=0; i<predecessor[current_node].children.size(); i++){
    process_leafs(predecessor, predecessor[current_node].children[i], e, found);
  }
  int node = current_node;
  if(node != predecessor[node].father){
    boost::tuples::tie(e,found) = edge( vertex(predecessor[node].father, G_pruned),vertex(node, G_pruned) , G_pruned);
    if(predecessor[node].price - weight_G_pruned[e] <= 0) {
      clear(predecessor, node);
    }
  }

}

// A function to dynamically sum up the prizes of vertices.
void price_collect(vector <Node> & predecessor, int & current_node, Edge &e, bool &found){
  for(unsigned int i=0; i<predecessor[current_node].children.size(); i++){
    price_collect(predecessor, predecessor[current_node].children[i], e, found);
  }
  int node = current_node;
  if(node != predecessor[node].father){
    boost::tuples::tie(e,found) = edge(vertex(predecessor[node].father, G_pruned),vertex(node, G_pruned) , G_pruned);
    if(predecessor[node].price - weight_G_pruned[e] > 0)
      predecessor[predecessor[node].father].price = predecessor[predecessor[node].father].price + predecessor[node].price - weight_G_pruned[e];
  }
}

// After reading the input network information from the input file, the algorithm constructs a
vector< Vertex > constructG(vector<int> & terminals, int &Root){

  // Distance: all-pairs-shortest-path distance matrix
  // perPath: List of arcs in all-pairs-shortest-path distance matrix
  vector <vector <vector<int> > > perPath;
  vector <vector<double> > Distance;
  perPath.resize (terminals.size());
  Distance.resize (terminals.size());
  for (unsigned int i = 0; i < terminals.size(); ++i) {
    perPath [i].resize(terminals.size());
    Distance [i].resize(terminals.size());
  }

  // Computing all-pairs-shortest-path distance matrix
  Vertex from; int current, pred, outer = 0, inner;
  std::vector<Vertex> p = vector<Vertex> (num_vertices(g));
  std::vector<double> d = vector<double> (num_vertices(g));

  for (std::vector<int>::iterator first=terminals.begin(); first!=terminals.end(); ++first){

    from = vertex(*first, g_adjusted);
    dijkstra_shortest_paths(g_adjusted, from,
                            predecessor_map(boost::make_iterator_property_map(p.begin(), get(boost::vertex_index, g_adjusted))).
                              distance_map(boost::make_iterator_property_map(d.begin(), get(boost::vertex_index, g_adjusted))));

    inner = outer + 1;
    for (std::vector<int>::iterator second = first+1; second!=terminals.end(); ++second) {
      current=*second; pred=p[current];
      Distance[outer][inner] = d[current];
      Distance[inner][outer] = d[current];
      while(pred!=current){perPath[outer][inner].push_back(current); current=pred; pred=p[current]; } perPath[outer][inner].push_back(current);
      inner++;
    }
    outer++;
  }

  // Heuristic Clustering, given large input network, the algorithm clusters input network into
  // smaller clusters, and solves the MST afterwards
  set<int> V;
  set<int> D;
  unsigned int root_index = -1;
  vector<int> node_labels;
  node_labels.resize(terminals.size());
  for(unsigned int i=0; i< node_labels.size(); i++){
    node_labels[i] = 0;}
  for(unsigned int i=0; i<terminals.size(); i++){
    if(terminals[i] == Root){
      root_index = i;}
    else{
      V.insert(i);}
  }
  node_labels[root_index] = INT_MAX;
  int clusterID=0; int targ;

  while(!V.empty()){
    clusterID++;  current = *V.begin();
    node_labels[current] =clusterID;
    V.erase(current);

    D.clear();
    for (unsigned int i=0; i < terminals.size(); i++) {
      targ = i;
      if(node_labels[targ] == 0 && current != targ && i != root_index){
        if(g[terminals[current]].c >= Distance[current][targ] && g[terminals[targ]].c >= Distance[current][targ]){
          if(g[terminals[targ]].c > 0){
            D.insert(targ); V.erase(targ);
          }
          node_labels[targ]=clusterID;
        }
      }
    }

    while(!D.empty()){
      current = *D.begin();  D.erase(current);
      for (unsigned int i=0; i < terminals.size(); i++) {
        targ = i;
        if(node_labels[targ] == 0 && current != targ && i != root_index){
          if(g[terminals[current]].c >= Distance[current][targ] && g[terminals[targ]].c >= Distance[current][targ]){
            if(g[terminals[targ]].c > 0){ D.insert(targ); V.erase(targ);}
            node_labels[targ]=clusterID;
          }
        }
      }
    }

  }

  // Identfying the vertex membership with respect to clusters
  vector<vector <int> > clusters(clusterID+1);
  for(unsigned int i=0; i< node_labels.size(); i++){
    for(int j=0; j<=clusterID; j++){
      if(node_labels[i] == j){
        clusters[j].push_back(i);
      }
    }
  }

  std::vector<int>::iterator it, itt; int num_clusters=0;
  for(unsigned int i=1; i<clusters.size(); i++){
    if(clusters[i].size() > 1) num_clusters++;
  }

  if(num_clusters == 0){
    //cout<<"There is no tree in construct G ()"<<endl;
    //return 0;
  }

  // Regrouping the singletone and dobletone clusters after clustering
  unsigned int threshold_num = 2; int min_index; double min_distance;
  for(unsigned int i=1; i<clusters.size(); i++){
    if( clusters[i].size() <= threshold_num){
      for (it=clusters[i].begin(); it!=clusters[i].end(); ++it){
        min_index = -1; min_distance = DBL_MAX;
        for(unsigned int j=1; j<clusters.size(); j++){
          if( clusters[j].size() > threshold_num){
            for (itt=clusters[j].begin(); itt!=clusters[j].end(); ++itt){
              if(min_distance > Distance[*it][*itt] -g[terminals[*it]].c - g[terminals[*itt]].c ){
                min_distance = Distance[*it][*itt] -g[terminals[*it]].c - g[terminals[*itt]].c; min_index = j;
              }
            }
          }
        }

        if (min_index != -1){
          clusters[min_index].push_back(*it);
          *it = -1;
        }
      }
    }

  }

  // Construct an artificial graph G, which is composed of all clusters determined
  // from Heuristic Clustering phase

  string str; unsigned int index1=-1, index2=-1;

  for (unsigned int l = 0; l<terminals.size(); l++){
    str=to_string(l);
    index1=idx_G(str);
    G[index1].c = g[terminals[l]].c;
    G[index1].name = g[terminals[l]].name;
  }

  for (unsigned int l = 0; l<terminals.size(); l++){
    if(l != root_index){
      str=to_string(l);
      index1=idx_G(str);
      add_edge(root_index, index1, Distance[root_index][index1], G);
    }
  }

  for(unsigned int i = 1; i < clusters.size(); i++){
    for (it=clusters[i].begin(); it!=clusters[i].end(); ++it){
      if(*it != -1){
        str=to_string(*it); index1=idx_G(str);
        for (itt=it+1; itt!=clusters[i].end(); ++itt){
          if(*itt != -1){
            str=to_string(*itt); index2=idx_G(str);
            add_edge(index1, index2, Distance[*it][*itt], G);
          }
        }
      }
    }
  }

  weight_G = get(edge_weight, G);
  vector < Vertex > spanning_tree_G(num_vertices(G));
  prim_minimum_spanning_tree(G, & spanning_tree_G[0]);

  Edge beg; Vertex sour, tar; double cost;
  Edge e; bool found;

  weight_g = get(edge_weight, g);

  vector<int> path; index1=0; index2=0;
  edge_iterator out_i, out_end; int add=0;

  // Solving the Minimum Spanning Tree on G
  for(unsigned int i = 0; i < spanning_tree_G.size(); ++i ){

    if(spanning_tree_G[i]!=i ){

      if(i> spanning_tree_G[i]) path=perPath[spanning_tree_G[i]][i];
      else path=perPath[i][spanning_tree_G[i]];

      for(unsigned int j=0; j<path.size()-1; j++){

        sour=vertex(path[j], g); tar= vertex(path[j+1], g);
        boost::tuples::tie(beg, found) = edge(sour, tar,g);
        cost=get(weight_g, beg);

        index1=idx_G_pruned(to_string(path[j])); index2=idx_G_pruned(to_string(path[j+1]));
        add=0;
        for (boost::tuples::tie(out_i, out_end) = out_edges(vertex(index1,G_pruned), G_pruned); out_i != out_end; ++out_i) {
          if(target(*out_i, G_pruned)==index2) add++;
        }
        if(!add) add_edge(index1, index2, cost, G_pruned);
      }
  }
  }

  weight_G_pruned = get(edge_weight, G_pruned);

  vector< Vertex >spanning_tree_G_pruned(num_vertices(G_pruned));
  prim_minimum_spanning_tree(G_pruned, &spanning_tree_G_pruned[0]);

  double total1=0;
  for(unsigned int i = 0; i < spanning_tree_G_pruned.size(); ++i ){
    if(spanning_tree_G_pruned[i] != i){
      sour= vertex(i,G_pruned); tar=vertex(spanning_tree_G_pruned[i], G_pruned);
      boost::tuples::tie(beg, found) = edge(sour, tar,G_pruned);
      total1+=get(weight_G_pruned, beg);
    }
  }

  return spanning_tree_G_pruned;

  }

// After obtaining MST tree, the algorithm prunes the leaf nodes
// which have prizes smaller than connection cost
double cut(int &Root, vector< Vertex > &spanning_tree_G_pruned,  vector< string > &tree_from,  vector< string > &tree_to,  vector< double > &tree_cost, map < string, double > &tree_terminals){

  weight_G_pruned = get(edge_weight, G_pruned);

  Edge e; bool found;

  Edge beg; Vertex sour, tar; int ancestor=-1; bool select=false;

  int root= -1;
  if(Root == -1){
    double max=0;
    vertex_iterator ei, ef;
    for(tie(ei, ef)= vertices(G_pruned); ei!=ef; ei++){
      if(g[boost::lexical_cast<int>(G_pruned[*ei].name)].c > max){
        root = *ei;
        max = g[boost::lexical_cast<int>(G_pruned[*ei].name)].c;
      }
    }
  } else {root = idx_G_pruned(to_string(Root));}

  select = true;

  bool ancestor_changed=true; unsigned int father, temp; ancestor = root;
  if (select){
    father=spanning_tree_G_pruned[ancestor];
    spanning_tree_G_pruned[ancestor]=ancestor;
    while(ancestor_changed){
      if(spanning_tree_G_pruned[father]==father){
        ancestor_changed=false;
        spanning_tree_G_pruned[father]=ancestor;
      }else{
        temp=spanning_tree_G_pruned[father];
        spanning_tree_G_pruned[father]=ancestor;
        ancestor=father; father=temp;
      }

    }

  }else{
    //cout <<"There is no tree"<<endl;
    return 0;
  }

  vector<Node> predecessor(num_vertices(G_pruned));
  if(select){
    for(unsigned int i = 0; i < spanning_tree_G_pruned.size(); ++i ){
      if(spanning_tree_G_pruned[i]!=i){
        predecessor[i].father=spanning_tree_G_pruned[i];
        predecessor[spanning_tree_G_pruned[i]].children.push_back(i);
      }else{predecessor[i].father=i;}
    }
  }

  for(unsigned int i = 0; i < predecessor.size(); ++i ){
    predecessor[i].size=predecessor[i].children.size();
    predecessor[i].price=g[boost::lexical_cast<int>(G_pruned[i].name)].c;

  }

  price_collect(predecessor, root, e, found);

  process_leafs(predecessor, root, e, found);

  weight_g = get(edge_weight, g);

  // Tree
  for(unsigned int i = 0; i < predecessor.size(); ++i ){
    if(predecessor[i].father != -1 &&  predecessor[i].father != (int) i ){
      sour= vertex(boost::lexical_cast<int>(G_pruned[i].name),g); tar = vertex(boost::lexical_cast<int>(G_pruned[predecessor[i].father].name),g);
      boost::tuples::tie(beg, found) = edge(sour, tar,g);
      tree_from.push_back(g[boost::lexical_cast<int>(G_pruned[i].name)].name);
      tree_to.push_back(g[boost::lexical_cast<int>(G_pruned[predecessor[i].father].name)].name);
      tree_cost.push_back(weight_g[beg]);
      tree_terminals[g[boost::lexical_cast<int>(G_pruned[i].name)].name] = g[boost::lexical_cast<int>(G_pruned[i].name)].c;
      tree_terminals[g[boost::lexical_cast<int>(G_pruned[predecessor[i].father].name)].name] = g[boost::lexical_cast<int>(G_pruned[predecessor[i].father].name)].c;
    }
  }

  double total = 0, lostPrice =0;
  int uncovered_nodes = 0;
  for(unsigned int i = 0; i < predecessor.size(); ++i ){
    if(predecessor[i].father != -1){
      if(predecessor[i].father != (int) i){
        sour= vertex(i,G_pruned); tar = vertex(predecessor[i].father,G_pruned);
        boost::tuples::tie(beg, found) = edge(sour, tar,G_pruned);
        total+=get(weight_G_pruned, beg);
      }
    }
  }

  // Lsit of nodes that are outside of final tree
  vector<int> calculatecost(num_vertices(g));
  for(unsigned int i = 0; i < predecessor.size(); ++i ){
    if(predecessor[i].father != -1 && predecessor[i].father != (int) i){
      sour= vertex(boost::lexical_cast<int>(G_pruned[i].name),g); tar = vertex(boost::lexical_cast<int>(G_pruned[predecessor[i].father].name),g);
      calculatecost[sour]=1; calculatecost[tar]=1;
    }
  }

  // Uncovered nodes
  for(unsigned int i = 0; i < num_vertices(g); ++i ){
    if(calculatecost[i] == 0 && (int) i != root ){
      lostPrice += g[i].c;
      uncovered_nodes++;
    }
  }

  // The list of Nodes in the final Tree
  for(unsigned int i = 0; i < num_vertices(g); ++i ){
    if(calculatecost[i] == 1){
    }
  }

  // Objective value
  return total + lostPrice;

}

//' Internal function \code{call_sr}
//'
//' This function is internally used to solve the PCST.
//'
//' @keywords internal
//'
//' @param from  A \code{CharacterVector} that corresponds to \code{head} nodes of the edges.
//' @param to A \code{CharacterVector} that corresponds the \code{tail} nodes of the edges.
//' @param cost A \code{NumericVector} which represents the edge weights.
//' @param node_names A \code{CharacterVector} demonstrates the names of the nodes.
//' @param node_prizes A \code{NumericVector} which corresponds to the node prizes.
//' @author Murodzhon Akhmedov
//'
// [[Rcpp::export]]
List call_sr(CharacterVector from, CharacterVector to,  NumericVector cost, CharacterVector node_names, NumericVector node_prizes)
{
  clear_variables();

  vector <int> terminals;

  read_input_graph(from, to, cost, node_names, node_prizes);

  Root = idx_g("DUMMY");

  double max_price=0; int max_price_index = -1;
  for(unsigned int i=0; i<num_vertices(g); i++){
    if (g[i].c > max_price){
      max_price = g[i].c;
      max_price_index = i;
    }
  }

  if(Root != -1){
    for(unsigned int i=0; i<num_vertices(g); i++){
      if( (int) i != Root && g[i].c >0){
        terminals.push_back(i);
      }
    }
  }else{
    Root = max_price_index;
    for(unsigned int i=0; i<num_vertices(g); i++){
      if( (int) i != Root && g[i].c >0){
        terminals.push_back(i);
      }
    }
  }

  terminals.push_back(Root);

  if(terminals.size() <=1){
    // There is no tree
    //return 0;
  }

  vector< Vertex > spanning_tree;
  spanning_tree = constructG(terminals, Root);

  vector< string > tree_from;
  vector< string > tree_to;
  vector< double > tree_cost;
  map < string, double > tree_terminals;
  double obj = cut(Root, spanning_tree, tree_from, tree_to, tree_cost, tree_terminals); if (obj == 0) return 0;

  CharacterVector tree_f(tree_from.size());
  CharacterVector tree_t(tree_to.size());
  NumericVector tree_c(tree_cost.size());
  CharacterVector tree_ter(tree_terminals.size());
  NumericVector tree_ter_p(tree_terminals.size());

  for(unsigned int i=0; i<tree_from.size(); i++){
    tree_f[i]=tree_from[i];
    tree_t[i]=tree_to[i];
    tree_c[i]=tree_cost[i];
  }

  int counter = 0;
  for (std::map<string, double>::iterator it=tree_terminals.begin(); it!=tree_terminals.end(); ++it){
    tree_ter[counter] = it->first;
    tree_ter_p[counter] = it->second;
    counter++;
  }

  List tree = List::create(tree_from, tree_to, tree_cost, tree_ter, tree_ter_p);

  return tree;
}


/*    Copyright (c) 2010-2019, Delft University of Technology
*    All rigths reserved
*
*    This file is part of the Tudat. Redistribution and use in source and
*    binary forms, with or without modification, are permitted exclusively
*    under the terms of the Modified BSD license. You should have received
*    a copy of the license with this file. If not, please or visit:
*    http://tudat.tudelft.nl/LICENSE.
*/

#include <iostream>
#include <fstream>

#include <boost/filesystem.hpp>

#include "pagmo/island.hpp"
#include "pagmo/io.hpp"
#include "pagmo/problem.hpp"
#include <pagmo/rng.hpp>
#include "pagmo/algorithms/sade.hpp"
#include "problems/multipleGravityAssist.h"
#include "tudat/astro/ephemerides/approximatePlanetPositions.h"
#include "tudat/astro/ephemerides/constantEphemeris.h"
#include "tudat/astro/gravitation/gravityFieldModel.h"
#include "tudat/astro/mission_segments/createTransferTrajectory.h"
#include "tudat/simulation/environment/body.h"

using namespace pagmo;
using namespace tudat;
using namespace ephemerides;
using namespace gravitation;
using namespace simulation_setup;

simulation_setup::NamedBodyMap getApproximatePlanetBodyMap( )
{

    NamedBodyMap bodyMap;
    bodyMap[ "Sun" ] = std::make_shared< Body >( );
    bodyMap[ "Mercury" ] = std::make_shared< Body >( );
    bodyMap[ "Venus" ] = std::make_shared< Body >( );
    bodyMap[ "Earth" ] = std::make_shared< Body >( );
    bodyMap[ "Mars" ] = std::make_shared< Body >( );
    bodyMap[ "Jupiter" ] = std::make_shared< Body >( );
    bodyMap[ "Saturn" ] = std::make_shared< Body >( );

    bodyMap[ "Sun" ]->setEphemeris( std::make_shared< ConstantEphemeris >( Eigen::Vector6d::Zero( ) ) );
    bodyMap[ "Mercury" ]->setEphemeris( std::make_shared< ApproximatePlanetPositions >(
                                            ApproximatePlanetPositionsBase::BodiesWithEphemerisData::mercury ) );
    bodyMap[ "Venus" ]->setEphemeris( std::make_shared< ApproximatePlanetPositions >(
                                          ApproximatePlanetPositionsBase::BodiesWithEphemerisData::venus ) );
    bodyMap[ "Earth" ]->setEphemeris( std::make_shared< ApproximatePlanetPositions >(
                                          ApproximatePlanetPositionsBase::BodiesWithEphemerisData::earthMoonBarycenter ) );
    bodyMap[ "Mars" ]->setEphemeris( std::make_shared< ApproximatePlanetPositions >(
                                          ApproximatePlanetPositionsBase::BodiesWithEphemerisData::earthMoonBarycenter ) );
    bodyMap[ "Jupiter" ]->setEphemeris( std::make_shared< ApproximatePlanetPositions >(
                                            ApproximatePlanetPositionsBase::BodiesWithEphemerisData::jupiter ) );
    bodyMap[ "Saturn" ]->setEphemeris( std::make_shared< ApproximatePlanetPositions >(
                                           ApproximatePlanetPositionsBase::BodiesWithEphemerisData::saturn ) );

    bodyMap[ "Sun" ]->setGravityFieldModel( std::make_shared< GravityFieldModel >( 1.32712428e20 ) );
    bodyMap[ "Mercury" ]->setGravityFieldModel( std::make_shared< GravityFieldModel >( 2.2321e13 ) );
    bodyMap[ "Venus" ]->setGravityFieldModel( std::make_shared< GravityFieldModel >( 3.24860e14 ) );
    bodyMap[ "Earth" ]->setGravityFieldModel( std::make_shared< GravityFieldModel >( 3.9860119e14 ) );
    bodyMap[ "Mars" ]->setGravityFieldModel( std::make_shared< GravityFieldModel >( 4.282837e13 ) );
    bodyMap[ "Jupiter" ]->setGravityFieldModel( std::make_shared< GravityFieldModel >( 1.267e17 ) );
    bodyMap[ "Saturn" ]->setGravityFieldModel( std::make_shared< GravityFieldModel >( 3.79e16 ) );

    return bodyMap;
}

namespace std
{
template <typename T>
ostream& operator<<(ostream& output, std::vector<T> const& values)
{
    std::cout<<"[ ";
    for (auto const& value : values)
    {
        output << value <<", ";
    }
    std::cout<<" ]"<<std::endl;
    return output;
}
}
//! Execute  main
int main( )
{
    //Set seed for reproducible results
    pagmo::random_device::set_seed( 123456789 );

    // Set transfer order
    std::vector< std::string > bodyOrder = {
        "Earth", "Venus", "Venus", "Earth", "Jupiter", "Saturn" };
    int numberOfNodes = bodyOrder.size( );

    // Create leg settings (all unpowered)
    std::vector< std::shared_ptr< TransferLegSettings > > transferLegSettings;
    transferLegSettings.resize( numberOfNodes - 1 );
    transferLegSettings[ 0 ] = unpoweredLeg( );
    transferLegSettings[ 1 ] = dsmVelocityBasedLeg( );
    transferLegSettings[ 2 ] = dsmVelocityBasedLeg( );
    transferLegSettings[ 3 ] = unpoweredLeg( );
    transferLegSettings[ 4 ] = unpoweredLeg( );

    // Define minimum periapsis altitudes for flybys;
    std::map< std::string, double >minimumPeriapses;
    minimumPeriapses[ "Venus" ] = 6351800.0;
    minimumPeriapses[ "Earth" ] = 6778000.0;
    minimumPeriapses[ "Mars" ] = 3696200.0;
    minimumPeriapses[ "Jupiter" ] =  600000000;

    std::vector< std::shared_ptr< TransferNodeSettings > > transferNodeSettings;
    transferNodeSettings.resize( numberOfNodes );
    transferNodeSettings[ 0 ] = escapeAndCaptureNode( std::numeric_limits< double >::infinity( ), 0.0 );
    transferNodeSettings[ 1 ] = swingbyNode( minimumPeriapses.at( bodyOrder.at( 1 ) ) );
    transferNodeSettings[ 2 ] = swingbyNode( minimumPeriapses.at( bodyOrder.at( 2 ) ) );
    transferNodeSettings[ 3 ] = swingbyNode( minimumPeriapses.at( bodyOrder.at( 3 ) ) );
    transferNodeSettings[ 4 ] = swingbyNode( minimumPeriapses.at( bodyOrder.at( 4 ) ) );
    transferNodeSettings[ 5 ] = captureAndInsertionNode( 1.0895e8 / 0.02, 0.98 );

    printTransferParameterDefinition( transferLegSettings, transferNodeSettings );
    simulation_setup::NamedBodyMap bodyMap = getApproximatePlanetBodyMap( );

    // Define search bounds: first parameter is start date, following parameters are leg durations
    std::vector< std::vector< double > > bounds( 2, std::vector< double >( 14, 0.0 ) );
    bounds[ 0 ][ 0 ] = -2000.0; //MJD2000
    bounds[ 1 ][ 0 ] = 0.0; //MJD2000
    bounds[ 0 ][ 1 ] = 50.0;
    bounds[ 1 ][ 1 ] = 500.0;
    bounds[ 0 ][ 2 ] = 100.0;
    bounds[ 1 ][ 2 ] = 500.0;
    bounds[ 0 ][ 3 ] = 50.0;
    bounds[ 1 ][ 3 ] = 500.0;
    bounds[ 0 ][ 4 ] = 500.0;
    bounds[ 1 ][ 4 ] = 2000.0;
    bounds[ 0 ][ 5 ] = 1000.0;
    bounds[ 1 ][ 5 ] = 10000.0;
//    bounds[ 0 ][ 6 ] = 0.0;
//    bounds[ 1 ][ 6 ] = 2.0E3;;
//    bounds[ 0 ][ 7 ] = 0.0;
//    bounds[ 1 ][ 7 ] = 2.0 * mathematical_constants::PI;

//    bounds[ 0 ][ 8 ] = 0.0;
//    bounds[ 1 ][ 8 ] = 2.0 * mathematical_constants::PI;

//    bounds[ 0 ][ 9 ] = 0.05;
//    bounds[ 1 ][ 9 ] = 0.95;
    int currentIndex = 5;
    for( int i = 1; i < 3; i++ )
    {
        // periapsis bounds
        bounds[ 0 ][ currentIndex + ( i - 1 ) * 4 + 1 ] = minimumPeriapses.at( bodyOrder.at( i ) );
        bounds[ 1 ][ currentIndex + ( i - 1 ) * 4 + 1 ] = 1.5 * minimumPeriapses.at( bodyOrder.at( i ) );

        // orbit orientation bounds
        bounds[ 0 ][ currentIndex + ( i - 1 ) * 4 + 2 ] = 0.0;
        bounds[ 1 ][ currentIndex + ( i - 1 ) * 4 + 2 ] = 2.0 * mathematical_constants::PI;

        // Swingby Delta V bounds
        bounds[ 0 ][ currentIndex + ( i - 1 ) * 4 + 3 ] = 0;
        bounds[ 1 ][ currentIndex + ( i - 1 ) * 4 + 3 ] = 2500.0;

        // DSM TOF fraction
        bounds[ 0 ][ currentIndex + ( i - 1 ) * 4 + 4 ] = 0.05;
        bounds[ 1 ][ currentIndex + ( i - 1 ) * 4 + 4 ] = 0.95;
    }

    // Create object to compute the problem fitness
    problem prob{ MultipleGravityAssist(
                    bodyMap, transferLegSettings, transferNodeSettings, bodyOrder, "Sun", bounds ) };

    // Select NSGA2 algorithm for priblem
    algorithm algo{sade( )};

    // Create an island with 1000 individuals
    island isl{algo, prob, 500 };

    // Evolve for 512 generations
    for( int i = 0 ; i < 10000; i++ )
    {

        isl.evolve( );
        while( isl.status( ) != pagmo::evolve_status::idle &&
               isl.status( ) != pagmo::evolve_status::idle_error )
        {
            isl.wait( );
        }

        isl.wait_check( ); // Raises errors

        if( i% 100 == 0 )
        {
            std::cout<<i<<" "<<isl.get_population().champion_f()[0]<<"***  "<<
                       isl.get_population().champion_x()<<std::endl;        // Write current iteration results to file
        }
//        printPopulationToFile( isl.get_population( ).get_x( ), "mo_mga_EVEEJ_" + std::to_string( i ), false );
//        printPopulationToFile( isl.get_population( ).get_f( ), "mo_mga_EVEEJ_" + std::to_string( i ), true );
    }

    return 0;

}


/**************************************************************************
 * Copyright (c) 2017-2019 by the mfmg authors                            *
 * All rights reserved.                                                   *
 *                                                                        *
 * This file is part of the mfmg library. mfmg is distributed under a BSD *
 * 3-clause license. For the licensing terms see the LICENSE file in the  *
 * top-level directory                                                    *
 *                                                                        *
 * SPDX-License-Identifier: BSD-3-Clause                                  *
 *************************************************************************/

#ifndef AMGE_HOST_TEMPLATES_HPP
#define AMGE_HOST_TEMPLATES_HPP

#include <mfmg/common/lanczos.templates.hpp>
#include <mfmg/common/utils.hpp>
#include <mfmg/dealii/amge_host.hpp>
#include <mfmg/dealii/anasazi.templates.hpp>
#include <mfmg/dealii/dealii_matrix_free_mesh_evaluator.hpp>

#include <deal.II/base/work_stream.h>
#include <deal.II/dofs/dof_accessor.h>
#include <deal.II/lac/arpack_solver.h>
#include <deal.II/lac/la_parallel_vector.h>
#include <deal.II/lac/lapack_full_matrix.h>
#include <deal.II/lac/precondition.h>
#include <deal.II/lac/solver_cg.h>
#include <deal.II/lac/sparse_direct.h>

#include <EpetraExt_MatrixMatrix.h>

namespace mfmg
{

struct Identity
{
  template <typename VectorType>
  void vmult(VectorType &dst, const VectorType &src) const
  {
    dst = src;
  }
};

/**
 * A class representing the operator on a given agglomerate
 * wrapping a MeshEvaluator object to restrict the interface
 * to the operations required on an agglomerate.
 */
template <typename MeshEvaluator>
struct MatrixFreeAgglomerateOperator
{
  using size_type = typename MeshEvaluator::size_type;

  /**
   * This constructor expects @p mesh_evaluator to be an object that performs
   * the actual operator evaluation. @p dof_handler is used to initialize the
   * appropriate data structures in the MeshEvaluator and is expected to be
   * initialized itself with respect to a given dealii::FiniteElement object
   * in that call.
   */
  template <typename DoFHandler>
  MatrixFreeAgglomerateOperator(MeshEvaluator const &mesh_evaluator,
                                DoFHandler &dof_handler,
                                dealii::AffineConstraints<double> &constraints)
      : _mesh_evaluator(mesh_evaluator.clone()), _dof_handler(dof_handler),
        _constraints(constraints)
  {
    _mesh_evaluator->matrix_free_initialize_agglomerate(dof_handler);
  }

  /**
   * Perform the operator evaluation on the agglomerate.
   */
  void vmult(dealii::Vector<double> &dst,
             dealii::Vector<double> const &src) const
  {
    _mesh_evaluator->matrix_free_evaluate_agglomerate(src, dst);
  }

  /**
   * Return the diagonal entries the matrix corresponding to the operator would
   * have. This data is necessary for certain smoothers to work.
   * @p _constraints is used to restrict the diagonal to the correct
   * (constrained) finite element subspace.
   */
  std::vector<double> get_diag_elements() const
  {
    return _mesh_evaluator->matrix_free_get_agglomerate_diagonal(_constraints);
  }

  /**
   * Return the dimension of the range of the agglomerate operator.
   */
  size_type m() const { return _dof_handler.n_dofs(); }

  /**
   * Return the dimension of the domain of the agglomerate operator.
   */
  size_type n() const { return _dof_handler.n_dofs(); }

private:
  /**
   * The actual operator wrapped.
   */
  std::unique_ptr<MeshEvaluator const> const _mesh_evaluator;

  /**
   * The dimension for the underlying mesh.
   */
  static int constexpr dim = MeshEvaluator::_dim;

  /**
   * The DoFHandler containing information about the degrees of freedom on the
   * agglomerate.
   */
  dealii::DoFHandler<dim> const &_dof_handler;

  /**
   * The constraints needed for restricting the vector returned by
   * get_diag_elements() to the correct subspace.
   */
  dealii::AffineConstraints<double> &_constraints;
};

template <int dim, typename MeshEvaluator, typename VectorType>
AMGe_host<dim, MeshEvaluator, VectorType>::AMGe_host(
    MPI_Comm comm, dealii::DoFHandler<dim> const &dof_handler,
    boost::property_tree::ptree const &eigensolver_params)
    : AMGe<dim, VectorType>(comm, dof_handler),
      _eigensolver_params(eigensolver_params)
{
}
} // namespace mfmg

namespace Anasazi
{
template <typename VectorType, typename MeshEvaluator>
class OperatorTraits<double, mfmg::MultiVector<VectorType>,
                     mfmg::MatrixFreeAgglomerateOperator<MeshEvaluator>>
{
  using MultiVectorType = mfmg::MultiVector<VectorType>;
  using OperatorType = mfmg::MatrixFreeAgglomerateOperator<MeshEvaluator>;

public:
  static void Apply(OperatorType const &op, MultiVectorType const &x,
                    MultiVectorType &y)
  {
    auto n_vectors = x.n_vectors();

    ASSERT(x.size() == y.size(), "");
    ASSERT(y.n_vectors() == n_vectors, "");

    for (int i = 0; i < n_vectors; i++)
      op.vmult(*y[i], *x[i]);
  }
};

} // namespace Anasazi

namespace mfmg
{
namespace
{
template <typename AgglomerateOperator>
void lanczos_compute_eigenvalues_and_eigenvectors(
    unsigned int n_eigenvectors, double tolerance,
    boost::property_tree::ptree const &eigensolver_params,
    AgglomerateOperator const &agglomerate_operator,
    dealii::Vector<double> const &initial_guess,
    std::vector<std::complex<double>> &eigenvalues,
    std::vector<dealii::Vector<double>> &eigenvectors)
{
  boost::property_tree::ptree lanczos_params;
  lanczos_params.put("num_eigenpairs", n_eigenvectors);
  // We are having trouble with Lanczos when tolerance is too tight.
  // Typically, it results in spurious eigenvalues (so far, only noticed 0).
  // This seems to be the result of producing too many Lanczos vectors. For
  // hierarchy_3d tests, any tolerance below 1e-5 (e.g., 1e-6) produces this
  // problem. Thus, we try to work around it here. It is still unclear how
  // robust this is.
  lanczos_params.put("tolerance", std::max(tolerance, 1e-4));
  lanczos_params.put("max_iterations",
                     eigensolver_params.get("max_iterations", 200));
  lanczos_params.put("percent_overshoot",
                     eigensolver_params.get("percent_overshoot", 5));
  bool is_deflated = eigensolver_params.get("is_deflated", false);
  if (is_deflated)
  {
    lanczos_params.put("is_deflated", true);
    lanczos_params.put("num_cycles", eigensolver_params.get<int>("num_cycles"));
    lanczos_params.put("num_eigenpairs_per_cycle",
                       eigensolver_params.get<int>("num_eigenpairs_per_cycle"));
  }

  Lanczos<AgglomerateOperator, dealii::Vector<double>> solver(
      agglomerate_operator);

  std::vector<double> real_eigenvalues;
  std::tie(real_eigenvalues, eigenvectors) =
      solver.solve(lanczos_params, initial_guess);
  ASSERT(n_eigenvectors == eigenvectors.size(),
         "Wrong number of computed eigenpairs");

  // Copy real eigenvalues to complex
  std::copy(real_eigenvalues.begin(), real_eigenvalues.end(),
            eigenvalues.begin());
}

template <typename AgglomerateOperator>
void anasazi_compute_eigenvalues_and_eigenvectors(
    unsigned int n_eigenvectors,
    boost::property_tree::ptree const &eigensolver_params,
    AgglomerateOperator const &agglomerate_operator,
    dealii::Vector<double> const &initial_guess,
    std::vector<dealii::Vector<double>> const &lobpcg_vectors,
    std::vector<std::complex<double>> &eigenvalues,
    std::vector<dealii::Vector<double>> &eigenvectors)
{
  AnasaziSolver<AgglomerateOperator, dealii::Vector<double>> solver(
      agglomerate_operator);

  std::vector<double> real_eigenvalues;
  std::vector<std::shared_ptr<dealii::Vector<double>>> lobpcg_initial_guess;
  // If the vectors in scratch_data do not exist or if the size of agglomerate
  // has changed, the initial guess for LOBPCG is the initial provided by the
  // user.
  if ((lobpcg_vectors.size() == 0) ||
      (lobpcg_vectors[0].size() != initial_guess.size()))
  {
    lobpcg_initial_guess.resize(1);
    lobpcg_initial_guess[0] =
        std::make_shared<dealii::Vector<double>>(initial_guess);
  }
  else
  {
    lobpcg_initial_guess.resize(n_eigenvectors);
    for (unsigned int i = 0; i < n_eigenvectors; ++i)
      lobpcg_initial_guess[i] =
          std::make_shared<dealii::Vector<double>>(lobpcg_vectors[i]);

    // If the initial vector has zero entries due to the constraints, we need
    // to modify the LOPBCG initial guess. Conversely if the LOBPCG initial
    // guess does have zero entries but the initial vector does not, we set
    // the entries in the LOPBCG initial guess.
    unsigned int const eigenvector_size = initial_guess.size();
    for (unsigned int i = 0; i < eigenvector_size; ++i)
    {
      if (initial_guess[i] == 0.)
      {
        for (unsigned int j = 0; j < n_eigenvectors; ++j)
        {
          (*lobpcg_initial_guess[j])[i] = 0.;
        }
      }
      else
      {
        for (unsigned int j = 0; j < n_eigenvectors; ++j)
        {
          if ((*lobpcg_initial_guess[j])[i] == 0.)
          {
            (*lobpcg_initial_guess[j])[i] = initial_guess[i];
          }
        }
      }
    }
  }
  std::tie(real_eigenvalues, eigenvectors) =
      solver.solve(eigensolver_params, lobpcg_initial_guess);
  ASSERT(n_eigenvectors == eigenvectors.size(),
         "Wrong number of computed eigenpairs");

  // Copy real eigenvalues to complex
  std::copy(real_eigenvalues.begin(), real_eigenvalues.end(),
            eigenvalues.begin());
}
} // namespace

template <int dim, typename MeshEvaluator, typename VectorType>
template <typename Triangulation>
std::tuple<std::vector<std::complex<double>>,
           std::vector<dealii::Vector<double>>,
           std::vector<typename VectorType::value_type>,
           std::vector<dealii::types::global_dof_index>>
AMGe_host<dim, MeshEvaluator, VectorType>::compute_local_eigenvectors(
    unsigned int n_eigenvectors, double tolerance,
    Triangulation const &agglomerate_triangulation,
    std::map<typename dealii::Triangulation<dim>::active_cell_iterator,
             typename dealii::DoFHandler<dim>::active_cell_iterator> const
        &patch_to_global_map,
    MeshEvaluator const &evaluator, LobpcgScratchData const &scratch_data,
    typename std::enable_if_t<is_matrix_free<MeshEvaluator>::value &&
                                  std::is_class<Triangulation>::value,
                              int>) const
{
  dealii::DoFHandler<dim> agglomerate_dof_handler(agglomerate_triangulation);
  dealii::AffineConstraints<double> agglomerate_constraints;

  using AgglomerateOperator = MatrixFreeAgglomerateOperator<MeshEvaluator>;
  AgglomerateOperator agglomerate_operator(evaluator, agglomerate_dof_handler,
                                           agglomerate_constraints);

  auto const diag_elements = agglomerate_operator.get_diag_elements();

  // Compute the eigenvalues and the eigenvectors
  unsigned int const n_dofs_agglomerate = agglomerate_operator.m();
  std::vector<std::complex<double>> eigenvalues(n_eigenvectors);
  std::vector<dealii::Vector<double>> eigenvectors(
      n_eigenvectors, dealii::Vector<double>(n_dofs_agglomerate));

  auto const eigensolver_type =
      _eigensolver_params.get<std::string>("type", "lanczos");
  dealii::Vector<double> initial_vector(n_dofs_agglomerate);
  evaluator.set_initial_guess(agglomerate_constraints, initial_vector);
  if (eigensolver_type == "lanczos")
  {
    lanczos_compute_eigenvalues_and_eigenvectors(
        n_eigenvectors, tolerance, _eigensolver_params, agglomerate_operator,
        initial_vector, eigenvalues, eigenvectors);
  }
  else if (eigensolver_type == "anasazi")
  {
    anasazi_compute_eigenvalues_and_eigenvectors(
        n_eigenvectors, _eigensolver_params, agglomerate_operator,
        initial_vector, scratch_data.lobpcg_init_guess, eigenvalues,
        eigenvectors);
  }
  else if (eigensolver_type == "arpack")
  {
    throw std::runtime_error(
        "ARPACK not available as eigensolver in matrix-free mode");
  }
  else if (eigensolver_type == "lapack")
  {
    throw std::runtime_error(
        "LAPACK not available as eigensolver in matrix-free mode");
  }
  else
  {
    ASSERT(false, "Unknown eigensolver type '" + eigensolver_type + "'");
  }

  // Compute the map between the local and the global dof indices.
  std::vector<dealii::types::global_dof_index> dof_indices_map =
      this->compute_dof_index_map(patch_to_global_map, agglomerate_dof_handler);

  return std::make_tuple(eigenvalues, eigenvectors, diag_elements,
                         dof_indices_map);
}

template <int dim, typename MeshEvaluator, typename VectorType>
template <typename Triangulation>
std::tuple<std::vector<std::complex<double>>,
           std::vector<dealii::Vector<double>>,
           std::vector<typename VectorType::value_type>,
           std::vector<dealii::types::global_dof_index>>
AMGe_host<dim, MeshEvaluator, VectorType>::compute_local_eigenvectors(
    unsigned int n_eigenvectors, double tolerance,
    Triangulation const &agglomerate_triangulation,
    std::map<typename dealii::Triangulation<dim>::active_cell_iterator,
             typename dealii::DoFHandler<dim>::active_cell_iterator> const
        &patch_to_global_map,
    MeshEvaluator const &evaluator, LobpcgScratchData const &scratch_data,
    typename std::enable_if_t<!is_matrix_free<MeshEvaluator>::value &&
                                  std::is_class<Triangulation>::value,
                              int>) const
{
  dealii::DoFHandler<dim> agglomerate_dof_handler(agglomerate_triangulation);
  dealii::AffineConstraints<double> agglomerate_constraints;
  using value_type = typename VectorType::value_type;
  dealii::SparsityPattern agglomerate_sparsity_pattern;
  dealii::SparseMatrix<value_type> agglomerate_system_matrix;

  // Call user function to build the system matrix
  evaluator.evaluate_agglomerate(
      agglomerate_dof_handler, agglomerate_constraints,
      agglomerate_sparsity_pattern, agglomerate_system_matrix);

  // Get the diagonal elements
  unsigned int const size = agglomerate_system_matrix.m();
  std::vector<ScalarType> diag_elements(size);
  for (unsigned int i = 0; i < size; ++i)
    diag_elements[i] = agglomerate_system_matrix.diag_element(i);

  // Shift eigenvalues away from zero
  double const average_diagonal =
      std::accumulate(diag_elements.begin(), diag_elements.end(), 0.) / size;
  for (unsigned int i = 0; i < size; ++i)
    agglomerate_system_matrix.diag_element(i) += average_diagonal;
  // Shift diagonal entries for constrained degrees of freedom, to avoid
  // using the corresponding eigenvectors
  for (auto const constraint : agglomerate_constraints.get_lines())
  {
    agglomerate_system_matrix.diag_element(constraint.index) = 200;
  }

  // Compute the eigenvalues and the eigenvectors
  unsigned int const n_dofs_agglomerate = agglomerate_system_matrix.m();
  std::vector<std::complex<double>> eigenvalues(n_eigenvectors);
  // Arpack only works with double not float
  std::vector<dealii::Vector<double>> eigenvectors(
      n_eigenvectors, dealii::Vector<double>(n_dofs_agglomerate));

  dealii::Vector<double> initial_vector(n_dofs_agglomerate);
  evaluator.set_initial_guess(agglomerate_constraints, initial_vector);
  auto const eigensolver_type =
      _eigensolver_params.get<std::string>("type", "arpack");
  if (eigensolver_type == "arpack")
  {
    // Make Identity mass matrix
    Identity agglomerate_mass_matrix;

    dealii::SparseDirectUMFPACK inv_system_matrix;
    inv_system_matrix.initialize(agglomerate_system_matrix);

    dealii::SolverControl solver_control(n_dofs_agglomerate, tolerance);
    unsigned int const n_arnoldi_vectors = 2 * n_eigenvectors + 2;
    bool const symmetric = true;
    // We want the eigenvalues of the smallest magnitudes but we need to ask
    // for the ones with the largest magnitudes because they are computed for
    // the inverse of the matrix we care about.
    auto const which_eigenvalues =
        dealii::ArpackSolver::WhichEigenvalues::largest_magnitude;
    dealii::ArpackSolver::AdditionalData additional_data(
        n_arnoldi_vectors, which_eigenvalues, symmetric);
    dealii::ArpackSolver solver(solver_control, additional_data);

    // Compute the eigenvectors. Arpack outputs eigenvectors with a L2 norm of
    // one.
    solver.set_initial_vector(initial_vector);
    solver.solve(agglomerate_system_matrix, agglomerate_mass_matrix,
                 inv_system_matrix, eigenvalues, eigenvectors);
  }
  else if (eigensolver_type == "lanczos")
  {
    lanczos_compute_eigenvalues_and_eigenvectors(
        n_eigenvectors, tolerance, _eigensolver_params,
        agglomerate_system_matrix, initial_vector, eigenvalues, eigenvectors);
  }
  else if (eigensolver_type == "anasazi")
  {
    anasazi_compute_eigenvalues_and_eigenvectors(
        n_eigenvectors, _eigensolver_params, agglomerate_system_matrix,
        initial_vector, scratch_data.lobpcg_init_guess, eigenvalues,
        eigenvectors);
  }
  else if (eigensolver_type == "lapack")
  {
    // Use Lapack to compute the eigenvalues
    dealii::LAPACKFullMatrix<double> full_matrix;
    full_matrix.copy_from(agglomerate_system_matrix);

    double const lower_bound = -0.5;
    double const upper_bound = 100.;
    double const tol = 1e-12;
    dealii::Vector<double> lapack_eigenvalues(size);
    dealii::FullMatrix<double> lapack_eigenvectors;
    full_matrix.compute_eigenvalues_symmetric(
        lower_bound, upper_bound, tol, lapack_eigenvalues, lapack_eigenvectors);

    // Copy the eigenvalues and the eigenvectors in the right format
    for (unsigned int i = 0; i < n_eigenvectors; ++i)
      eigenvalues[i] = lapack_eigenvalues[i];

    for (unsigned int i = 0; i < n_eigenvectors; ++i)
      for (unsigned int j = 0; j < n_dofs_agglomerate; ++j)
        eigenvectors[i][j] = lapack_eigenvectors[j][i];
  }
  else
  {
    ASSERT(false, "Unknown eigensolver type '" + eigensolver_type + "'");
  }

  // Shift eigenvalues back
  for (unsigned int i = 0; i < n_eigenvectors; ++i)
    eigenvalues[i] -= average_diagonal;

  // Compute the map between the local and the global dof indices.
  std::vector<dealii::types::global_dof_index> dof_indices_map =
      this->compute_dof_index_map(patch_to_global_map, agglomerate_dof_handler);

  return std::make_tuple(eigenvalues, eigenvectors, diag_elements,
                         dof_indices_map);
}

template <int dim, typename MeshEvaluator, typename VectorType>
void AMGe_host<dim, MeshEvaluator, VectorType>::setup_restrictor(
    boost::property_tree::ptree const &agglomerate_ptree,
    unsigned int const n_eigenvectors, double const tolerance,
    MeshEvaluator const &evaluator,
    dealii::LinearAlgebra::distributed::Vector<
        typename VectorType::value_type> const &locally_relevant_global_diag,
    dealii::TrilinosWrappers::SparseMatrix &restriction_sparse_matrix)
{
  // Flag the cells to build agglomerates.
  unsigned int const n_agglomerates =
      this->build_agglomerates(agglomerate_ptree);

  // Parallel part of the setup.
  std::vector<unsigned int> agglomerate_ids(n_agglomerates);
  std::iota(agglomerate_ids.begin(), agglomerate_ids.end(), 1);
  std::vector<dealii::Vector<double>> eigenvectors;
  std::vector<std::vector<ScalarType>> diag_elements;
  std::vector<std::vector<dealii::types::global_dof_index>> dof_indices_maps;
  std::vector<unsigned int> n_local_eigenvectors;
  LobpcgScratchData scratch_data;
  CopyData copy_data;

  dealii::WorkStream::run(
      agglomerate_ids.begin(), agglomerate_ids.end(),
      [&](std::vector<unsigned int>::iterator const &agg_id,
          LobpcgScratchData &local_scratch_data, CopyData &local_copy_data) {
        this->local_worker(n_eigenvectors, tolerance, evaluator, agg_id,
                           local_scratch_data, local_copy_data);
      },
      [&](CopyData const &local_copy_data) {
        this->copy_local_to_global(local_copy_data, eigenvectors, diag_elements,
                                   dof_indices_maps, n_local_eigenvectors);
      },
      scratch_data, copy_data);

  AMGe<dim, VectorType>::compute_restriction_sparse_matrix(
      eigenvectors, diag_elements, dof_indices_maps, n_local_eigenvectors,
      locally_relevant_global_diag, restriction_sparse_matrix);

  // When checking the restriction matrix, we check that the sum of the local
  // diagonals is the global diagonals. This is not true for matrix-free because
  // the constraints values are set arbitrarily.
  if (std::is_base_of<DealIIMatrixFreeMeshEvaluator<dim>,
                      MeshEvaluator>::value == false)
  {
    check_restriction_matrix(this->_comm, eigenvectors, dof_indices_maps,
                             locally_relevant_global_diag, diag_elements,
                             n_local_eigenvectors);
  }
}

template <int dim, typename MeshEvaluator, typename VectorType>
void AMGe_host<dim, MeshEvaluator, VectorType>::setup_restrictor(
    boost::property_tree::ptree const &agglomerate_ptree,
    unsigned int const n_eigenvectors, double const tolerance,
    MeshEvaluator const &evaluator,
    dealii::LinearAlgebra::distributed::Vector<
        typename VectorType::value_type> const &locally_relevant_global_diag,
    std::shared_ptr<dealii::TrilinosWrappers::SparseMatrix>
        restriction_sparse_matrix,
    std::unique_ptr<dealii::TrilinosWrappers::SparseMatrix>
        &eigenvector_sparse_matrix,
    std::unique_ptr<dealii::TrilinosWrappers::SparseMatrix>
        &delta_eigenvector_matrix,
    std::vector<double> &eigenvalues)
{
  // Flag the cells to build agglomerates.
  unsigned int const n_agglomerates =
      this->build_agglomerates(agglomerate_ptree);

  // Parallel part of the setup.
  std::vector<unsigned int> agglomerate_ids(n_agglomerates);
  std::iota(agglomerate_ids.begin(), agglomerate_ids.end(), 1);
  std::vector<dealii::Vector<double>> eigenvectors;
  std::vector<std::vector<ScalarType>> diag_elements;
  std::vector<std::vector<dealii::types::global_dof_index>> dof_indices_maps;
  std::vector<unsigned int> n_local_eigenvectors;
  LobpcgScratchData scratch_data;
  CopyData copy_data;

  dealii::WorkStream::run(
      agglomerate_ids.begin(), agglomerate_ids.end(),
      [&](std::vector<unsigned int>::iterator const &agg_id,
          LobpcgScratchData &local_scratch_data, CopyData &local_copy_data) {
        this->local_worker(n_eigenvectors, tolerance, evaluator, agg_id,
                           local_scratch_data, local_copy_data);
      },
      [&](CopyData const &local_copy_data) {
        this->copy_local_to_global_eig(local_copy_data, eigenvalues,
                                       eigenvectors, diag_elements,
                                       dof_indices_maps, n_local_eigenvectors);
      },
      scratch_data, copy_data);

  AMGe<dim, VectorType>::compute_restriction_sparse_matrix(
      eigenvectors, diag_elements, dof_indices_maps, n_local_eigenvectors,
      locally_relevant_global_diag, restriction_sparse_matrix,
      eigenvector_sparse_matrix, delta_eigenvector_matrix);
}

template <int dim, typename MeshEvaluator, typename VectorType>
void AMGe_host<dim, MeshEvaluator, VectorType>::local_worker(
    unsigned int const n_eigenvectors, double const tolerance,
    MeshEvaluator const &evaluator,
    std::vector<unsigned int>::iterator const &agg_id,
    LobpcgScratchData &scratch_data, CopyData &copy_data)
{
  dealii::Triangulation<dim> agglomerate_triangulation;
  std::map<typename dealii::Triangulation<dim>::active_cell_iterator,
           typename dealii::DoFHandler<dim>::active_cell_iterator>
      agglomerate_to_global_tria_map;

  this->build_agglomerate_triangulation(*agg_id, agglomerate_triangulation,
                                        agglomerate_to_global_tria_map);

  std::tie(copy_data.local_eigenvalues, copy_data.local_eigenvectors,
           copy_data.diag_elements, copy_data.local_dof_indices_map) =
      compute_local_eigenvectors(
          n_eigenvectors, tolerance, agglomerate_triangulation,
          agglomerate_to_global_tria_map, evaluator, scratch_data);

  if (_eigensolver_params.get("type", "lanczos") == "anasazi")
  {
    if (_eigensolver_params.get("use_initial_guess", false))
    {
      // Copy the eigenvectors to be used as initial guess for LOBPCG
      scratch_data.lobpcg_init_guess = copy_data.local_eigenvectors;
    }
  }
}

template <int dim, typename MeshEvaluator, typename VectorType>
void AMGe_host<dim, MeshEvaluator, VectorType>::copy_local_to_global(
    CopyData const &copy_data,
    std::vector<dealii::Vector<double>> &eigenvectors,
    std::vector<std::vector<typename VectorType::value_type>> &diag_elements,
    std::vector<std::vector<dealii::types::global_dof_index>> &dof_indices_maps,
    std::vector<unsigned int> &n_local_eigenvectors)
{
  eigenvectors.insert(eigenvectors.end(), copy_data.local_eigenvectors.begin(),
                      copy_data.local_eigenvectors.end());

  diag_elements.push_back(copy_data.diag_elements);

  dof_indices_maps.push_back(copy_data.local_dof_indices_map);

  n_local_eigenvectors.push_back(copy_data.local_eigenvectors.size());
}

template <int dim, typename MeshEvaluator, typename VectorType>
void AMGe_host<dim, MeshEvaluator, VectorType>::copy_local_to_global_eig(
    CopyData const &copy_data, std::vector<double> &eigenvalues,
    std::vector<dealii::Vector<double>> &eigenvectors,
    std::vector<std::vector<typename VectorType::value_type>> &diag_elements,
    std::vector<std::vector<dealii::types::global_dof_index>> &dof_indices_maps,
    std::vector<unsigned int> &n_local_eigenvectors)
{
  copy_local_to_global(copy_data, eigenvectors, diag_elements, dof_indices_maps,
                       n_local_eigenvectors);

  std::transform(copy_data.local_eigenvalues.begin(),
                 copy_data.local_eigenvalues.end(),
                 std::back_inserter(eigenvalues),
                 [](std::complex<double> const &z) { return z.real(); });
}
} // namespace mfmg

#endif


/**
 * \file
 * \copyright
 * Copyright (c) 2012-2022, OpenGeoSys Community (http://www.opengeosys.org)
 *            Distributed under a Modified BSD License.
 *              See accompanying file LICENSE.txt or
 *              http://www.opengeosys.org/project/license
 */

#include "Coulomb.h"

#include <boost/math/constants/constants.hpp>

#include "BaseLib/Error.h"
#include "LogPenalty.h"
#include "MathLib/MathTools.h"
#include "NumLib/Exceptions.h"

namespace MaterialLib
{
namespace Fracture
{
namespace Coulomb
{
struct MaterialPropertyValues
{
    double Kn = 0.0;
    double Ks = 0.0;
    double phi = 0.0;  // friction angle
    double psi = 0.0;  // dilation angle
    double c = 0.0;

    template <typename MaterialProperties>
    MaterialPropertyValues(MaterialProperties const& mp,
                           double const t,
                           ParameterLib::SpatialPosition const& x)
    {
        Kn = mp.normal_stiffness(t, x)[0];
        Ks = mp.shear_stiffness(t, x)[0];
        auto constexpr degree =
            boost::math::constants::degree<double>();  // pi/180
        phi = mp.friction_angle(t, x)[0] * degree;
        psi = mp.dilatancy_angle(t, x)[0] * degree;
        c = mp.cohesion(t, x)[0];
    }
};

template <int DisplacementDim>
void Coulomb<DisplacementDim>::computeConstitutiveRelation(
    double const t,
    ParameterLib::SpatialPosition const& x,
    double const aperture0,
    Eigen::Ref<Eigen::VectorXd const>
    /*sigma0*/,
    Eigen::Ref<Eigen::VectorXd const>
        w_prev,
    Eigen::Ref<Eigen::VectorXd const>
        w,
    Eigen::Ref<Eigen::VectorXd const>
        sigma_prev,
    Eigen::Ref<Eigen::VectorXd>
        sigma,
    Eigen::Ref<Eigen::MatrixXd>
        Kep,
    typename FractureModelBase<DisplacementDim>::MaterialStateVariables&
        material_state_variables)
{
    assert(dynamic_cast<StateVariables<DisplacementDim> const*>(
               &material_state_variables) != nullptr);

    StateVariables<DisplacementDim>& state =
        static_cast<StateVariables<DisplacementDim>&>(material_state_variables);

    MaterialPropertyValues const mat(_mp, t, x);

    const int index_ns = DisplacementDim - 1;
    double const aperture = w[index_ns] + aperture0;

    Eigen::MatrixXd Ke;
    {  // Elastic tangent stiffness
        Ke = Eigen::MatrixXd::Zero(DisplacementDim, DisplacementDim);
        for (int i = 0; i < index_ns; i++)
        {
            Ke(i, i) = mat.Ks;
        }

        Ke(index_ns, index_ns) = mat.Kn;
    }

    // Total plastic aperture compression
    // NOTE: Initial condition sigma0 seems to be associated with an initial
    // condition of the w0 = 0. Therefore the initial state is not associated
    // with a plastic aperture change.
    {  // Exact elastic predictor
        sigma.noalias() = Ke * (w - w_prev);

        sigma.coeffRef(index_ns) *=
            logPenaltyDerivative(aperture0, aperture, _penalty_aperture_cutoff);
        sigma.noalias() += sigma_prev;
    }

    // correction for an opening fracture
    if (_tension_cutoff && sigma[DisplacementDim - 1] >= 0)
    {
        Kep.setZero();
        sigma.setZero();
        state.w_p = w;
        material_state_variables.setTensileStress(true);
        return;
    }

    auto yield_function = [&mat](Eigen::VectorXd const& s)
    {
        double const sigma_n = s[DisplacementDim - 1];
        Eigen::VectorXd const sigma_s = s.head(DisplacementDim - 1);
        double const mag_tau = sigma_s.norm();  // magnitude
        return mag_tau + sigma_n * std::tan(mat.phi) - mat.c;
    };

    {  // Exit if still in elastic range by checking the shear yield function.
        double const Fs = yield_function(sigma);
        material_state_variables.setShearYieldFunctionValue(Fs);
        if (Fs < .0)
        {
            Kep = Ke;
            Kep(index_ns, index_ns) *= logPenaltyDerivative(
                aperture0, aperture, _penalty_aperture_cutoff);
            return;
        }
    }

    auto yield_function_derivative = [&mat](Eigen::VectorXd const& s)
    {
        Eigen::Matrix<double, DisplacementDim, 1> dFs_dS;
        dFs_dS.template head<DisplacementDim - 1>().noalias() =
            s.template head<DisplacementDim - 1>().normalized();
        dFs_dS.coeffRef(DisplacementDim - 1) = std::tan(mat.phi);
        return dFs_dS;
    };

    // plastic potential function: Qs = |tau| + Sn * tan da
    auto plastic_potential_derivative = [&mat](Eigen::VectorXd const& s)
    {
        Eigen::Matrix<double, DisplacementDim, 1> dQs_dS;
        dQs_dS.template head<DisplacementDim - 1>().noalias() =
            s.template head<DisplacementDim - 1>().normalized();
        dQs_dS.coeffRef(DisplacementDim - 1) = std::tan(mat.psi);
        return dQs_dS;
    };

    {  // Newton

        Eigen::FullPivLU<Eigen::Matrix<double, 1, 1, Eigen::RowMajor>>
            linear_solver;
        using ResidualVectorType = Eigen::Matrix<double, 1, 1, Eigen::RowMajor>;
        using JacobianMatrix = Eigen::Matrix<double, 1, 1, Eigen::RowMajor>;

        JacobianMatrix jacobian;
        ResidualVectorType solution;
        solution << 0;

        auto const update_residual = [&](ResidualVectorType& residual)
        { residual[0] = yield_function(sigma); };

        auto const update_jacobian = [&](JacobianMatrix& jacobian)
        {
            jacobian(0, 0) = -yield_function_derivative(sigma).transpose() *
                             Ke * plastic_potential_derivative(sigma);
        };

        auto const update_solution = [&](ResidualVectorType const& increment)
        {
            solution += increment;
            /*DBUG("analytical = {:g}",
                 Fs / (mat.Ks + mat.Kn * std::tan(mat.psi) * std::tan(mat.phi)))
                 */
            state.w_p = state.w_p_prev +
                        solution[0] * plastic_potential_derivative(sigma);

            sigma.noalias() = Ke * (w - w_prev - state.w_p + state.w_p_prev);

            sigma.coeffRef(index_ns) *= logPenaltyDerivative(
                aperture0, aperture, _penalty_aperture_cutoff);
            sigma.noalias() += sigma_prev;
        };

        auto newton_solver =
            NumLib::NewtonRaphson<decltype(linear_solver), JacobianMatrix,
                                  decltype(update_jacobian), ResidualVectorType,
                                  decltype(update_residual),
                                  decltype(update_solution)>(
                linear_solver, update_jacobian, update_residual,
                update_solution, _nonlinear_solver_parameters);

        auto const success_iterations = newton_solver.solve(jacobian);

        if (!success_iterations)
        {
            throw NumLib::AssemblyException(
                "FractureModel/Coulomb local nonlinear solver didn't "
                "converge.");
        }

        // Solution containing lambda is not needed; w_p and sigma already
        // up to date.
    }

    {  // Update material state shear yield function value.
        double const Fs = yield_function(sigma);
        material_state_variables.setShearYieldFunctionValue(Fs);
    }

    Ke(index_ns, index_ns) *=
        logPenaltyDerivative(aperture0, aperture, _penalty_aperture_cutoff);
    Eigen::RowVectorXd const A = yield_function_derivative(sigma).transpose() *
                                 Ke /
                                 (yield_function_derivative(sigma).transpose() *
                                  Ke * plastic_potential_derivative(sigma));
    Kep = Ke - Ke * plastic_potential_derivative(sigma) * A;
}

template class Coulomb<2>;
template class Coulomb<3>;

}  // namespace Coulomb
}  // namespace Fracture
}  // namespace MaterialLib


/*****************************************************************************
*
* Rokko: Integrated Interface for libraries of eigenvalue decomposition
*
* Copyright (C) 2017 by Rokko Developers https://github.com/t-sakashita/rokko
*
* Distributed under 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 ROKKO_LAPACK_UNGQR_HPP
#define ROKKO_LAPACK_UNGQR_HPP

#include <complex>
#ifdef I
# undef I
#endif
#include <stdexcept>
#include <boost/type_traits/is_same.hpp>
#include <lapacke.h>
#include <rokko/traits/value_t.hpp>
#include "complex_cast.hpp"

namespace rokko {
namespace lapack {

namespace {

template<typename T> struct ungqr_dispatch;

template<>
struct ungqr_dispatch<float> {
  template<typename MATRIX, typename VECTOR>
  static lapack_int ungqr(int matrix_layout, lapack_int m, lapack_int n, lapack_int k,
                          MATRIX& a, VECTOR& tau) {
    return LAPACKE_sorgqr(matrix_layout, m, n, k, storage(a), ld(a), storage(tau));
  }
  template<typename MATRIX, typename VECTOR>
  static lapack_int ungqr(int matrix_layout, lapack_int m, lapack_int n, lapack_int k,
                          MATRIX& a, VECTOR& tau, VECTOR& work) {
    return LAPACKE_sorgqr(matrix_layout, m, n, k, storage(a), ld(a), storage(tau),
                          storage(work), size(work));
  }
};

template<>
struct ungqr_dispatch<double> {
  template<typename MATRIX, typename VECTOR>
  static lapack_int ungqr(int matrix_layout, lapack_int m, lapack_int n, lapack_int k,
                          MATRIX& a, VECTOR& tau) {
    return LAPACKE_dorgqr(matrix_layout, m, n, k, storage(a), ld(a), storage(tau));
  }
  template<typename MATRIX, typename VECTOR>
  static lapack_int ungqr(int matrix_layout, lapack_int m, lapack_int n, lapack_int k,
                          MATRIX& a, VECTOR& tau, VECTOR& work) {
    return LAPACKE_dorgqr(matrix_layout, m, n, k, storage(a), ld(a), storage(tau),
                          storage(work), size(work));
  }
};

template<>
struct ungqr_dispatch<std::complex<float>> {
  template<typename MATRIX, typename VECTOR>
  static lapack_int ungqr(int matrix_layout, lapack_int m, lapack_int n, lapack_int k,
                          MATRIX& a, VECTOR& tau) {
    return LAPACKE_cungqr(matrix_layout, m, n, k, complex_cast(storage(a)), ld(a),
                          complex_cast(storage(tau)));
  }
  template<typename MATRIX, typename VECTOR>
  static lapack_int ungqr(int matrix_layout, lapack_int m, lapack_int n, lapack_int k,
                          MATRIX& a, VECTOR& tau, VECTOR& work) {
    return LAPACKE_cungqr(matrix_layout, m, n, k, complex_cast(storage(a)), ld(a),
                          complex_cast(storage(tau)), complex_cast(storage(work)),
                          size(work));
  }
};

template<>
struct ungqr_dispatch<std::complex<double>> {
  template<typename MATRIX, typename VECTOR>
  static lapack_int ungqr(int matrix_layout, lapack_int m, lapack_int n, lapack_int k,
                          MATRIX& a, VECTOR& tau) {
    return LAPACKE_zungqr(matrix_layout, m, n, k, complex_cast(storage(a)), ld(a),
                          complex_cast(storage(tau)));
  }
  template<typename MATRIX, typename VECTOR>
  static lapack_int ungqr(int matrix_layout, lapack_int m, lapack_int n, lapack_int k,
                          MATRIX& a, VECTOR& tau, VECTOR& work) {
    return LAPACKE_zungqr(matrix_layout, m, n, k, complex_cast(storage(a)), ld(a),
                          complex_cast(storage(tau)), complex_cast(storage(work)),
                          size(work));
  }
};

}

template<typename MATRIX, typename VECTOR>
lapack_int ungqr(lapack_int k, MATRIX& a, VECTOR const& tau) {
  static_assert(std::is_same<value_t<MATRIX>, value_t<VECTOR>>::value, "");
  lapack_int m = rows(a);
  lapack_int n = cols(a);
  if (size(tau) != k)
    throw std::invalid_argument("vector tau size mismatch");
  return ungqr_dispatch<value_t<MATRIX>>
    ::ungqr((is_col_major(a) ? LAPACK_COL_MAJOR : LAPACK_ROW_MAJOR), m, n, k, a, tau);
}

template<typename MATRIX, typename VECTOR>
lapack_int ungqr(lapack_int k, MATRIX& a, VECTOR const& tau, VECTOR& work) {
  static_assert(std::is_same<value_t<MATRIX>, value_t<VECTOR>>::value, "");
  lapack_int m = rows(a);
  lapack_int n = cols(a);
  if (size(tau) != k)
    throw std::invalid_argument("vector tau size mismatch");
  if (size(work) < std::max(1, n))
    throw std::invalid_argument("vector work size mismatch");
  return ungqr_dispatch<value_t<MATRIX>>
    ::ungqr((is_col_major(a) ? LAPACK_COL_MAJOR : LAPACK_ROW_MAJOR), m, n, k, a, tau, work);
}

template<typename MATRIX, typename VECTOR>
lapack_int orgqr(lapack_int k, MATRIX& a, VECTOR const& tau) {
  return ungqr(k, a, tau);
}

template<typename MATRIX, typename VECTOR>
lapack_int orgqr(lapack_int k, MATRIX& a, VECTOR const& tau, VECTOR& work) {
  return ungqr(k, a, tau, work);
};

} // end namespace lapack
} // end namespace rokko

#endif // ROKKO_LAPACK_UNGQR_HPP