#include <iostream>
#include <fstream>
#include <string>
#include <stdio.h>
#include <time.h>

// OpenCV headers
#include <cv.h>
#include <highgui.h>

#include <opencv2/contrib/contrib.hpp>
#include <opencv2/core/core.hpp>
#include <opencv2/highgui/highgui.hpp>
#include <opencv2/video/tracking.hpp>

#include <opencv_workbench/syllo/syllo.h>
//
#include <opencv_workbench/utils/AnnotationParser.h>
#include <opencv_workbench/plot/Plot.h>
#include <opencv_workbench/track/Dynamics.h>
#include <opencv_workbench/track/KalmanFilter.h>
#include <opencv_workbench/track/EKF.h>

#include <boost/random.hpp>
#include <boost/random/normal_distribution.hpp>

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

int main(int argc, char *argv[])
{
     cout << "Kalman Filter Test" << endl;

     double t0 = 0;
     double dt = 0.1;
     double tend = 10;

     // Setup boost RNG
     boost::mt19937 rng;
     boost::normal_distribution<> nd(0,1.0);
     boost::variate_generator<boost::mt19937&,
                              boost::normal_distribution<> > var_nor(rng,nd);

     ///////////////////////////////////////////////////////
     // Setup Cart Model
     Dynamics::state_5d_type input;
     input[0] = 5;
     input[1] = 3.14159265359/20;//0

     Dynamics dyn;
     dyn.set_time(t0, dt, tend);
     dyn.set_model(Dynamics::cart);
     dyn.set_input(input);
     dyn.set_process_noise(0.0);
     dyn.set_measurement_noise(0.1);

     dyn.compute_trajectory();
     std::vector<cv::Point2d> truth = dyn.truth_points();
     std::vector<cv::Point2d> measured = dyn.measured_points();

     ////////////////////////////////////////////////////////
     // Setup Kalman Filter
     // Kalman filter setup...

     syllo::KalmanFilter kf;
     Eigen::MatrixXf A;     // system matrix
     Eigen::MatrixXf B;     // input matrix
     Eigen::MatrixXf H;     // measurement matrix
     Eigen::MatrixXf Q;     // process noise matrix
     Eigen::MatrixXf R;     // measurement noise matrix
     Eigen::MatrixXf x0;    // initial position
     Eigen::MatrixXf covar; // initial state covariance matrix

     A.resize(4,4);
     B.resize(4,2);
     H.resize(2,4);
     Q.resize(4,4);
     R.resize(2,2);
     x0.resize(4,1);
     covar.resize(4,4);

     double T = dt;
     A << 1, 0, T, 0,
          0, 1, 0, T,
          0, 0, 1, 0,
          0, 0, 0, 1;

     B << 0, 0,
          1, 0,
          0, 0,
          0, 1;

     H << 1, 0, 0, 0,
          0, 1, 0, 0;

     Q = Eigen::MatrixXf::Identity(A.rows(), A.cols()) * 1e-4;

     //R << 10, 0,
     //     0, 10;
     R << .001, 0,
          0, .001;

     //x = [x, y, x_dot, y_dot]
     x0 << 0,
           0,
           0,
           0;

     covar << 0.1, 0, 0, 0,
              0, 0.1, 0, 0,
              0, 0, 0.1, 0,
              0, 0, 0, 0.1;

     kf.setModel(A,B,H,Q,R);
     kf.init(x0, covar);

     // Setup OpenCV's kalman filter tracker
     cv::KalmanFilter opencv_kf = cv::KalmanFilter(4, 2, 0);
     cv::Mat_<float> transition_matrix = cv::Mat_<float>(4,4);
     transition_matrix  << 1,0,T,0,
                           0,1,0,T,
                           0,0,1,0,
                           0,0,0,1;

     opencv_kf.transitionMatrix = transition_matrix;

     opencv_kf.statePre.at<float>(0) = 0;
     opencv_kf.statePre.at<float>(1) = 0;
     opencv_kf.statePre.at<float>(2) = 2;
     opencv_kf.statePre.at<float>(3) = 0;
     cv::setIdentity(opencv_kf.measurementMatrix);
     cv::setIdentity(opencv_kf.processNoiseCov, cv::Scalar::all(1e-4));
     cv::setIdentity(opencv_kf.measurementNoiseCov, cv::Scalar::all(0.001));
     cv::setIdentity(opencv_kf.errorCovPost, cv::Scalar::all(0.1));

     /////////////////////////////////////////////////////////
     // Process the measured points with the kalman filter
     std::vector<cv::Point2d> kf_points;
     std::vector<cv::Point2d> opencv_kf_points;
     std::vector<cv::Point2d> kf_var_x;
     std::vector<cv::Point2d> kf_var_y;
     std::vector<cv::Point2d> error;
     std::vector<cv::Point2d>::iterator it = measured.begin();
     std::vector<cv::Point2d>::iterator it_truth = truth.begin();
     double t = t0;
     for(; it != measured.end(); it++) {
          Eigen::MatrixXf u, z;
          u.resize(2,1);
          u << 0,
               0;

          z.resize(2,1);
          z << it->x , it->y;

          kf.predict(u);
          kf.update(z);

          //////////////////////////////////////////
          // Opencv
          cv::Mat prediction = opencv_kf.predict();
          cv::Point predictPt(prediction.at<float>(0),prediction.at<float>(1));

          // correct tracker update using the newly computed centroid.
          cv::Mat_<float> measurement(2,1);
          measurement(0) = it->x;
          measurement(1) = it->y;

          cv::Mat estimated = opencv_kf.correct(measurement);
          cv::Point2d est_centroid = cv::Point2d(estimated.at<float>(0),estimated.at<float>(1));
          opencv_kf_points.push_back(est_centroid);
          ////////////////////////////////////////////

          Eigen::MatrixXf state = kf.state();
          kf_points.push_back(cv::Point2d(state(0,0),state(1,0)));

          Eigen::MatrixXf covar = kf.covariance();
          kf_var_x.push_back(cv::Point2d(t,covar(0,0)));
          kf_var_y.push_back(cv::Point2d(t,covar(1,1)));

          double err = sqrt( pow(state(0,0) - it_truth->x, 2) + pow(state(1,0) - it_truth->y, 2) );
          error.push_back(cv::Point2d(t,err));

          it_truth++;
          t += dt;

#if 0
          std::string temp;
          std::cin >> temp;
#endif
     }

     /////////////////////////////////////////////////////////
     // Plot the tracks
     std::vector< std::vector<cv::Point2d> > vectors;
     const std::string title = "Tracks";
     std::vector<std::string> labels;
     std::vector<std::string> styles;

     vectors.push_back(truth);
     labels.push_back("Truth");
     styles.push_back("points");

     vectors.push_back(measured);
     labels.push_back("Measured");
     styles.push_back("points");

     vectors.push_back(kf_points);
     labels.push_back("KF");
     styles.push_back("linespoints");

     vectors.push_back(opencv_kf_points);
     labels.push_back("OPENCV KF");
     styles.push_back("linespoints");

     syllo::Plot plot;
     plot.plot(vectors, title, labels, styles);

     ////////////////////////////////////////////////////
     vectors.clear() ; labels.clear(); styles.clear();
     const std::string title_covar = "Covariance";

     vectors.push_back(kf_var_x);
     labels.push_back("X Var");
     styles.push_back("points");

     vectors.push_back(kf_var_y);
     labels.push_back("Y Var");
     styles.push_back("points");

     vectors.push_back(error);
     labels.push_back("Error");
     styles.push_back("points");

     syllo::Plot plot_covar;
     plot_covar.plot(vectors, title_covar, labels, styles);
     return 0;
}


/**************************************************************************
** This file is a part of our work (Siggraph'16 paper, binary, code and dataset):
**
** Roto++: Accelerating Professional Rotoscoping using Shape Manifolds
** Wenbin Li, Fabio Viola, Jonathan Starck, Gabriel J. Brostow and Neill D.F. Campbell
**
** w.li AT cs.ucl.ac.uk
** http://visual.cs.ucl.ac.uk/pubs/rotopp
**
** Copyright (c) 2016, Wenbin Li
** All rights reserved.
**
** Redistribution and use in source and binary forms, with or without
** modification, are permitted provided that the following conditions are met:
**
** -- Redistributions of source code and data must retain the above
**    copyright notice, this list of conditions and the following disclaimer.
** -- Redistributions in binary form must reproduce the above copyright
**    notice, this list of conditions and the following disclaimer in the
**    documentation and/or other materials provided with the distribution.
**
** THIS WORK AND THE RELATED SOFTWARE, SOURCE CODE AND DATA IS PROVIDED BY
** THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED
** WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF
** MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN
** NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT,
** INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
** BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF
** USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
** THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
** NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE,
** EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
***************************************************************************/

#ifndef GPLVM_HPP
#define GPLVM_HPP

#include "ceres/ceres.h"
#include "glog/logging.h"
#include "ceres/solver.h"

using ceres::AutoDiffCostFunction;
using ceres::CostFunction;
using ceres::Problem;
using ceres::Solver;

#include <Eigen/Core>
#include <Eigen/Dense>
#include <Eigen/SparseCore>
#include <Eigen/Eigenvalues>

#include <array>
#include <random>
#include <iostream>
#include <cmath>
#include <fstream>
#include <cassert>
#include <chrono>

#include "include/rotoSolver/kernel.hpp"
#include "include/rotoSolver/baseDefs.hpp"

using namespace std;

class NoiseHyperPriorFunction
{
public:
    virtual double getHyperPrior(const double logBeta) const = 0;
    virtual double getHyperPriorDerivative(const double logBeta) const = 0;
};

class GaussianNoiseHyperPriorFunction : public NoiseHyperPriorFunction
{
private:
    double _logOffset;
    double _logScaleSq;

public:
    GaussianNoiseHyperPriorFunction(double logOffset = 0.0, double logScale = 1.0) :
        _logOffset(logOffset), _logScaleSq(logScale * logScale)
    {}

    virtual double getHyperPrior(const double logBeta) const
    {
        double b = logBeta - _logOffset;
        return (0.5 * _logScaleSq * b * b);
    }

    virtual double getHyperPriorDerivative(const double logBeta) const
    {
        return (_logScaleSq * logBeta);
    }
};

class GPLVM {
public:
    typedef Eigen::VectorXd Vector;
    typedef Eigen::MatrixXd Matrix;
    typedef Eigen::SparseMatrix<double> SparseMatrix;
    explicit GPLVM(const Matrix& Y, const int Q,
                   shared_ptr<Kernel> kernel,
                   const double initBeta = 1,
                   shared_ptr<NoiseHyperPriorFunction> noisePriorFunc =
     shared_ptr<NoiseHyperPriorFunction>(new GaussianNoiseHyperPriorFunction())) :
        _Y(Y), _N(_Y.rows()), _Q(Q), _kernel(kernel), _beta(initBeta),
        _noisePriorFunc(noisePriorFunc)
    {
        assert (_N > 0);
        assert (_Q > 0);

        const double N = double(_Y.rows());

        _Ymean = _Y.array().colwise().mean();
        _Y = _Y.array().rowwise() - _Ymean.array().colwise().mean();
        _Ystd = (_Y.array().pow(2.0).colwise().sum() / (N-1.0)).sqrt();
        _Y = _Y.array().rowwise() / _Ystd.array().colwise().mean();

#ifdef CHECK_NORMALISATION
        {
            vdbg(_Ymean);
            vdbg(_Ystd);
            vdbg(_Y.transpose());

            Matrix yy_mean = _Y;
            yy_mean = yy_mean.array().rowwise() * _Ystd.array().colwise().mean();
            yy_mean = yy_mean.array().rowwise() + _Ymean.array().colwise().mean();

            vdbg((yy_mean - Y).norm());
        }
#endif // CHECK_NORMALISATION

        _X = GetDefaultXInit();

        resetCachedValues();
    }

    void Print() const
    {
        std::cout << "GPLVM: ";
        std::cout << " N = " << _N;
        std::cout << " Q = " << _Q;
        std::cout << " beta = " << _beta << std::endl;
        _kernel->Print();
    }

    void SetInternalsFromMatlabStructure(std::map<string, shared_ptr<Eigen::MatrixXd> >& S)
    {
        _Y = *S["Y"];
        _Ymean = *S["Y_mean"];
        _Ystd = *S["Y_std"];

        _X = *S["X"];

        resetCachedValues();
    }

