#ifndef NORMALN_H
#define NORMALN_H

#include <cmath>
#include <random>

#include <eigen3/Eigen/Dense>

#include "../../Assertions.h"
#include "../Random.h"

namespace Distribution {

    class NormalDistributionN {

    private:

        const Eigen::VectorXd mu;
        const Eigen::MatrixXd sigma;

        const double _a;
        const Eigen::MatrixXd _sigmaInv;

        const Eigen::SelfAdjointEigenSolver<Eigen::MatrixXd> eigenSolver;
        Eigen::MatrixXd transform; //can i make this const?

        RandomGenerator gen;
        std::normal_distribution<> dist;

    public:

        /** ctor */
        NormalDistributionN(const Eigen::VectorXd mu, const Eigen::MatrixXd sigma) :
            mu(mu), sigma(sigma), _a( 1.0 / std::sqrt( (sigma * 2.0 * M_PI).determinant() ) ), _sigmaInv(sigma.inverse()), eigenSolver(sigma) {

            transform = eigenSolver.eigenvectors() * eigenSolver.eigenvalues().cwiseSqrt().asDiagonal();
        }

        /** get probability for the given value */
        double getProbability(const Eigen::VectorXd val) const {
            const double b = ((val-this->mu).transpose() * this->_sigmaInv * (val-this->mu));
            return this->_a * std::exp(-b/2.0);
        }

        /** get a randomly drawn sample from the given normalN distribution*/
        Eigen::VectorXd draw() {
            return this->mu + this->transform * Eigen::VectorXd{ this->mu.size() }.unaryExpr([&](double x) { return dist(gen); });
        }

        /** get the mean vector */
        const Eigen::VectorXd getMu(){
            return this->mu;
        }

        /** get covariance matrix */
        const Eigen::MatrixXd getSigma(){
            return this->sigma;
        }

        const Eigen::MatrixXd getSigmaInv(){
            return this->_sigmaInv;
        }

        /** return a NormalN based on given data */
        static NormalDistributionN getNormalNFromSamples(const Eigen::MatrixXd& data) {

                const int numElements = data.rows();
                Assert::notEqual(numElements, 1, "data is just 1 value, thats not enough for getting the distribution!");
                Assert::notEqual(numElements, 0, "data is empty, thats not enough for getting the distribution!");

                const Eigen::VectorXd mean = data.colwise().mean();
                const Eigen::MatrixXd centered = data.rowwise() - mean.transpose();
                const Eigen::MatrixXd cov = (centered.adjoint() * centered) / double(data.rows() - 1);

                return NormalDistributionN(mean, cov);
        }

        /** return a NormalN based on given data and a given mean vector mu*/
        static NormalDistributionN getNormalNFromSamplesAndMean(const Eigen::MatrixXd& data, const Eigen::VectorXd mean) {

                const int numElements = data.rows();
                Assert::notEqual(numElements, 1, "data is just 1 value, thats not enough for getting the distribution!");
                Assert::notEqual(numElements, 0, "data is empty, thats not enough for getting the distribution!");

                const Eigen::MatrixXd centered = data.rowwise() - mean.transpose();
                const Eigen::MatrixXd cov = (centered.adjoint() * centered) / double(data.rows() - 1);

                return NormalDistributionN(mean, cov);
        }

    };

}
#endif // NORMALN_H