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

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

/*  CCOPYRIGHT  */

#ifndef MRFSHAPEMODEL_H
#define MRFSHAPEMODEL_H

#include "shapemodel.h"

class MRFShapeModel;
BOOST_CLASS_EXPORT_KEY(MRFShapeModel)

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

    void SetWeight(double weight);
    void SetMeanFraction(double meanfrac);

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

    virtual std::vector<double> GetConditional(const std::vector<int> &deltas, int vert) const;

private:
    double m_weight;
    double m_meanFraction;
    std::vector<int> m_mean;

    double TrianglePotential(int delta1, int delta2, int delta3) const;
    double MeanPotential(int delta, int mean) const;

    std::vector<int> UpdateMean(const std::vector<int> &currentmean,
                                const std::vector<std::vector<int> > &currentdeltas,
                                int &diffs) const;

    friend class boost::serialization::access;
    template<class Archive>
    void serialize(Archive &ar, const unsigned int version);
};

template<class Archive>
void MRFShapeModel::serialize(Archive &ar, const unsigned int version)
{
    BOOST_LOG_TRIVIAL(debug) << "(De)serialising MRFShapeModel";

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

    ar & m_weight;
    ar & m_meanFraction;
    ar & m_mean;
}


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

        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, MRFShapeModel *m, const unsigned int version)
    {
        BOOST_LOG_TRIVIAL(debug) << "Loading construction info for MRFShapeModel";

        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)MRFShapeModel(*shape, modalitynames, profilepoints, profilespacing, usenormalisationmasks);
    }
}
}

#endif // MRFSHAPEMODEL_H

