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

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

/*  CCOPYRIGHT  */

#ifndef MVNSHAPEMODEL_H
#define MVNSHAPEMODEL_H

#include "gibbsshapemodel.h"
#include <memory>

#include "armawrap/newmat.h"

class MVNShapeModel;
BOOST_CLASS_EXPORT_KEY(MVNShapeModel)

namespace boost
{
namespace serialization
{
    template<class Archive>
    inline void save_construct_data(Archive &ar, const MVNShapeModel *m, const unsigned int version);
}
}

class MVNShapeModel : public GibbsShapeModel
{
public:
    MVNShapeModel(const Shape &shape, const std::vector<std::string> &modalitynames, int profilepoints, double profilespacing,
                  bool usenormalisationmasks);

    void SetShapeN0(int n0);
    void SetShapeAlpha(double alpha);
    void SetShapeCovariancePrior(double stdev, double smoothness);
    void UseShapeVariability(bool use);

protected:
    virtual void TrainShapeImplementation(const std::unordered_map<std::string, std::vector<std::vector<NEWMAT::ColumnVector> > > &trainingdata);
    virtual std::vector<double> GetConditionalContinuous(const NEWMAT::ColumnVector &x, int vert) const;

private:
    // mu0 is zero, training sample mean isn't!
    int m_n0;
    double m_alpha;
    double m_betaStdev;
    double m_betaSmoothness;

    bool m_useShapeVariability = true;

    double m_stAlpha;
    NEWMAT::ColumnVector m_stMu;
    NEWMAT::Matrix m_stLambda;

    NEWMAT::Matrix GetBeta() const;

    friend class boost::serialization::access;

    template<class Archive>
    void save(Archive &ar, const unsigned int version) const;

    template<class Archive>
    void load(Archive &ar, const unsigned int version);

    BOOST_SERIALIZATION_SPLIT_MEMBER()
};

template<class Archive>
void MVNShapeModel::save(Archive &ar, const unsigned int version) const
{
    BOOST_LOG_TRIVIAL(debug) << "Serialising MVNShapeModel";

    ar << boost::serialization::base_object<GibbsShapeModel>(*this);

    ar << m_n0;
    ar << m_alpha;
    ar << m_betaStdev;
    ar << m_betaSmoothness;
    ar << m_stAlpha;
    ar << m_stMu;

    ar << m_stLambda;
}

template<class Archive>
void MVNShapeModel::load(Archive &ar, const unsigned int version)
{
    BOOST_LOG_TRIVIAL(info) << "Deserialising MVNShapeModel";

    ar >> boost::serialization::base_object<GibbsShapeModel>(*this);

    ar >> m_n0;
    ar >> m_alpha;
    ar >> m_betaStdev;
    ar >> m_betaSmoothness;
    ar >> m_stAlpha;
    ar >> m_stMu;

    ar >> m_stLambda;
}

namespace boost
{
namespace serialization
{
    template<class Archive>
    inline void save_construct_data(Archive &ar, const MVNShapeModel *m, const unsigned int version)
    {
        BOOST_LOG_TRIVIAL(debug) << "Saving contruction info for MVNShapeModel";

        const std::vector<std::string> modalitynames = m->GetModalityNames();
        int points = m->GetProfilePoints();
        double spacing = m->GetProfileSpacing();
        bool usenormalisationmasks = m->GetUseNormalisationMasks();
        boost::shared_ptr<const Shape> shape = m->GetShape();

        ar << modalitynames;
        ar << points;
        ar << spacing;
        ar << usenormalisationmasks;
        ar << shape;
    }

    template<class Archive>
    inline void load_construct_data(Archive &ar, MVNShapeModel *m, const unsigned int version)
    {
        BOOST_LOG_TRIVIAL(debug) << "Loading construction info for MVNShapeModel";

        std::vector<std::string> modalitynames;
        int profilepoints;
        double profilespacing;
        bool usenormalisationmasks;
        boost::shared_ptr<Shape> shape;

        ar >> modalitynames;
        ar >> profilepoints;
        ar >> profilespacing;
        ar >> usenormalisationmasks;
        ar >> shape;

        ::new(m)MVNShapeModel(*shape, modalitynames, profilepoints, profilespacing, usenormalisationmasks);
    }
}
}


#endif // MVNSHAPEMODEL_H
