/*  Multimodal Image Segmentation Tool (MIST)  */
/*  Eelke Visser  */

/*  Copyright (c) 2016 University of Oxford  */

/*  CCOPYRIGHT  */

#ifndef STATS_H
#define STATS_H

#include "armawrap/newmat.h"
#include <vector>
#include <stdexcept>

namespace Stats
{
    double logsumexp(double a, double b);
    double logsumexp(const std::vector<double> &x);
    NEWMAT::ColumnVector mean(const std::vector<NEWMAT::ColumnVector> &x);
    NEWMAT::Matrix cov(const std::vector<NEWMAT::ColumnVector> &x);
    NEWMAT::Matrix pinv(const NEWMAT::Matrix &x, double tolfrac, int &rank);

    double lognormal(double x, double mu, double lambda);
    double logstudent(double x, double mu, double lambda, double alpha);

    class StatsException : public std::logic_error
    {
    public:
        StatsException(const std::string &what) : logic_error(what) { }
    };

    class VectorDistribution
    {
    public:
        double P(const NEWMAT::ColumnVector &x) const;
        double LogP(const NEWMAT::ColumnVector &x) const;

        virtual ~VectorDistribution() { }

    protected:
        // Length of x
        int m_k = 0;

        virtual double ComputeP(const NEWMAT::ColumnVector &x) const;
        virtual double ComputeLogP(const NEWMAT::ColumnVector &x) const = 0;
    };

    class MatrixDistribution
    {
    public:
        double P(const NEWMAT::Matrix &x) const;
        double LogP(const NEWMAT::Matrix &x) const;

        virtual ~MatrixDistribution() { }

    protected:
        // Length of x
        int m_k = 0;

        virtual double ComputeP(const NEWMAT::Matrix &x) const;
        virtual double ComputeLogP(const NEWMAT::Matrix &x) const = 0;
    };

    class MVN : public VectorDistribution
    {
    public:
        MVN(const NEWMAT::ColumnVector &mu, const NEWMAT::Matrix &lambda);
        virtual double ComputeLogP(const NEWMAT::ColumnVector &x) const;

    private:
        NEWMAT::ColumnVector m_mu;
        NEWMAT::Matrix m_lambda;
        double m_logConstant;
    };

    class MVStudent : public VectorDistribution
    {
    public:
        MVStudent(const NEWMAT::ColumnVector &mu, const NEWMAT::Matrix &lambda, double alpha);
        virtual double ComputeLogP(const NEWMAT::ColumnVector &x) const;

    private:
        NEWMAT::ColumnVector m_mu;
        NEWMAT::Matrix m_lambda;
        double m_alpha;
        double m_logConstant;
    };

    class Wishart : public MatrixDistribution
    {
    public:
        Wishart(double alpha, const NEWMAT::Matrix &beta);
        virtual double ComputeLogP(const NEWMAT::Matrix &x) const;

    private:
        double m_alpha;
        NEWMAT::Matrix m_beta;
        double m_logConstant;
    };

    class MVNWishart
    {
    public:
        MVNWishart(const NEWMAT::ColumnVector &mu, double lambda, double alpha, const NEWMAT::Matrix &beta);
        double P(const NEWMAT::ColumnVector &x, const NEWMAT::Matrix &y) const;
        double LogP(const NEWMAT::ColumnVector &x, const NEWMAT::Matrix &y) const;

    private:
        NEWMAT::ColumnVector m_mu;
        double m_lambda;
        Wishart m_wishart;
    };

    class Dirichlet : public VectorDistribution
    {
    public:
        Dirichlet(const NEWMAT::ColumnVector &alpha);
        virtual double ComputeLogP(const NEWMAT::ColumnVector &x) const;

    private:
        NEWMAT::ColumnVector m_alpha;
        double m_logConstant;
    };
}
#endif // STATS_H
