#ifndef NORMALN_H
#define NORMALN_H

#include <cmath>
#include <random>
#include <vector>

#include <eigen3/Eigen/Dense>

#include "../../Assertions.h"
#include "../random/RandomGenerator.h"
#include "../../geo/Point2.h"
#include "ChiSquared.h"

namespace Distribution {

    class NormalDistributionN {

    private:

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

		double _a;
		Eigen::MatrixXd _sigmaInv;

		Eigen::MatrixXd transform;

		Random::RandomGenerator gen;
		std::normal_distribution<> dist;

    public:

		/** empty ctor */
		NormalDistributionN() {
			;
		}

		/** 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()) {

			const Eigen::SelfAdjointEigenSolver<Eigen::MatrixXd> 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() const {
            return this->mu;
        }

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

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

        void setSigma(Eigen::MatrixXd sigma){
            this->sigma = sigma;
        }

        void setMu(Eigen::VectorXd mu){
            this->mu = mu;
        }

		/** return a NormalN based on given data */
		static NormalDistributionN getNormalNFromSamples(const std::vector<std::vector<double>>& data) {

			int numAttrs = data[0].size();
			Eigen::MatrixXd eCovar = Eigen::MatrixXd::Zero(numAttrs, numAttrs);
			Eigen::VectorXd eSum = Eigen::VectorXd::Zero(numAttrs);
			Eigen::VectorXd eAvg = Eigen::VectorXd::Zero(numAttrs);

			// mean
			for (const std::vector<double>& vec : data) {
				const Eigen::Map<Eigen::VectorXd> eVec((double*)vec.data(), vec.size());
				eSum += eVec;
			}
			eAvg = eSum / data.size();

			// covariance
			for (const std::vector<double>& vec : data) {
				const Eigen::Map<Eigen::VectorXd> eVec((double*)vec.data(), vec.size());
				const Eigen::VectorXd eTmp = (eVec - eAvg);
				eCovar += eTmp * eTmp.transpose();
			}
			eCovar /= data.size();

			return NormalDistributionN(eAvg, eCovar);

		}


        /** 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();
                Eigen::MatrixXd cov = (centered.adjoint() * centered) / double(data.rows() - 1);

                return NormalDistributionN(mean, cov);
        }

		std::vector<Point2> getContour2(const double percentWithin) const {

			const int degreesOfFreedom = 2;	// 2D distribution
			const ChiSquared<double> chi(degreesOfFreedom);

			// https://people.richland.edu/james/lecture/m170/tbl-chi.html
			Assert::isNear(0.103, chi.getInvCDF(0.05), 0.01, "error within chi-squared distribution");
			Assert::isNear(0.211, chi.getInvCDF(0.10), 0.01, "error within chi-squared distribution");
			Assert::isNear(4.605, chi.getInvCDF(0.90), 0.01, "error within chi-squared distribution");
			Assert::isNear(5.991, chi.getInvCDF(0.95), 0.03, "error within chi-squared distribution");

			// size of the ellipse using confidence intervals
			const float mul = chi.getInvCDF(percentWithin);

			std::vector<Point2> res;

			//std::cout << sigma << std::endl;
			Eigen::SelfAdjointEigenSolver<Eigen::MatrixXd> solver(this->sigma);

			Eigen::Vector2d evec1 = solver.eigenvectors().col(0);
			Eigen::Vector2d evec2 = solver.eigenvectors().col(1);

			double eval1 = solver.eigenvalues()(0);
			double eval2 = solver.eigenvalues()(1);

			for (int deg = 0; deg <= 360; deg += 5) {
				const float rad = deg / 180.0f * M_PI;
				Eigen::Vector2d pos =
						std::cos(rad) * std::sqrt(mul * eval1) * evec1 +
				        std::sin(rad) * std::sqrt(mul * eval2) * evec2 +
				        this->mu;
				res.push_back(Point2(pos(0), pos(1)));
			}

			return res;

		}

    };

}
#endif // NORMALN_H