    inline int NumParameters() const
    {
        return NumPreDataParameters() + (_N * _Q);
    }

    inline int NumPreDataParameters() const
    {
        return (_kernel->NumberOfParameters() + 1);
    }

    int GetN() const { return _N; }

    int GetQ() const { return _Q; }

    // Cannot return ref due to normalisation so don't call too often..
    Matrix GetY() const
    {
        Matrix Y(_Y);
        Y = Y.array().rowwise() * _Ystd.array().colwise().mean();
        Y = Y.array().rowwise() + _Ymean.array().colwise().mean();
        return Y;
    }

    Matrix GetX() const { return _X; }

    double GetBeta() const { return _beta; }

    void SetBeta(const double newBeta)
    {
        _beta = newBeta;

        resetCachedValues();
    }

    template<typename Derived>
    Matrix Predict(const Eigen::MatrixBase<Derived>& x,
                   Vector* y_var_ptr = NULL,
                   std::vector<Matrix>* dy_dx_vec_ptr = NULL) const
    {
        assert (x.cols() == _Q);
        assert (x.rows() > 0);

        Matrix K_x_X = _kernel->GetKernelMatrix(x, _X, NULL);

        const Matrix& KinvY = GetKinvY();

        Matrix y_mean = K_x_X * KinvY;

        if (y_var_ptr) {
            Vector& y_var = *(y_var_ptr);
            y_var.resize(x.rows());

            const Matrix& Kinv = GetKinv();

            Vector K_xx_xx_diag = _kernel->GetDiagonalOfKernel(x);

            y_var = K_xx_xx_diag.array() -
                    (K_x_X.array() * (K_x_X * Kinv.transpose()).array()).rowwise().sum().array();
            y_var.array() += (1.0 / _beta);
        }

        y_mean = y_mean.array().rowwise() * _Ystd.array().colwise().mean();
        y_mean = y_mean.array().rowwise() + _Ymean.array().colwise().mean();

        if (dy_dx_vec_ptr)
        {
            std::vector<Matrix>& dy_dx_vec = *(dy_dx_vec_ptr);
            dy_dx_vec.clear();

            const Matrix& Kinv = GetKinv();

            const int D = _Y.cols();
            const int P = x.rows();

            dy_dx_vec.reserve(P);

            for (int p = 0; p < P; ++p)
            {
                shared_ptr<SparseMatrix> dK_dx = _kernel->GetGradientWrtSecondTerm(_X, x.row(p), K_x_X.row(p).transpose());

                Matrix dy_dx = (*dK_dx) * KinvY;

                dy_dx = dy_dx.array().rowwise() * _Ystd.array().colwise().mean();

                dy_dx_vec.push_back(dy_dx);
            }
        }

        return y_mean;
    }

    void SaveToMatFile(std::string matFilename) const
    {
        std::cerr << "UNSUPPORTED: UNABLE TO WRITE GPLVM MATLAB DATA TO \"" << matFilename << "\"." << std::endl;
    }

private:
    Matrix _Y;
    Matrix _Ymean;
    Matrix _Ystd;
    const int _N;
    const int _Q;
    shared_ptr<Kernel> _kernel;
    double _beta;
    Matrix _X;
    shared_ptr<NoiseHyperPriorFunction> _noisePriorFunc;

private:
    // Cached values (hence mutable)..
    mutable shared_ptr<Matrix> _KinvY_ptr;
    mutable shared_ptr<Matrix> _Kinv_ptr;

    void resetCachedValues()
    {
        _KinvY_ptr.reset();
        _Kinv_ptr.reset();
    }

protected:

    class CostFunction : public ceres::FirstOrderFunction
    {
    private:
        GPLVM* _gplvm;
        shared_ptr<Kernel> _kernel;

    public:
        explicit CostFunction(GPLVM* gplvm) :
            _gplvm(gplvm),
            _kernel(gplvm->_kernel->createCopy())
        {
        }

        virtual bool Evaluate(const double* parameters,
                              double* cost,
                              double* gradient) const {
            const int numParams = _gplvm->NumPreDataParameters();
            const int numKernelParams = _kernel->NumberOfParameters();
            const int N = _gplvm->_N;
            const int Q = _gplvm->_Q;
            const Matrix& Y = _gplvm->_Y;

            const double logBeta = parameters[0];
            const double beta = exp(logBeta);
            const Eigen::Map< const Vector > logKernelParams(&parameters[1], numKernelParams);

            _kernel->SetParameters(logKernelParams);

            const Eigen::Map< const Matrix > X(&parameters[numParams], N, Q);
            const Eigen::Map< const Vector > Xvec(&parameters[numParams], N * Q);

            // Create noise free copy of K and prepare extra data
            Matrix DDmat;
            Matrix KK = _kernel->GetKernelMatrix(X, X, &DDmat);
            Matrix K(KK);

            const double Nd = N;
            const double Dd = Y.cols();

            K.diagonal().array() += (1.0 / beta);
            Matrix L = K.llt().matrixL();

            double logDet = 2.0 * (L.diagonal().array().log()).sum();

            Matrix LinvY = K.llt().matrixL().solve(Y);
            double trace_Kinv_YYt = LinvY.squaredNorm();

            double hyperPriorNoise = _gplvm->_noisePriorFunc->getHyperPrior(logBeta);

            double hyperPriorParams = 0.5 * (logKernelParams.array() * logKernelParams.array()).sum();

            double hyperPriorData = 0.5 * X.squaredNorm();

            double LogLike = (0.5 * Dd * Nd) * log(2.0 * M_PI);
            LogLike += 0.5 * Dd * logDet;
            LogLike += 0.5 * trace_Kinv_YYt;
            LogLike += hyperPriorParams;
            LogLike += hyperPriorNoise;

            LogLike += hyperPriorData;

            cost[0] = LogLike;

            if (gradient != NULL)
            {
                // The beta gradient
                double& dL_dlogBeta = gradient[0];

                // The kernel params gradient
                Eigen::Map< Eigen::VectorXd > dL_dlogKernParams(&gradient[1], numKernelParams);

                // The final data gradient
                Eigen::Map< Eigen::VectorXd > dL_dX(&gradient[numParams], N*Q);

                Matrix dK_dbeta(Matrix::Identity(N, N));
                dK_dbeta.diagonal() /= - beta;

                // std::vector< shared_ptr< Matrix > >
                auto dK_dlogKernParams = _kernel->GetGradientWrtParams(X, KK, &DDmat);

                // shared_ptr< Eigen::SparseMatrix<double> >
                auto dK_dX_ptr = _kernel->GetGradientWrtData(X, KK, &DDmat);

                Matrix KinvY = K.llt().matrixL().transpose().solve(LinvY);

                Matrix dL_dK = - 0.5 * (KinvY * KinvY.transpose() - (Dd * K.llt().solve(Matrix::Identity(N, N))));

                dL_dlogBeta = (dL_dK.array() * dK_dbeta.array()).sum();
                dL_dlogBeta += _gplvm->_noisePriorFunc->getHyperPriorDerivative(logBeta);

                for (int i = 0; i < numKernelParams; ++i)
                {
                    dL_dlogKernParams[i] = (dL_dK.array() * (*(dK_dlogKernParams[i])).array()).sum() + logKernelParams[i];
                }

                Eigen::Map< Vector > dL_dK_vec(dL_dK.data(), N*N);

                dL_dX = (*dK_dX_ptr) * dL_dK_vec;

                // Add hyperPriorData gradient
                dL_dX += Xvec;
            }

            return true;
        }

        virtual int NumParameters() const
        {
            return _gplvm->NumParameters();
        }
    };

private:

    Matrix GetDefaultXInit() const
    {
        Eigen::SelfAdjointEigenSolver<Eigen::MatrixXd> es(_Y * _Y.transpose());
        Matrix X0 = es.eigenvectors().rowwise().reverse().block(0, 0, _N, _Q);

        const double scaling = ((X0.array().pow(2.0).colwise().sum() / (double(_N)-1.0)).sqrt()).mean();
        X0 /= scaling;

        return X0;
    }

    Vector GetParameterVector() const
    {
        Vector p(NumParameters());

        p[0] = log(_beta);

        Vector logKernParams;
        _kernel->GetParameters(logKernParams);
        p.block(1, 0, _kernel->NumberOfParameters(), 1) = logKernParams;

        const Eigen::Map< const Vector > Xvec(_X.data(), _N*_Q);
        p.block(NumPreDataParameters(), 0, _N*_Q, 1) = Xvec;

        return p;
    }

    void SetFromParameterVector(const Vector& p)
    {
        _beta = exp(p[0]);

        Vector logKernParams = p.block(1, 0, _kernel->NumberOfParameters(), 1);
        _kernel->SetParameters(logKernParams);

        Eigen::Map< Vector > Xvec(_X.data(), _N*_Q);
        Xvec = p.block(NumPreDataParameters(), 0, _N*_Q, 1);

        resetCachedValues();
    }

    const Matrix& GetKinvY() const
    {
        if (!_KinvY_ptr.get())
        {
            std::cout << "Computing KinvY cached value.." << std::endl;

            Matrix DDmat;
            Matrix K = _kernel->GetKernelMatrix(_X, _X, &DDmat);

            K.diagonal().array() += (1.0 / _beta);

            Matrix LinvY = K.llt().matrixL().solve(_Y);

            shared_ptr<Matrix> tmp(new Matrix(K.llt().matrixL().transpose().solve(LinvY)));

            _KinvY_ptr.swap(tmp);
        }

        return (*_KinvY_ptr);
    }

    const Matrix& GetKinv() const
    {
        if (!_Kinv_ptr.get()) {
            std::cout << "Computing Kinv cached value.." << std::endl;

            Matrix DDmat;
            Matrix K = _kernel->GetKernelMatrix(_X, _X, &DDmat);

            K.diagonal().array() += (1.0 / _beta);

            Matrix Linv = K.llt().matrixL().solve(Matrix::Identity(_N, _N));
            shared_ptr<Matrix> tmp(new Matrix(Linv.transpose() * Linv));

            _Kinv_ptr.swap(tmp);

#ifdef TEST_OTHER_INV_METHOD
//            vdbg(K * (*_Kinv_ptr));
//            vdbg((*_Kinv_ptr) * K);

//            vdbg((K * (*_Kinv_ptr) - Matrix::Identity(_N,_N)).norm());
//            vdbg(((*_Kinv_ptr) * K - Matrix::Identity(_N,_N)).norm());
            vdbg((K * Kinv - Matrix::Identity(_N,_N)).norm());
            vdbg((Kinv * K - Matrix::Identity(_N,_N)).norm());

            start = std::chrono::high_resolution_clock::now();
            Matrix Kinv2 = K.inverse();
            stop = std::chrono::high_resolution_clock::now();
            vdbg(std::chrono::duration_cast<std::chrono::microseconds>(stop-start).count());

//            vdbg(K * Kinv);
//            vdbg(Kinv * K);

            vdbg((K * Kinv2 - Matrix::Identity(_N,_N)).norm());
            vdbg((Kinv2 * K - Matrix::Identity(_N,_N)).norm());
#endif // TEST_OTHER_INV_METHOD
        }

        return (*_Kinv_ptr);
    }

public:

    void TestGplvmGradients(const double step = 1e-8)
    {
        using std::cout; using std::endl;
        cout << endl << " TESTING GPLVM GRADIENTS ";
        cout << endl << " ======================= " << endl;

        CostFunction costFunc(this);

        Vector parameters(GetParameterVector());

        parameters.setRandom();

        double C0 = 0.0;
        Vector gradients(NumParameters());

        bool r = costFunc.Evaluate(parameters.data(), &C0, gradients.data());
        assert (r);

        vdbg(C0);
        vdbg(gradients.transpose());

        Vector estGradients(NumParameters());
        estGradients.setZero();

        for (int i = 0, I = NumParameters(); i < I; ++i)
        {
            Vector p(parameters);
            p[i] += step;
            double C = 0;

            bool r = costFunc.Evaluate(p.data(), &C, NULL);
            assert (r);

            estGradients[i] = (C - C0) / step;
        }

        vdbg(estGradients.transpose());

        vdbg((estGradients - gradients).transpose());
        vdbg((estGradients - gradients).norm() / gradients.norm());

        assert ((estGradients - gradients).norm() / gradients.norm() < 1e-4);

        cout << " ======================= " << endl << endl;
    }

    void TestPredictionGradients(Matrix x = Matrix(0,0), const double step = 1e-8)
    {
        using std::cout; using std::endl;
        cout << endl << " TESTING PREDICTION GRADIENTS ";
        cout << endl << " ============================ " << endl;

        if (x.size() == 0)
        {
            x = GetX();
        }

        Vector yVar;
        std::vector<Matrix> dy_dx_vec;
        Matrix y = Predict(x, &yVar, &dy_dx_vec);

        const int P = x.rows();

        for (int p = 0; p < P; ++p)
        {
            Matrix xx = x.row(p);
            xx = xx + Matrix::Constant(xx.rows(), xx.cols(), step);
            Matrix yy = (Predict(xx) - y.row(p));
            yy /= step;

            vdbg((yy - dy_dx_vec[p]).norm());

            assert ((yy - dy_dx_vec[p]).norm() / (dy_dx_vec[p]).norm() < (10 * step));
        }

        cout << " ============================ " << endl << endl;
    }

    // Optimise to convergence or until the first of the specified limits is reached..
    void LearnParameters(const double maxSolverTimeInSeconds = 1.0,
                         const int maxNumIterations = 10000,
                         const bool showProgress = false)
    {
        ceres::GradientProblem problem(new CostFunction(this));

        Vector parameters(GetParameterVector());

        ceres::GradientProblemSolver::Options options;
        options.minimizer_progress_to_stdout = showProgress;

        options.function_tolerance = 1e-6;
        options.gradient_tolerance = 1e-4;

        options.max_num_iterations = maxNumIterations;
        options.max_solver_time_in_seconds = maxSolverTimeInSeconds;

        ceres::GradientProblemSolver::Summary summary;
        ceres::Solve(options, problem, parameters.data(), &summary);

        SetFromParameterVector(parameters);
    }

};

#endif // GPLVM_HPP


/*
 *
 * Copyright (c) Kresimir Fresl 2002, 2003
 *
 * 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)
 *
 * Author acknowledges the support of the Faculty of Civil Engineering,
 * University of Zagreb, Croatia.
 *
 */

#ifndef BOOST_NUMERIC_BINDINGS_CBLAS2_OVERLOADS_HPP
#define BOOST_NUMERIC_BINDINGS_CBLAS2_OVERLOADS_HPP

#include <complex>
#include <boost/numeric/bindings/atlas/cblas_inc.hpp>
#include <boost/numeric/bindings/traits/type.hpp>

namespace boost { namespace numeric { namespace bindings {

  namespace atlas { namespace detail {

    // y <- alpha * op (A) * x + beta * y

    inline
    void gemv (CBLAS_ORDER const Order,
               CBLAS_TRANSPOSE const TransA, int const M, int const N,
               float const alpha, float const* A, int const lda,
               float const* X, int const incX,
               float const beta, float* Y, int const incY)
    {
      cblas_sgemv (Order, TransA, M, N, alpha, A, lda,
                   X, incX,
                   beta, Y, incY);
    }

    inline
    void gemv (CBLAS_ORDER const Order,
               CBLAS_TRANSPOSE const TransA, int const M, int const N,
               double const alpha, double const* A, int const lda,
               double const* X, int const incX,
               double const beta, double* Y, int const incY)
    {
      cblas_dgemv (Order, TransA, M, N, alpha, A, lda,
                   X, incX,
                   beta, Y, incY);
    }

    inline
    void gemv (CBLAS_ORDER const Order,
               CBLAS_TRANSPOSE const TransA, int const M, int const N,
               traits::complex_f const& alpha,
               traits::complex_f const* A, int const lda,
               traits::complex_f const* X, int const incX,
               traits::complex_f const& beta,
               traits::complex_f* Y, int const incY)
    {
      cblas_cgemv (Order, TransA, M, N,
                   static_cast<void const*> (&alpha),
                   static_cast<void const*> (A), lda,
                   static_cast<void const*> (X), incX,
                   static_cast<void const*> (&beta),
                   static_cast<void*> (Y), incY);
    }

    inline
    void gemv (CBLAS_ORDER const Order,
               CBLAS_TRANSPOSE const TransA, int const M, int const N,
               traits::complex_d const& alpha,
               traits::complex_d const* A, int const lda,
               traits::complex_d const* X, int const incX,
               traits::complex_d const& beta,
               traits::complex_d* Y, int const incY)
    {
      cblas_zgemv (Order, TransA, M, N,
                   static_cast<void const*> (&alpha),
                   static_cast<void const*> (A), lda,
                   static_cast<void const*> (X), incX,
                   static_cast<void const*> (&beta),
                   static_cast<void*> (Y), incY);
    }

    // y <- alpha * A * x + beta * y
    // A real symmetric

    inline
    void symv (CBLAS_ORDER const Order, CBLAS_UPLO const Uplo,
               int const N, float const alpha, float const* A,
               int const lda, float const* X, int const incX,
               float const beta, float* Y, int const incY)
    {
      cblas_ssymv (Order, Uplo, N, alpha, A, lda,
                   X, incX, beta, Y, incY);
    }

    inline
    void symv (CBLAS_ORDER const Order, CBLAS_UPLO const Uplo,
               int const N, double const alpha, double const* A,
               int const lda, double const* X, int const incX,
               double const beta, double* Y, int const incY)
    {
      cblas_dsymv (Order, Uplo, N, alpha, A, lda,
                   X, incX, beta, Y, incY);
    }

    // y <- alpha * A * x + beta * y
    // A real symmetric in packed form

    inline
    void spmv (CBLAS_ORDER const Order, CBLAS_UPLO const Uplo,
               int const N, float const alpha, float const* Ap,
               float const* X, int const incX,
               float const beta, float* Y, int const incY)
    {
      cblas_sspmv (Order, Uplo, N, alpha, Ap, X, incX, beta, Y, incY);
    }

    inline
    void spmv (CBLAS_ORDER const Order, CBLAS_UPLO const Uplo,
               int const N, double const alpha, double const* Ap,
               double const* X, int const incX,
               double const beta, double* Y, int const incY)
    {
      cblas_dspmv (Order, Uplo, N, alpha, Ap, X, incX, beta, Y, incY);
    }

    // y <- alpha * A * x + beta * y
    // A complex hermitian

    inline
    void hemv (CBLAS_ORDER const Order, CBLAS_UPLO const Uplo,
               int const N, traits::complex_f const& alpha,
               traits::complex_f const* A, int const lda,
               traits::complex_f const* X, int const incX,
               traits::complex_f const& beta,
               traits::complex_f* Y, int const incY)
    {
      cblas_chemv (Order, Uplo, N,
                   static_cast<void const*> (&alpha),
                   static_cast<void const*> (A), lda,
                   static_cast<void const*> (X), incX,
                   static_cast<void const*> (&beta),
                   static_cast<void*> (Y), incY);
    }

    inline
    void hemv (CBLAS_ORDER const Order, CBLAS_UPLO const Uplo,
               int const N, traits::complex_d const& alpha,
               traits::complex_d const* A, int const lda,
               traits::complex_d const* X, int const incX,
               traits::complex_d const& beta,
               traits::complex_d* Y, int const incY)
    {
      cblas_zhemv (Order, Uplo, N,
                   static_cast<void const*> (&alpha),
                   static_cast<void const*> (A), lda,
                   static_cast<void const*> (X), incX,
                   static_cast<void const*> (&beta),
                   static_cast<void*> (Y), incY);
    }

    // y <- alpha * A * x + beta * y
    // A complex hermitian in packed form

    inline
    void hpmv (CBLAS_ORDER const Order, CBLAS_UPLO const Uplo,
               int const N, traits::complex_f const& alpha,
               traits::complex_f const* Ap,
               traits::complex_f const* X, int const incX,
               traits::complex_f const& beta,
               traits::complex_f* Y, int const incY)
    {
      cblas_chpmv (Order, Uplo, N,
                   static_cast<void const*> (&alpha),
                   static_cast<void const*> (Ap),
                   static_cast<void const*> (X), incX,
                   static_cast<void const*> (&beta),
                   static_cast<void*> (Y), incY);
    }

    inline
    void hpmv (CBLAS_ORDER const Order, CBLAS_UPLO const Uplo,
               int const N, traits::complex_d const& alpha,
               traits::complex_d const* Ap,
               traits::complex_d const* X, int const incX,
               traits::complex_d const& beta,
               traits::complex_d* Y, int const incY)
    {
      cblas_zhpmv (Order, Uplo, N,
                   static_cast<void const*> (&alpha),
                   static_cast<void const*> (Ap),
                   static_cast<void const*> (X), incX,
                   static_cast<void const*> (&beta),
                   static_cast<void*> (Y), incY);
    }

    // A <- alpha * x * y^T + A
    // .. real types:    calls cblas_xger
    // .. complex types: calls cblas_xgeru

    inline
    void ger (CBLAS_ORDER const Order, int const M, int const N,
              float const alpha, float const* X, int const incX,
              float const* Y, int const incY, float* A, int const lda)
    {
      cblas_sger (Order, M, N, alpha, X, incX, Y, incY, A, lda);
    }

    inline
    void ger (CBLAS_ORDER const Order, int const M, int const N,
              double const alpha, double const* X, int const incX,
              double const* Y, int const incY, double* A, int const lda)
    {
      cblas_dger (Order, M, N, alpha, X, incX, Y, incY, A, lda);
    }

    inline
    void ger (CBLAS_ORDER const Order, int const M, int const N,
              traits::complex_f const& alpha,
              traits::complex_f const* X, int const incX,
              traits::complex_f const* Y, int const incY,
              traits::complex_f* A, int const lda)
    {
      cblas_cgeru (Order, M, N,
                   static_cast<void const*> (&alpha),
                   static_cast<void const*> (X), incX,
                   static_cast<void const*> (Y), incY,
                   static_cast<void*> (A), lda);
    }

    inline
    void ger (CBLAS_ORDER const Order, int const M, int const N,
              traits::complex_d const& alpha,
              traits::complex_d const* X, int const incX,
              traits::complex_d const* Y, int const incY,
              traits::complex_d* A, int const lda)
    {
      cblas_zgeru (Order, M, N,
                   static_cast<void const*> (&alpha),
                   static_cast<void const*> (X), incX,
                   static_cast<void const*> (Y), incY,
                   static_cast<void*> (A), lda);
    }

    // A <- alpha * x * y^T + A
    // .. complex types only

    inline
    void geru (CBLAS_ORDER const Order, int const M, int const N,
               traits::complex_f const& alpha,
               traits::complex_f const* X, int const incX,
               traits::complex_f const* Y, int const incY,
               traits::complex_f* A, int const lda)
    {
      cblas_cgeru (Order, M, N,
                   static_cast<void const*> (&alpha),
                   static_cast<void const*> (X), incX,
                   static_cast<void const*> (Y), incY,
                   static_cast<void*> (A), lda);
    }

    inline
    void geru (CBLAS_ORDER const Order, int const M, int const N,
               traits::complex_d const& alpha,
               traits::complex_d const* X, int const incX,
               traits::complex_d const* Y, int const incY,
               traits::complex_d* A, int const lda)
    {
      cblas_zgeru (Order, M, N,
                   static_cast<void const*> (&alpha),
                   static_cast<void const*> (X), incX,
                   static_cast<void const*> (Y), incY,
                   static_cast<void*> (A), lda);
    }

    // A <- alpha * x * y^H + A
    // .. complex types only

    inline
    void gerc (CBLAS_ORDER const Order, int const M, int const N,
               traits::complex_f const& alpha,
               traits::complex_f const* X, int const incX,
               traits::complex_f const* Y, int const incY,
               traits::complex_f* A, int const lda)
    {
      cblas_cgerc (Order, M, N,
                   static_cast<void const*> (&alpha),
                   static_cast<void const*> (X), incX,
                   static_cast<void const*> (Y), incY,
                   static_cast<void*> (A), lda);
    }

    inline
    void gerc (CBLAS_ORDER const Order, int const M, int const N,
               traits::complex_d const& alpha,
               traits::complex_d const* X, int const incX,
               traits::complex_d const* Y, int const incY,
               traits::complex_d* A, int const lda)
    {
      cblas_zgerc (Order, M, N,
                   static_cast<void const*> (&alpha),
                   static_cast<void const*> (X), incX,
                   static_cast<void const*> (Y), incY,
                   static_cast<void*> (A), lda);
    }

    // A <- alpha * x * x^T + A
    // A real symmetric

    inline
    void syr (CBLAS_ORDER const Order, CBLAS_UPLO const Uplo,
              int const N, float const alpha,
              float const* X, int const incX, float* A, int const lda)
    {
      cblas_ssyr (Order, Uplo, N, alpha, X, incX, A, lda);
    }

    inline
    void syr (CBLAS_ORDER const Order, CBLAS_UPLO const Uplo,
              int const N, double const alpha,
              double const* X, int const incX, double* A, int const lda)
    {
      cblas_dsyr (Order, Uplo, N, alpha, X, incX, A, lda);
    }

    // A <- alpha * x * x^T + A
    // A real symmetric in packed form

    inline
    void spr (CBLAS_ORDER const Order, CBLAS_UPLO const Uplo,
              int const N, float const alpha,
              float const* X, int const incX, float* Ap)
    {
      cblas_sspr (Order, Uplo, N, alpha, X, incX, Ap);
    }

    inline
    void spr (CBLAS_ORDER const Order, CBLAS_UPLO const Uplo,
              int const N, double const alpha,
              double const* X, int const incX, double* Ap)
    {
      cblas_dspr (Order, Uplo, N, alpha, X, incX, Ap);
    }

    // A <- alpha * x * y^T + alpha * y * x^T + A
    // A real symmetric

    inline
    void syr2 (CBLAS_ORDER const Order, CBLAS_UPLO const Uplo,
               int const N, float const alpha,
               float const* X, int const incX,
               float const* Y, int const incY,
               float* A, int const lda)
    {
      cblas_ssyr2 (Order, Uplo, N, alpha, X, incX, Y, incY, A, lda);
    }

    inline
    void syr2 (CBLAS_ORDER const Order, CBLAS_UPLO const Uplo,
               int const N, double const alpha,
               double const* X, int const incX,
               double const* Y, int const incY,
               double* A, int const lda)
    {
      cblas_dsyr2 (Order, Uplo, N, alpha, X, incX, Y, incY, A, lda);
    }

    // A <- alpha * x * y^T + alpha * y * x^T + A
    // A real symmetric in packed form

    inline
    void spr2 (CBLAS_ORDER const Order, CBLAS_UPLO const Uplo,
               int const N, float const alpha,
               float const* X, int const incX,
               float const* Y, int const incY, float* Ap)
    {
      cblas_sspr2 (Order, Uplo, N, alpha, X, incX, Y, incY, Ap);
    }

    inline
    void spr2 (CBLAS_ORDER const Order, CBLAS_UPLO const Uplo,
               int const N, double const alpha,
               double const* X, int const incX,
               double const* Y, int const incY, double* Ap)
    {
      cblas_dspr2 (Order, Uplo, N, alpha, X, incX, Y, incY, Ap);
    }

    // A <- alpha * x * x^H + A
    // A hermitian

    inline
    void her (CBLAS_ORDER const Order, CBLAS_UPLO const Uplo,
              int const N, float const alpha,
              traits::complex_f const* X, int const incX,
              traits::complex_f* A, int const lda)
    {
      cblas_cher (Order, Uplo, N, alpha,
                  static_cast<void const*> (X), incX,
                  static_cast<void*> (A), lda);
    }

    inline
    void her (CBLAS_ORDER const Order, CBLAS_UPLO const Uplo,
              int const N, double const alpha,
              traits::complex_d const* X, int const incX,
              traits::complex_d* A, int const lda)
    {
      cblas_zher (Order, Uplo, N, alpha,
                  static_cast<void const*> (X), incX,
                  static_cast<void*> (A), lda);
    }

    // A <- alpha * x * x^H + A
    // A hermitian in packed form

    inline
    void hpr (CBLAS_ORDER const Order, CBLAS_UPLO const Uplo,
              int const N, float const alpha,
              traits::complex_f const* X, int const incX,
              traits::complex_f* Ap)
    {
      cblas_chpr (Order, Uplo, N, alpha,
                  static_cast<void const*> (X), incX,
                  static_cast<void*> (Ap));
    }

    inline
    void hpr (CBLAS_ORDER const Order, CBLAS_UPLO const Uplo,
              int const N, double const alpha,
              traits::complex_d const* X, int const incX,
              traits::complex_d* Ap)
    {
      cblas_zhpr (Order, Uplo, N, alpha,
                  static_cast<void const*> (X), incX,
                  static_cast<void*> (Ap));
    }

    // A <- alpha * x * y^H + y * (alpha * x)^H + A
    // A hermitian

    inline
    void her2 (CBLAS_ORDER const Order, CBLAS_UPLO const Uplo,
               int const N, traits::complex_f const& alpha,
               traits::complex_f const* X, int const incX,
               traits::complex_f const* Y, int const incY,
               traits::complex_f* A, int const lda)
    {
      cblas_cher2 (Order, Uplo, N,
                   static_cast<void const*> (&alpha),
                   static_cast<void const*> (X), incX,
                   static_cast<void const*> (Y), incY,
                   static_cast<void*> (A), lda);
    }

    inline
    void her2 (CBLAS_ORDER const Order, CBLAS_UPLO const Uplo,
               int const N, traits::complex_d const& alpha,
               traits::complex_d const* X, int const incX,
               traits::complex_d const* Y, int const incY,
               traits::complex_d* A, int const lda)
    {
      cblas_zher2 (Order, Uplo, N,
                   static_cast<void const*> (&alpha),
                   static_cast<void const*> (X), incX,
                   static_cast<void const*> (Y), incY,
                   static_cast<void*> (A), lda);
    }

    // A <- alpha * x * y^H + y * (alpha * x)^H + A
    // A hermitian in packed form

    inline
    void hpr2 (CBLAS_ORDER const Order, CBLAS_UPLO const Uplo,
               int const N, traits::complex_f const& alpha,
               traits::complex_f const* X, int const incX,
               traits::complex_f const* Y, int const incY,
               traits::complex_f* Ap)
    {
      cblas_chpr2 (Order, Uplo, N,
                   static_cast<void const*> (&alpha),
                   static_cast<void const*> (X), incX,
                   static_cast<void const*> (Y), incY,
                   static_cast<void*> (Ap));
    }

    inline
    void hpr2 (CBLAS_ORDER const Order, CBLAS_UPLO const Uplo,
               int const N, traits::complex_d const& alpha,
               traits::complex_d const* X, int const incX,
               traits::complex_d const* Y, int const incY,
               traits::complex_d* Ap)
    {
      cblas_zhpr2 (Order, Uplo, N,
                   static_cast<void const*> (&alpha),
                   static_cast<void const*> (X), incX,
                   static_cast<void const*> (Y), incY,
                   static_cast<void*> (Ap));
    }

  }} // namepaces detail & atlas

}}}

#endif // BOOST_NUMERIC_BINDINGS_CBLAS2_OVERLOADS_HPP


/**
 * @author Tomas Polasek, David Hrusa
 * @date 1.14.2020
 * @version 1.0
 * @brief gui rendering and helper classes.
 */

#include "TreeStats.h"

#include <queue>

#include <boost/random.hpp>
#include <boost/math/distributions.hpp>

#include "TreeScene.h"
#include "TreeRenderSystemRT.h"

namespace treestat
{

namespace impl
{

/// @brief Internal implementation of the DistributionEngine.
struct DistributionEngineImpl
{
    /// @brief Randomness engine used.
    using EngineT = boost::mt19937;

    /// @brief Type used for variate generators.
    template <typename DistT>
    using VariateGenerator = boost::variate_generator<EngineT, DistT>;

    /// Instance of the engine.
    EngineT engine{ };
}; // struct DistributionEngineImpl

/// @brief Normal Gaussian distribution.
struct NormalDistribution : public DistributionProperties
{
    /// Unique identifier of this distribution.
    static constexpr auto IDENTIFIER{ "Normal" };
    /// @brief Type used for sampling the distribution.
    using DistT = boost::math::normal_distribution<RealType>;

    /// @brief Initialize normal distribution with mean and standard deviation.
    NormalDistribution(DistributionEngineImpl &engine,
        RealType mean = 0, RealType sigma = 1);

    /// @brief Initialize normal distribution from serialized representation.
    NormalDistribution(DistributionEngineImpl &engine,
        const std::string &serialized);

    // Implement interface:
    virtual std::string serialize() const override final;
    virtual void deserialize(const std::string &serialized) override final;
    virtual void setEngine(DistributionEngineImpl &engine) override final;
    virtual RealType cdf(const RealType &x) const override final;
    virtual RealType sample() override final;

    /// @brief Intitialize the internal structure using provided parameters.
    void initialize(RealType mean, RealType sigma);

    /// Distribution implementation.
    DistT distribution{ };
    /// Sample generator for the distribution.
    DistributionEngineImpl::EngineT generator{ };
}; // struct NormalDistribution

NormalDistribution::NormalDistribution(DistributionEngineImpl &engine,
    RealType mean, RealType sigma) :
    generator{ engine.engine }
{ initialize(mean, sigma); }

NormalDistribution::NormalDistribution(DistributionEngineImpl &engine,
    const std::string &serialized) :
    generator{ engine.engine }
{ deserialize(serialized); }

std::string NormalDistribution::serialize() const
{
    return treeio::json{
        { "distribution", IDENTIFIER },
        { "mean", distribution.mean() },
        { "sigma", distribution.standard_deviation() },
    }.dump();
}

void NormalDistribution::deserialize(const std::string &serialized)
{
    const auto data{ treeio::json::parse(serialized) };

    if (data.at("distribution").get<std::string>() != IDENTIFIER)
    { throw StochasticException("Deserializing with unknown distribution identifier!"); }

    initialize(
        data.at("mean").get<RealType>(),
        data.at("sigma").get<RealType>()
    );
}

void NormalDistribution::setEngine(DistributionEngineImpl &engine)
{ generator = engine.engine; }

NormalDistribution::RealType NormalDistribution::cdf(const RealType &x) const
{ return boost::math::cdf<RealType>(distribution, x); }

NormalDistribution::RealType NormalDistribution::sample()
{ return boost::math::quantile<RealType>(distribution, generator()); }

void NormalDistribution::initialize(RealType mean, RealType sigma)
{ distribution = DistT{ mean, sigma }; }

std::shared_ptr<DistributionProperties>
    DistributionProperties::normalDistribution(DistributionEngineImpl &engine,
        RealType mean, RealType sigma)
{ return std::make_shared<NormalDistribution>(engine, mean, sigma); }

std::shared_ptr<DistributionProperties>
    DistributionProperties::deserializeDistribution(
        DistributionEngineImpl &engine, const std::string &serialized)
{
    const auto data{ treeio::json::parse(serialized) };
    const auto name{ data.at("distribution").get<std::string>() };

    if (name == NormalDistribution::IDENTIFIER)
    { return std::make_shared<NormalDistribution>(engine, serialized); }
    else
    { throw StochasticException("Unknown distribution is being deserialized!"); }
}

} // namespace impl

StochasticEngine::StochasticEngine(uint32_t seed)
{ reset(seed); }
StochasticEngine::~StochasticEngine()
{ /* Automatic */ }

void StochasticEngine::reset(uint32_t seed)
{
    if (!mImpl)
    { mImpl = std::make_shared<impl::DistributionEngineImpl>(); }

    if (seed)
    { mImpl->engine.seed(seed); }
    else
    { mImpl->engine.seed(); }
}

std::string ImageData::valueTypeToStr(const ValueType &type)
{
    switch (type)
    {
        default:
        case ValueType::UInt:
        { return "UInt"; }
        case ValueType::Float:
        { return "Float"; }
    }
}

void ImageData::saveTo(treeio::json &out) const
{
    out["image"] = {
        { "width", width },
        { "height", height },
        { "channels", channels },
        { "valueType", valueTypeToStr(valueType) },
        { "data", data.empty() ? std::string{ } : treeutil::encodeBinaryJSON(data) },
    };
}

std::size_t branchingToOrdinal(const Branching &branching)
{
    switch (branching)
    {
        default:
        case Branching::Monopodial:
        { return 0u; }
        case Branching::SympodialMonochasial:
        { return 1u; }
        case Branching::SympodialDichasial:
        { return 2u; }
    }
}

Branching ordinalToBranching(std::size_t ordinal)
{
    switch (ordinal)
    {
        default:
        case 0u:
        { return Branching::Monopodial; }
        case 1u:
        { return Branching::SympodialMonochasial; }
        case 2u:
        { return Branching::SympodialDichasial; }
    }
}

std::size_t ramificationToOrdinal(const Ramification &ramification)
{
    switch (ramification)
    {
        default:
        case Ramification::Continuous:
        { return 0u; }
        case Ramification::Rhythmic:
        { return 1u; }
        case Ramification::Diffuse:
        { return 2u; }
    }
}

Ramification ordinalToRamification(std::size_t ordinal)
{
    switch (ordinal)
    {
        default:
        case 0u:
        { return Ramification::Continuous; }
        case 1u:
        { return Ramification::Rhythmic; }
        case 2u:
        { return Ramification::Diffuse; }
    }
}

void TreeStatValues::saveTo(treeio::json &out) const
{
    forEach([&out] (const auto &var)
    {
        auto &varOut{ out[var.name] };
        varOut["name"] = var.name;
        varOut["description"] = var.description;

        var.var.saveTo(varOut);
        var.hist.saveTo(varOut);
    });

    auto &commonOut{ out["common"] };
    commonOut = {
        { "treeAge", treeAge },
        { "trunkLength", trunkLength },
        { "interNodeLength", interNodeLengthEstimate },
        { "interNodeVariance", interNodeLengthVariance },
        { "leafCount", leafCount },
        { "branching", branchingToOrdinal(branching) },
        { "ramification", ramificationToOrdinal(ramification) },
    };

    auto &visualOut{ out["visual"] };
    shadowImprint.saveTo(visualOut["shadowImprint"]);
}

TreeStats::TreeStats()
{ /* Automatic */ }

TreeStats::TreeStats(const treeio::ArrayTree &tree)
{ calculateStatistics(tree); }

TreeStats::~TreeStats()
{ /* Automatic */ }

void printChains(std::size_t currentChainIdx, const treeutil::TreeChains &chains, std::size_t maxDepth,
    const std::string &indent = "", std::size_t currentDepth = 0u)
{
    const auto &currentChain{ chains.chains()[currentChainIdx] };

    const auto newline{ std::string{ "\n" } + indent + "\t" };
    Info << indent << "[ Chain " << currentChainIdx
        << newline << "depth: " << currentDepth
        << "; nodeCount: " << currentChain.nodes.size()
        << "; gOrder: " << currentChain.graveliusOrder
        << "; gDepth: " << currentChain.graveliusDepth
        << newline << "children (" << currentChain.childChains.size() << "): [ \n";

    if (currentDepth + 1u < maxDepth)
    {
        if (currentChain.childChains.empty())
        { Info << indent << "\tNo Child Chains"; }
        else
        {
            for (const auto &childChainIdx: currentChain.childChains)
            { printChains(childChainIdx, chains, maxDepth, indent + "\t", currentDepth + 1u); Info << "\n"; }
        }
    }
    else
    { Info << indent << "\tMaximum Depth Reached"; }

    Info << newline << " ]"
        << "\n" << indent << " ]";

    if (currentDepth == 0u)
    { Info << std::endl; }
}

void TreeStats::calculateStatistics(const treeio::ArrayTree &tree)
{
    treeutil::Timer profilingTimer{ };

    /* Phase 0 Initialize structures */
    mStats = { };
    if (tree.empty()) { return; }

    /* Phase 1 Prepare helper tree data */
    const auto [chains, chainDataStorage]{ prepareTree(mStats, mSettings, tree) };
    if (chainDataStorage.empty()) { return; }

    const auto timePreparation{ profilingTimer.reset() };

    /* Phase 2 Calculate basic statistics */
    calculateBasicStats(mStats, mSettings, chains, chainDataStorage);

    const auto timeBasicStats{ profilingTimer.reset() };

    /* Phase 3 Calculate visual statistics */
    calculateVisualStats(mStats, mSettings, tree);

    const auto timeVisualStats{ profilingTimer.reset() };

    /* Phase 4 Calculate derived statistics */
    calculateDerivedStats(mStats, mSettings, chains, chainDataStorage);

    const auto timeDerivedStats{ profilingTimer.reset() };

    /* Phase 5 Calculate growth statistics */
    calculateGrowthStats(mStats, mSettings, chains, chainDataStorage);

    const auto timeGrowthStats{ profilingTimer.reset() };

    /* Phase 6 Finalize and clean up */
    finalizeStats(mStats, mSettings);

    const auto timeFinalizeStats{ profilingTimer.reset() };

    const auto timeTotalNoVisual{
        timePreparation + timeBasicStats +
        timeDerivedStats + timeGrowthStats +
        timeFinalizeStats
    };
    const auto timeTotal{ timeTotalNoVisual + timeVisualStats };

    Info << "[Prof] timePreparation: " << timePreparation << std::endl;
    Info << "[Prof] timeBasicStats: " << timeBasicStats << std::endl;
    Info << "[Prof] timeVisualStats: " << timeVisualStats << std::endl;
    Info << "[Prof] timeDerivedStats: " << timeDerivedStats << std::endl;
    Info << "[Prof] timeGrowthStats: " << timeGrowthStats << std::endl;
    Info << "[Prof] timeFinalizeStats: " << timeFinalizeStats << std::endl;
    Info << "[Prof] timeTotalNoVisual: " << timeTotalNoVisual << std::endl;
    Info << "[Prof] timeTotal: " << timeTotal << std::endl;
}

void TreeStats::saveStatisticsToDynamic(treeio::ArrayTree &tree)
{
    auto &dynamic{ tree.metaData().dynamicData() };

    auto &statsOut{ dynamic["stats"] };
    mStats.saveTo(statsOut);
}

std::pair<treeutil::TreeChains, std::vector<TreeStats::ChainData>>
    TreeStats::prepareTree(TreeStatValues &stats, const StatsSettings &settings,
        const treeio::ArrayTree &tree)
{
    treeutil::TreeChains chains{ tree };

    //printChains(0u, chains, 3u);

    if (chains.chains().empty())
    { return { }; }

    std::vector<ChainData> chainDataStorage{ };
    chainDataStorage.resize(chains.chains().size());

    // Calculate forward pass information - root -> leaves.
    prepareTreeForwardPass(chains, chainDataStorage, stats, settings);

    // Calculate backward pass information - leaves -> root.
    prepareTreeBackwardPass(chains, chainDataStorage, stats, settings);

    return { chains, chainDataStorage };
}

void TreeStats::prepareTreeForwardPass(treeutil::TreeChains &chains, std::vector<ChainData> &chainDataStorage,
    TreeStatValues &stats, const StatsSettings &settings)
{
    const auto &chainsTree{ chains.internalTree() };

    for (std::size_t chainIdx = 0u; chainIdx < chains.chains().size(); ++chainIdx)
    { // Go through chains and pre-calculate data.
        const auto &chain{ chains.chains()[chainIdx] };
        auto &chainData{ chainDataStorage[chainIdx] };

        float length{ 0.0f };
        float volume{ 0.0f };
        float angleSum{ 0.0f };
        float angleAbsSum{ 0.0f };
        float lengthAverageThickness{ 0.0f };
        Vector3D startToEnd{ };
        if (chain.nodes.size() > 1u)
        {
            const auto &firstNodeData{ chainsTree.getNode(chain.nodes[0u]).data() };
            const auto &secondNodeData{ chainsTree.getNode(chain.nodes[1u]).data() };
            const auto &lastNodeData{ chainsTree.getNode(chain.nodes.back()).data() };

            chainData.segmentData.reserve(chain.nodes.size() - 1u);
            auto lastNodePos{ firstNodeData.pos };
            auto lastNodeThickness{ firstNodeData.thickness };
            startToEnd = (lastNodeData.pos - lastNodePos);
            auto lastToCurrent{ (secondNodeData.pos - lastNodePos).normalized() };

            chainData.minThickness = std::min<float>(
                chainData.minThickness, firstNodeData.thickness);
            chainData.maxThickness = std::max<float>(
                chainData.maxThickness, firstNodeData.thickness);
            chainData.firstDirection = lastToCurrent;

            for (auto it = chain.nodes.begin() + 1u; it != chain.nodes.end(); ++it)
            {
                const auto currentNodePos{ chainsTree.getNode(*it).data().pos };
                const auto currentDirection{ (currentNodePos - lastNodePos).normalized() };
                const auto lastToCurrentAngle{
                    treeutil::angleBetweenNormVectorsRad<float>(lastToCurrent, currentDirection)
                };

                const auto segmentLength{ lastNodePos.distanceTo(currentNodePos) };
                length += segmentLength;
                angleSum += lastToCurrentAngle;
                angleAbsSum += std::abs(lastToCurrentAngle);

                auto thickness{ chainsTree.getNode(*it).data().thickness };
                if (std::isnan(thickness) || std::isinf(thickness))
                { thickness = 0.0f; }

                volume += treeutil::circularConeFrustumVolume(
                    segmentLength, lastNodeThickness, thickness);

                lengthAverageThickness += segmentLength * thickness;
                stats.segmentThickness.observeSample(thickness);
                chainData.minThickness = std::min<float>(
                    chainData.minThickness, thickness);
                chainData.maxThickness = std::max<float>(
                    chainData.maxThickness, thickness);

                chainData.segmentData.emplace_back(SegmentData{
                    segmentLength, lastNodeThickness, thickness
                });

                lastNodePos = currentNodePos;
                lastNodeThickness = thickness;
                lastToCurrent = currentDirection;

                // Try to find a better first direction, when first few nodes are at the exact same position.
                if (!treeutil::aboveEpsilon(chainData.firstDirection.length()))
                { chainData.firstDirection = lastToCurrent; }
            }

            chainData.lastDirection = lastToCurrent;
        }

        chainData.length = length;
        chainData.angleSum = angleSum;
        chainData.angleAbsSum = angleAbsSum;
        chainData.lengthAverageThickness = lengthAverageThickness;
        chainData.volume = volume;
        chainData.startToEnd = startToEnd;
    }
}

void TreeStats::prepareTreeBackwardPass(treeutil::TreeChains &chains, std::vector<ChainData> &chainDataStorage,
    TreeStatValues &stats, const StatsSettings &settings)
{
    using ChainIdxT = treeutil::TreeChains::NodeChain::ChainIdxT;
    static constexpr auto INVALID_CHAIN_IDX{ treeutil::TreeChains::NodeChain::INVALID_CHAIN_IDX };

    std::vector<std::size_t> finalizedChildChains{ };
    finalizedChildChains.resize(chains.chains().size(), 0u);
    std::queue<ChainIdxT> toProcess{ };
    for (const auto &leafChainIdx : chains.leafChains())
    { toProcess.push(leafChainIdx); }

    while (!toProcess.empty())
    { // Process all chains from leaves to root.
        const auto currentChainIdx{ toProcess.front() }; toProcess.pop();
        auto &currentChain{ chains.chains()[currentChainIdx] };
        auto &currentChainData{ chainDataStorage[currentChainIdx] };

        for (const auto &ccIdx : currentChain.childChains)
        { // Aggregate data from child chains.
            const auto &childChain{ chains.chains()[ccIdx] };
            const auto &childChainData{ chainDataStorage[ccIdx] };

            currentChainData.subtreeLength += childChainData.subtreeLength;
            currentChainData.subtreeVolume += childChainData.subtreeVolume;
            currentChainData.subtreeLeafCount += childChainData.subtreeLeafCount;
        }

        // Set fixed value for leaf chains.
        if (currentChain.childChains.empty())
        {
            currentChainData.subtreeLeafCount = 1u;
            currentChainData.subtreeLength = currentChainData.length;
            currentChainData.subtreeVolume = currentChainData.volume;
        }

        // Check parent node finalization..
        const auto parentChainIdx{ currentChain.parentChain };
        if (parentChainIdx != INVALID_CHAIN_IDX)
        { // We have parent -> Not currently at the root node.
            const auto &parentChain{ chains.chains()[parentChainIdx] };
            finalizedChildChains[parentChainIdx]++;

            // Check if all of the parents children have necessary calculations finished.
            if (finalizedChildChains[parentChainIdx] >= parentChain.childChains.size())
            { // All children are finished -> Add parent for processing.
                toProcess.push(parentChainIdx);
            }
        }
    }
}

void TreeStats::calculateBasicStats(TreeStatValues &stats, const StatsSettings &settings,
    const treeutil::TreeChains &chains, const std::vector<ChainData> &chainDataStorage)
{
    for (std::size_t chainIdx = 0u; chainIdx < chains.chains().size(); ++chainIdx)
    { // Go through the tree using non-branching sequences of segments -> chains.
        const auto &chain{ chains.chains()[chainIdx] };
        const auto &chainData{ chainDataStorage[chainIdx] };

        const auto segmentCount{ static_cast<uint32_t>(chain.nodes.size() - 1u) };
        const auto chainDepth{ static_cast<uint32_t>(chain.chainDepth) };

        stats.segmentsPerChain.observeSample(segmentCount);
        stats.chainsPerDepth.observeSample(chainDepth);
        if (treeutil::aboveEpsilon(chainData.length))
        { stats.chainLength.observeSample(chainData.length); }
        if (treeutil::aboveEpsilon(chainData.startToEnd.length()))
        { stats.chainTotalLength.observeSample(chainData.startToEnd.length()); }
        stats.chainDeformation.observeSample(chainData.angleAbsSum);

        for (const auto &segment : chainData.segmentData)
        { // Approximate volume of each segment with circular cone frustum volume.
            const auto volume{
                treeutil::circularConeFrustumVolume(
                    segment.length, segment.startThickness, segment.endThickness)
            };
            stats.segmentVolume.observeSample({ segment.startThickness, volume });
            stats.segmentVolume.observeSample({ segment.endThickness, volume });
        }

        if (chain.nodes.size() > 2u && treeutil::aboveEpsilon(chainData.startToEnd.length()) &&
            treeutil::aboveEpsilon(chainData.length))
        { stats.chainStraightness.observeSample(chainData.startToEnd.length() / chainData.length); }
        if (treeutil::aboveEpsilon(chainData.startToEnd.length()))
        {
            stats.chainSlope.observeSample(
                treeutil::angleBetweenVectorsRad<float>(
                    chainData.startToEnd, Vector3D{ 0.0f, 1.0f, 0.0f }
                )
            );
        }

        if (chainData.minThickness < std::numeric_limits<float>::max())
        { stats.chainMinThickness.observeSample(chainData.minThickness); }
        if (chainData.maxThickness > std::numeric_limits<float>::min())
        { stats.chainMaxThickness.observeSample(chainData.maxThickness); }
        if (chainData.minThickness < std::numeric_limits<float>::max() &&
            chainData.maxThickness > std::numeric_limits<float>::min())
        { stats.chainMinMaxThicknessRatio.observeSample(chainData.minThickness / chainData.maxThickness); }

        if (chain.parentChain != chain.INVALID_CHAIN_IDX)
        { // Calculate child-parent statistics.
            const auto &parentChain{ chains.chains()[chain.parentChain] };
            const auto &parentChainData{ chainDataStorage[chain.parentChain] };

            const auto parentChildAngle{ treeutil::angleBetweenVectorsRad<float>(
                parentChainData.startToEnd, chainData.startToEnd
            ) };

            if (treeutil::aboveEpsilon(chainData.length) && treeutil::aboveEpsilon(parentChainData.length))
            { stats.chainLengthRatio.observeSample(chainData.length / parentChainData.length); }
            if (treeutil::aboveEpsilon(chainData.angleSum) && treeutil::aboveEpsilon(parentChainData.angleSum))
            { stats.chainAngleSumDelta.observeSample(chainData.angleSum - parentChainData.angleSum); }
            stats.chainParentChildAngle.observeSample(parentChildAngle);

            const auto tropismBranch{ Vector3D::dotProduct(
                (chainData.startToEnd - chainData.firstDirection).normalized(),
                StatsSettings::DEFAULT_HORIZON_UP_DIRECTION.normalized()
            ) };
            stats.chainTropismBranch.observeSample(tropismBranch);

            const auto tropismHorizon{ Vector3D::dotProduct(
                chainData.firstDirection.normalized(),
                StatsSettings::DEFAULT_HORIZON_UP_DIRECTION.normalized()
            ) };
            stats.chainTropismHorizon.observeSample(tropismHorizon);

            const auto tropismParent{ Vector3D::dotProduct(
                    parentChainData.lastDirection.normalized(),
                    chainData.startToEnd.normalized()
            ) };
            stats.chainTropismParent.observeSample(tropismParent);
        }
        if (chain.childChains.size() > 1u)
        { // Calculate child statistics.
            auto minChildAngle{ std::numeric_limits<float>::max() };
            auto maxChildAngle{ std::numeric_limits<float>::lowest() };
            auto minChildAngleTotal{ std::numeric_limits<float>::max() };
            auto maxChildAngleTotal{ std::numeric_limits<float>::lowest() };

            auto minParentAngle{ std::numeric_limits<float>::max() };
            auto maxParentAngle{ std::numeric_limits<float>::lowest() };
            auto minParentAngleTotal{ std::numeric_limits<float>::max() };
            auto maxParentAngleTotal{ std::numeric_limits<float>::lowest() };

            for (std::size_t idx1 = 0u; idx1 < chain.childChains.size(); ++idx1)
            {
                const auto &cChainIdx1{ chain.childChains[idx1] };
                const auto &cChainData1{ chainDataStorage[cChainIdx1] };

                auto minSiblingAngle{ std::numeric_limits<float>::max() };
                auto maxSiblingAngle{ std::numeric_limits<float>::lowest() };
                auto minSiblingAngleTotal{ std::numeric_limits<float>::max() };
                auto maxSiblingAngleTotal{ std::numeric_limits<float>::lowest() };

                const auto parentChainAngle{ treeutil::angleBetweenVectorsRad<float>(
                    chainData.lastDirection, cChainData1.firstDirection)};
                const auto parentChainAngleTotal{ treeutil::angleBetweenVectorsRad<float>(
                    chainData.startToEnd, cChainData1.startToEnd)};

                if (treeutil::aboveEpsilon(cChainData1.firstDirection.length()) &&
                    treeutil::aboveEpsilon(chainData.lastDirection.length()))
                {
                    stats.chainParentAngle.observeSample(parentChainAngle);
                    minParentAngle = std::min<float>(minParentAngle, parentChainAngle);
                    maxParentAngle = std::max<float>(maxChildAngle, parentChainAngle);
                }

                if (treeutil::aboveEpsilon(cChainData1.startToEnd.length()) &&
                    treeutil::aboveEpsilon(chainData.startToEnd.length()))
                {
                    stats.chainParentAngleTotal.observeSample(parentChainAngleTotal);
                    minParentAngleTotal = std::min<float>(minParentAngleTotal, parentChainAngleTotal);
                    maxParentAngleTotal = std::max<float>(maxParentAngleTotal, parentChainAngleTotal);
                }

                for (std::size_t idx2 = 0u; idx2 < chain.childChains.size(); ++idx2)
                {
                    const auto &cChainIdx2{ chain.childChains[idx2] };
                    const auto &cChainData2{ chainDataStorage[cChainIdx2] };

                    const auto childChainAngle{ treeutil::angleBetweenVectorsRad<float>(
                        cChainData1.firstDirection, cChainData2.firstDirection)};
                    const auto childChainAngleTotal{ treeutil::angleBetweenVectorsRad<float>(
                        cChainData1.startToEnd, cChainData2.startToEnd)};

                    if (treeutil::aboveEpsilon(cChainData1.firstDirection.length()) &&
                        treeutil::aboveEpsilon(cChainData2.firstDirection.length()))
                    {
                        stats.chainSiblingAngle.observeSample(childChainAngle);
                        minSiblingAngle = std::min<float>(minSiblingAngle, childChainAngle);
                        maxSiblingAngle = std::max<float>(maxSiblingAngle, childChainAngle);
                    }

                    if (treeutil::aboveEpsilon(cChainData1.startToEnd.length()) &&
                        treeutil::aboveEpsilon(cChainData2.startToEnd.length()))
                    {
                        stats.chainSiblingAngleTotal.observeSample(childChainAngleTotal);
                        minSiblingAngleTotal = std::min<float>(minSiblingAngleTotal, childChainAngleTotal);
                        maxSiblingAngleTotal = std::max<float>(maxSiblingAngleTotal, childChainAngleTotal);
                    }

                    if (treeutil::aboveEpsilon(chainData.subtreeLeafCount))
                    {
                        const auto leafAsymmetry{
                            (cChainData1.subtreeLeafCount - cChainData2.subtreeLeafCount) /
                            static_cast<float>(chainData.subtreeLeafCount)
                        };
                        stats.chainAsymmetryLeaf.observeSample(leafAsymmetry);
                    }

                    if (treeutil::aboveEpsilon(chainData.subtreeLength))
                    {
                        const auto lengthAsymmetry{
                            (cChainData1.subtreeLength - cChainData2.subtreeLength) /
                            static_cast<float>(chainData.subtreeLength)
                        };
                        stats.chainAsymmetryLength.observeSample(lengthAsymmetry);
                    }

                    if (treeutil::aboveEpsilon(chainData.subtreeVolume))
                    {
                        const auto volumeAsymmetry{
                            (cChainData1.subtreeVolume - cChainData2.subtreeVolume) /
                            static_cast<float>(chainData.subtreeVolume)
                        };
                        stats.chainAsymmetryVolume.observeSample(volumeAsymmetry);
                    }
                }
                if (minSiblingAngle < std::numeric_limits<float>::max())
                {
                    stats.chainMinSiblingAngle.observeSample(minSiblingAngle);
                    minChildAngle = std::min<float>(minChildAngle, minSiblingAngle);
                }
                if (maxSiblingAngle > std::numeric_limits<float>::lowest())
                {
                    stats.chainMaxSiblingAngle.observeSample(maxSiblingAngle);
                    maxChildAngle = std::max<float>(maxChildAngle, maxSiblingAngle);
                }
                if (minSiblingAngleTotal < std::numeric_limits<float>::max())
                {
                    stats.chainMinSiblingAngleTotal.observeSample(minSiblingAngleTotal);
                    minChildAngleTotal = std::min<float>(minChildAngleTotal, minSiblingAngleTotal);
                }
                if (maxSiblingAngleTotal > std::numeric_limits<float>::lowest())
                {
                    stats.chainMaxSiblingAngleTotal.observeSample(maxSiblingAngleTotal);
                    maxChildAngleTotal = std::max<float>(maxChildAngleTotal, maxSiblingAngleTotal);
                }
            }

            if (minChildAngle < std::numeric_limits<float>::max())
            { stats.chainMinChildAngle.observeSample(minChildAngle); }
            if (maxChildAngle > std::numeric_limits<float>::lowest())
            { stats.chainMaxChildAngle.observeSample(maxChildAngle); }
            if (minChildAngleTotal < std::numeric_limits<float>::max())
            { stats.chainMinChildAngleTotal.observeSample(minChildAngleTotal); }
            if (maxChildAngleTotal > std::numeric_limits<float>::lowest())
            { stats.chainMaxChildAngleTotal.observeSample(maxChildAngleTotal); }

            if (minParentAngle < std::numeric_limits<float>::max())
            { stats.chainMinParentAngle.observeSample(minParentAngle); }
            if (maxParentAngle > std::numeric_limits<float>::lowest())
            { stats.chainMaxParentAngle.observeSample(maxParentAngle); }
            if (minParentAngleTotal < std::numeric_limits<float>::max())
            { stats.chainMinParentAngleTotal.observeSample(minParentAngleTotal); }
            if (maxParentAngleTotal > std::numeric_limits<float>::lowest())
            { stats.chainMaxParentAngleTotal.observeSample(maxParentAngleTotal); }
        }
    }
}

void TreeStats::calculateGrowthStats(TreeStatValues &stats, const StatsSettings &settings,
    const treeutil::TreeChains &chains, const std::vector<ChainData> &chainDataStorage)
{
    const auto branchingPointDelta{ stats.interNodeLengthEstimate * settings.growthBranchingPointDelta };
    const auto compactChains{ chains.generateCompactChains(branchingPointDelta) };
    const auto &chainsTree{ chains.internalTree() };
    const auto &fullChains{ chains.chains() };

    for (const auto &compactChain : compactChains)
    { // Pre-calculate statistics on the compact chains, skipping leaves.
        if (compactChain.childChains.empty())
        { continue; }

        const auto maxCompactedChainLength{ calculateMaxCompactChainLength(
            fullChains, chainsTree, compactChain) };
        const auto internodesInChain{ maxCompactedChainLength / stats.interNodeLengthEstimate };
        stats.chainInternodeCount.observeSample(internodesInChain);
    }

    // Calculate required statistics from the first run.
    const auto averageInternodeStats{ stats.chainInternodeCount.var.properties() };

    for (const auto &compactChain : compactChains)
    { // Go through all compacted chains, skipping leaves.
        if (compactChain.childChains.empty())
        { continue; }
        if (compactChain.childChains.size() == 1u)
        { Warning << "Found compact chain with one child chain!" << std::endl; }

        std::vector<float> branchContinuationWeights{ };
        std::vector<float> branchSizeWeights{ };

        for (const auto &childChainRec : compactChain.childChains)
        { // Calculate weights for child chains.
            const auto parentChainInfo{ compactChain.compactedChains[childChainRec.originIdx] };
            const auto parentChainIdx{ parentChainInfo.chainIdx };
            const auto &parentChain{ fullChains[parentChainIdx] };
            const auto &parentChainData{ chainDataStorage[parentChainIdx] };

            const auto childChainIdx{ childChainRec.chainIdx };
            const auto &childChain{ fullChains[childChainIdx] };
            const auto &childChainData{ chainDataStorage[childChainIdx] };

            const auto childContinuationWeight{ Vector3D::dotProduct(
                parentChainData.lastDirection, childChainData.firstDirection) };
            branchContinuationWeights.push_back(childContinuationWeight);

            const auto childSizeWeight{ childChainData.subtreeLeafCount };
            branchSizeWeights.push_back(childSizeWeight);
        }

        // Determine branching categories.
        const auto [ mpFitness, mpSameAxis, mpHigherAxis ]{
            calculateMonopodialFitness(branchContinuationWeights, branchSizeWeights) };
        const auto [ smFitness, smSameAxis, smHigherAxis ]{
            calculateSympodialMonochasialFitness(branchContinuationWeights, branchSizeWeights) };
        const auto [ sdFitness, sdSameAxis, sdHigherAxis ]{
            calculateSympodialDichasialFitness(branchContinuationWeights, branchSizeWeights) };

        // Observe branching results.
        stats.chainMonopodialBranchingFitness.observeSample(mpFitness);
        stats.aggregateBranchingFitness.observeSample(
            { branchingToOrdinal(Branching::Monopodial), mpFitness });
        stats.chainSympodialMonochasialBranchingFitness.observeSample(smFitness);
        stats.aggregateBranchingFitness.observeSample(
            { branchingToOrdinal(Branching::SympodialMonochasial), smFitness });
        stats.chainSympodialDichasialBranchingFitness.observeSample(sdFitness);
        stats.aggregateBranchingFitness.observeSample(
            { branchingToOrdinal(Branching::SympodialDichasial), sdFitness });

        // Calculate how many internodes does this chain contain before branching.
        const auto maxCompactedChainLength{ calculateMaxCompactChainLength(
            fullChains, chainsTree, compactChain) };
        const auto internodesInChain{ maxCompactedChainLength / stats.interNodeLengthEstimate };

        // Determine ramification categories.
        const auto continuousFitness{ calculateContinuousFitness(internodesInChain,
            averageInternodeStats.mean, averageInternodeStats.variance,
            averageInternodeStats.min, averageInternodeStats.max) };
        const auto rhythmicFitness{ calculateRhythmicFitness(internodesInChain,
            averageInternodeStats.mean, averageInternodeStats.variance,
            averageInternodeStats.min, averageInternodeStats.max) };
        const auto diffuseFitness{ calculateDiffuseFitness(internodesInChain,
            averageInternodeStats.mean, averageInternodeStats.variance,
            averageInternodeStats.min, averageInternodeStats.max) };

        // Observe ramification results.
        stats.chainContinuousRamificationFitness.observeSample(continuousFitness);
        stats.aggregateRamificationFitness.observeSample(
            { ramificationToOrdinal(Ramification::Continuous), continuousFitness });
        stats.chainRhythmicRamificationFitness.observeSample(rhythmicFitness);
        stats.aggregateRamificationFitness.observeSample(
            { ramificationToOrdinal(Ramification::Rhythmic), rhythmicFitness });
        stats.chainDiffuseRamificationFitness.observeSample(diffuseFitness);
        stats.aggregateRamificationFitness.observeSample(
            { ramificationToOrdinal(Ramification::Diffuse), diffuseFitness });
    }
}

/*(
 * prumer kmene
 * krivost retezu
 * oproti horizontu
 * Minimalni pocet featur ktere koreluji
 * vetvici uhel pro ruzne urovne by mel byt konstantni
 *   vypocitat rozdil oproti prumerne hodnote -> variance
 *   vynechat prvni uroven
 * Vypocitat featury pouze pro uroven 1 a 2
 * Vaha - tluste vetve pod kmenem jsou dulezitejsi.
 *      - vahovat pomoci objemu.
 * Apicalni dominance explicitni feature
 *
 * Pomery bb v shadow obrazu a ve 3D
 *
 * Overleaf:
 *  Tabulka
 *  Delky/tloustky, uhly, krivosti, objemy.
 *  Globalni, lokalni
 */

float TreeStats::calculateMaxCompactChainLength(
    const treeutil::TreeChains::ChainStorage &fullChains,
    const treeutil::TreeChains::InternalArrayTree &chainsTree,
    const treeutil::TreeChains::CompactNodeChain &compactChain)
{
    auto maxCompactedChainLength{ std::numeric_limits<float>::lowest() };

    /// @brief Helper for calculating lengths of compacted chains.
    struct ChainLengthHelper
    {
        /// Next index to be processed.
        std::size_t nextIdx{ 0u };
        /// Currently accumulated length.
        float accumulatedLength{ 0.0f };
    }; // struct ChainLengthHelper

    std::queue<ChainLengthHelper> compactedChainQueue{ };
    compactedChainQueue.push(ChainLengthHelper{ 0u, 0.0f });
    while (!compactedChainQueue.empty())
    {
        const auto chainLengthHelper{ compactedChainQueue.front() }; compactedChainQueue.pop();
        const auto subChainRec{ compactChain.compactedChains[chainLengthHelper.nextIdx] };
        const auto &subChain{ fullChains[subChainRec.chainIdx] };

        const auto chainLength{
            chainLengthHelper.accumulatedLength +
            subChain.calculateChainLength(chainsTree)
        };

        auto foundChild{ false };
        for (std::size_t iii = 0u; iii < compactChain.compactedChains.size(); ++iii)
        {
            const auto &subChainChildRec{ compactChain.compactedChains[iii] };
            if (subChainChildRec.originIdx == chainLengthHelper.nextIdx)
            { foundChild = true; compactedChainQueue.push(ChainLengthHelper{ iii, chainLength }); }
        }

        if (!foundChild)
        { maxCompactedChainLength = std::max(maxCompactedChainLength, chainLength); }
    }

    return maxCompactedChainLength;
}

std::tuple<float, TreeStats::BranchIdxList, TreeStats::BranchIdxList>
TreeStats::calculateMonopodialFitness(
        const std::vector<float> &branchContinuationWeights,
        const std::vector<float> &branchSizeWeights)
{
    /*
     * Monopodial:
     *
     *  \ | /
     *   \|/
     *    |
     *
     *  * One primary branch with high size/continuation - same axis.
     *  * Two or more child branches with the same size weights - higher axis.
     */

    BranchIdxList sameAxis{ };
    BranchIdxList higherAxis{ };

    const auto [minPrimary, argMinPrimary, maxPrimary, argMaxPrimary]{
        treeutil::argMinMax(branchSizeWeights) };
    const auto [minContinuation, argMinContinuation, maxContinuation, argMaxContinuation]{
        treeutil::argMinMax(branchContinuationWeights) };
    const auto allSame{ !treeutil::aboveEpsilon(maxPrimary - minPrimary) };
    const auto primaryBranchIdx{ allSame ? argMaxContinuation : argMaxPrimary };
    const auto maxDifference{
        allSame ?
        1.0f :
        (maxPrimary - minPrimary)
    };

    const auto secondaryBranchCount{ branchSizeWeights.size() - 1u };
    auto meanSecondarySizeWeight{ 0.0f };
    auto secondMaxPrimary{ minPrimary };
    for (std::size_t iii = 0u; iii < branchSizeWeights.size(); ++iii)
    {
        if (iii == primaryBranchIdx) { continue; }
        meanSecondarySizeWeight += branchSizeWeights[iii];
        secondMaxPrimary = std::max(secondMaxPrimary, branchSizeWeights[iii]);
    }
    meanSecondarySizeWeight /= static_cast<float>(secondaryBranchCount);

    const auto primaryBranchFitness{
        (allSame ? 1.0f : ((branchSizeWeights[primaryBranchIdx] - secondMaxPrimary) / maxDifference)) *
        ((branchContinuationWeights[primaryBranchIdx] + 1.0f) / 2.0f)
    };
    sameAxis.push_back(primaryBranchIdx);

    auto secondaryBranchFitness{ 0.0f };
    for (std::size_t iii = 0u; iii < branchSizeWeights.size(); ++iii)
    {
        if (iii == primaryBranchIdx) { continue; }
        const auto branchFitness{ branchSizeWeights[iii] - meanSecondarySizeWeight };
        secondaryBranchFitness += 1.0f - (branchFitness * branchFitness) / (maxDifference * maxDifference);
        higherAxis.push_back(iii);
    }
    secondaryBranchFitness /= static_cast<float>(std::max<std::size_t>(2u, secondaryBranchCount));

    const auto monopodialFitness{ primaryBranchFitness * secondaryBranchFitness };

    return { monopodialFitness, sameAxis, higherAxis };
}

std::tuple<float, TreeStats::BranchIdxList, TreeStats::BranchIdxList>
TreeStats::calculateSympodialMonochasialFitness(
    const std::vector<float> &branchContinuationWeights,
    const std::vector<float> &branchSizeWeights)
{
    /*
     * Sympodial Monochasial:
     *
     *  \ |/
     *   \|
     *    |
     *
     *  * One primary branch - same axis.
     *  * One secondary branch - higher axis.
     */

    BranchIdxList sameAxis{ };
    BranchIdxList higherAxis{ };

    const auto [minPrimary, argMinPrimary, maxPrimary, argMaxPrimary]{
        treeutil::argMinMax(branchSizeWeights) };
    const auto [minContinuation, argMinContinuation, maxContinuation, argMaxContinuation]{
        treeutil::argMinMax(branchContinuationWeights) };
    const auto allSame{ !treeutil::aboveEpsilon(maxPrimary - minPrimary) };
    const auto primaryBranchIdx{ allSame ? argMaxContinuation : argMaxPrimary };
    const auto maxDifference{
        allSame ?
        1.0f :
        (maxPrimary - minPrimary)
    };

    const auto secondaryBranchCount{ branchSizeWeights.size() - 1u };
    auto secondMaxPrimary{ minPrimary };
    std::size_t argSecondMaxPrimary{ 1u };
    for (std::size_t iii = 0u; iii < branchSizeWeights.size(); ++iii)
    {
        if (iii == primaryBranchIdx) { continue; }
        const auto sizeValue{ branchSizeWeights[iii] };
        if (secondMaxPrimary >= sizeValue)
        { secondMaxPrimary = sizeValue; argSecondMaxPrimary = iii; }
    }
    const auto secondaryBranchIdx{ argSecondMaxPrimary };

    const auto primaryBranchFitness{
        (allSame ? 1.0f : ((branchSizeWeights[primaryBranchIdx] - secondMaxPrimary) / maxDifference)) *
        ((branchContinuationWeights[primaryBranchIdx] + 1.0f) / 2.0f)
    };
    sameAxis.push_back(primaryBranchIdx);
    higherAxis.push_back(secondaryBranchIdx);

    auto secondaryBranchFitnessPenalization{ 0.0f };
    for (std::size_t iii = 0u; iii < branchSizeWeights.size(); ++iii)
    {
        if (iii == primaryBranchIdx || iii == secondaryBranchIdx) { continue; }
        const auto branchFitness{ branchSizeWeights[iii] };
        secondaryBranchFitnessPenalization += branchFitness / maxDifference;
        higherAxis.push_back(iii);
    }
    const auto secondaryBranchFitness{
        secondaryBranchCount >= 2u ?
        1.0f - secondaryBranchFitnessPenalization / static_cast<float>(secondaryBranchCount - 1u) :
        1.0f
    };

    const auto sympodialMonochasialFitness{ primaryBranchFitness * secondaryBranchFitness };

    return { sympodialMonochasialFitness, sameAxis, higherAxis };
}

std::tuple<float, TreeStats::BranchIdxList, TreeStats::BranchIdxList>
TreeStats::calculateSympodialDichasialFitness(
    const std::vector<float> &branchContinuationWeights,
    const std::vector<float> &branchSizeWeights)
{
    /*
     * Sympodial Dichasial:
     *
     *  \   /
     *   \|/
     *    |
     *
     *  * Two secondary branches with same size weight, low continuation - higher axis.
     *  * Optional one lower size weight branch - same axis.
     */

    BranchIdxList sameAxis{ };
    BranchIdxList higherAxis{ };

    const auto [minPrimary, argMinPrimary, maxPrimary, argMaxPrimary]{
        treeutil::argMinMax(branchSizeWeights) };
    const auto allSame{ !treeutil::aboveEpsilon(maxPrimary - minPrimary) };
    const auto primaryBranchIdx{ argMaxPrimary };
    const auto maxDifference{
        allSame ?
        1.0f :
        (maxPrimary - minPrimary)
    };

    const auto secondaryBranchCount{ branchSizeWeights.size() - 1u };
    auto secondMaxPrimary{ minPrimary };
    std::size_t argSecondMaxPrimary{ 1u };
    for (std::size_t iii = 0u; iii < branchSizeWeights.size(); ++iii)
    {
        if (iii == primaryBranchIdx) { continue; }
        const auto sizeValue{ branchSizeWeights[iii] };
        if (secondMaxPrimary >= sizeValue)
        { secondMaxPrimary = sizeValue; argSecondMaxPrimary = iii; }
    }
    const auto secondaryBranchIdx{ argSecondMaxPrimary };

    const auto primaryBranchFitness{
        (secondMaxPrimary / branchSizeWeights[primaryBranchIdx]) *
        (1.0f - (std::abs(
            branchContinuationWeights[primaryBranchIdx] -
            branchContinuationWeights[secondaryBranchIdx]
        ) / 2.0f))
    };
    higherAxis.push_back(primaryBranchIdx);
    higherAxis.push_back(secondaryBranchIdx);

    auto secondaryBranchFitnessPenalization{ 0.0f };
    for (std::size_t iii = 0u; iii < branchSizeWeights.size(); ++iii)
    {
        if (iii == primaryBranchIdx || iii == secondaryBranchIdx) { continue; }
        const auto branchFitness{ branchSizeWeights[iii] };
        secondaryBranchFitnessPenalization += (branchFitness / maxDifference) *
                                              ((branchContinuationWeights[iii] + 1.0f) / 2.0f);
        sameAxis.push_back(iii);
    }
    const auto secondaryBranchFitness{
        secondaryBranchCount >= 2u ?
        1.0f - secondaryBranchFitnessPenalization / static_cast<float>(secondaryBranchCount - 1u) :
        1.0f
    };

    const auto sympodialDichasialFitness{ primaryBranchFitness * secondaryBranchFitness };

    return { sympodialDichasialFitness, sameAxis, higherAxis };
}

ImageData copyModality(const treert::RayTracer::OutputModality<float> &modality)
{
    ImageData result{ };

    result.width = modality.width;
    result.height = modality.height;
    result.channels = 1u;
    result.valueType = ImageData::ValueType::Float;

    const auto dataSize{ modality.dataBuffer.size() * sizeof(float) };
    result.data.resize(dataSize);
    std::memcpy(result.data.data(), modality.dataBuffer.data(), dataSize);

    return result;
}

ImageData copyModality(const treert::RayTracer::OutputModality<Vector3D> &modality)
{
    ImageData result{ };

    result.width = modality.width;
    result.height = modality.height;
    result.channels = 3u;
    result.valueType = ImageData::ValueType::Float;

    const auto dataSize{ modality.dataBuffer.size() * sizeof(Vector3D) };
    result.data.resize(dataSize);
    std::memcpy(result.data.data(), modality.dataBuffer.data(), dataSize);

    return result;
}

ImageData copyModality(const treert::RayTracer::OutputModality<uint32_t> &modality)
{
    ImageData result{ };

    result.width = modality.width;
    result.height = modality.height;
    result.channels = 1u;
    result.valueType = ImageData::ValueType::UInt;

    const auto dataSize{ modality.dataBuffer.size() * sizeof(uint32_t) };
    result.data.resize(dataSize);
    std::memcpy(result.data.data(), modality.dataBuffer.data(), dataSize);

    return result;
}

template <typename ModalityT>
ImageData createEmpty(std::size_t w, std::size_t h)
{ ModalityT modality(w, h, 1); return copyModality(modality); }

void TreeStats::calculateVisualStats(TreeStatValues &stats,
    const StatsSettings &settings, const treeio::ArrayTree &inputTree)
{
    if (!settings.visualEnabled)
    {
        stats.shadowImprint = createEmpty<treert::RayTracer::VolumeModality>(1u, 1u);
        return;
    }

    // Prepare ray-tracer with tree reconstruction placed in the scene.
    const auto rayTracerPtr{
        treerndr::RenderSystemRT::prepareTreeRayTracer(
            inputTree, settings.visualTreeScale)
    };
    auto &rayTracer{ *rayTracerPtr };

    // Setup ray-tracer.
    rayTracer.setVerbose(settings.visualVerbose);
    rayTracer.setSampling(settings.visualSampleCount);

    // Prepare modalities.
    rayTracer.traceModalities({ treert::TracingModality::Volume });
    treert::RayTracer::VolumeModality tracedAccumulator{ };

    // Prepare progress printing for long jobs.
    treeutil::ProgressBar progressBar{ "Views " };
    static constexpr auto PROGRESS_PRINT_STEP{ 10u };
    treeutil::ProgressPrinter progressPrinter{ progressBar,
        std::max<std::size_t>(PROGRESS_PRINT_STEP, settings.visualViewCount),
        PROGRESS_PRINT_STEP
    };

    // Ray-trace requested views.
    const auto basePos{ settings.visualTreeScale / 2.0f };
    for (std::size_t viewIdx = 0u; viewIdx < settings.visualViewCount; ++viewIdx)
    { // Ray-trace each view independently.
        const auto t{ viewIdx / static_cast<float>(settings.visualViewCount) };

        // Position the camera.
        const auto cameraPosition{
            Vector3D{
                std::sin(t * 2.0f * treeutil::PI<float>) * basePos,
                basePos,
                std::cos(t * 2.0f * treeutil::PI<float>) * basePos
            }
        };
        const auto cameraFocus{ Vector3D{ 0.0f, basePos, 0.0f } };

        // Prepare ray-tracing context for current view.
        const auto ctx{
            rayTracer.generateOrthoContext(
                settings.visualViewWidth, settings.visualViewHeight,
                cameraPosition, cameraFocus
            )
        };

        // Perform ray tracing.
        rayTracer.traceRays(ctx);

        // Accumulate results.
        const auto &tracedData{ rayTracer.volumeModality() };
        if (tracedAccumulator.dataBuffer.empty())
        { tracedAccumulator = tracedData; }
        else
        {
            for (std::size_t jjj = 0u; jjj < tracedAccumulator.dataBuffer.size(); ++jjj)
            { tracedAccumulator.dataBuffer[jjj] += tracedData.dataBuffer[jjj]; }
        }

        // Report on progress.
        if (settings.visualViewCount >= PROGRESS_PRINT_STEP)
        { progressPrinter.printProgress(Info, viewIdx + 1u); }
    }

    // Finalize processing of modalities.
    for (auto &volume : tracedAccumulator.dataBuffer)
    { volume /= static_cast<float>(settings.visualViewCount); }

    // Save results into buffers.
    stats.shadowImprint = copyModality(tracedAccumulator);

    if (settings.visualExportResults)
    { // Export the results.
        tracedAccumulator.exportToFile(settings.visualExportPath + "/shadowImprint.png");
    }
}

void TreeStats::calculateDerivedStats(TreeStatValues &stats, const StatsSettings &settings,
    const treeutil::TreeChains &chains, const std::vector<ChainData> &chainDataStorage)
{
    const auto &chainLengths{ stats.chainLength.var.observations() };
    // Get trunk chain length, which is always the first.
    stats.trunkLength = chainLengths.empty() ? 0.0f : chainLengths.front();

    // Estimate inter-node length, without including the trunk.
    const auto [lengthMean, lengthVariance]{
        chainLengths.empty() ?
            std::make_pair(0.0f, 0.0f) :
            utils::sampleMeanVariance<float>(
                chainLengths.begin() + 1u,
                chainLengths.end()
            )
    };
    stats.interNodeLengthEstimate = lengthMean;
    stats.interNodeLengthVariance = lengthVariance;

    // Find the deepest leaf chain.
    std::size_t deepestLeafChain{ 0u };
    for (const auto &leafChainIdx : chains.leafChains())
    { deepestLeafChain = std::max(deepestLeafChain, chains.chains()[leafChainIdx].chainDepth); }

    // Estimate crown age by using the deepest leaf chain, without the trunk.
    const auto crownAge{ deepestLeafChain - 1u };
    // Estimate trunk age by dividing it into corresponding inter-node sizes.
    const auto trunkAge{ static_cast<std::size_t>(std::ceil(stats.trunkLength / stats.interNodeLengthEstimate)) };

    // Total age is the age of the crown + age of the trunk...
    stats.treeAge = crownAge + trunkAge;

    // Determine number of leaves.
    stats.leafCount = chains.leafChains().size();
}

float TreeStats::calculateContinuousFitness(float internodeCount, float mean, float var, float min, float max)
{
    /*
     * Continuous ramification is defined by regular branching, specifically
     * 1 internode per branch.
     */

    // Calculate fitness as distance from optimal value - 1 internode per branching.
    const auto diff{ std::abs(1.0f - internodeCount) };
    const auto fitness{ treeutil::smoothstep<float>(1.0f - diff) };

    return fitness;
}

float TreeStats::calculateRhythmicFitness(float internodeCount, float mean, float var, float min, float max)
{
    /*
     * Rhythmic ramification is defined by regular branching, which is other
     * than 1 internode per branch.
     */

    // Calculate regularity as distance from optimal value - average internodes per branching.
    const auto regularityDiff{ std::abs(mean - internodeCount) };
    const auto regularityFitness{ treeutil::smoothstep<float>(1.0f - regularityDiff) };

    // Calculate rhythmic nature as distance from 1 internode per branch.
    const auto rhythmicDiff{ std::abs(1.0f - internodeCount) };
    const auto rhythmicFitness{ treeutil::smoothstep<float>(rhythmicDiff) };

    // Resulting fitness is combination of the two attributes.
    const auto fitness{ regularityFitness * rhythmicFitness };
    return fitness;
}

float TreeStats::calculateDiffuseFitness(float internodeCount, float mean, float var, float min, float max)
{
    /*
     * Diffuse ramification is defined by irregular branching.
     */

    // Calculate irregularity as distance from optimal value - average internodes per branching.
    const auto diff{ std::abs(mean - internodeCount) };
    const auto fitness{ treeutil::smoothstep<float>(diff) };

    return fitness;
}

void TreeStats::finalizeStats(TreeStatValues &stats, const StatsSettings &settings)
{
    // Calculate the rest of the histograms.
    stats.forEach([] (auto &var)
    { var.var.properties(); if (!var.histPrepared) { var.calculateHistogram(MIN_BUCKETS); } });

    // Calculate properties and clean up unnecessary data:
    stats.forEach([] (auto &var)
    // TODO - Clear observations after saving?
    { /* var.var.clearObservations(true); */ });

    // Determine aggregate stats.
    stats.branching = ordinalToBranching(stats.aggregateBranchingFitness.hist.maxBucket().first);
    stats.ramification = ordinalToRamification(stats.aggregateRamificationFitness.hist.maxBucket().first);
}

TreeStats::StatsSettings &TreeStats::settings()
{ return mSettings; }
const TreeStats::StatsSettings &TreeStats::settings() const
{ return mSettings; }

} // namespace treestat