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

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

/*  CCOPYRIGHT  */

#ifndef SHAPEMODEL_H
#define SHAPEMODEL_H

#include "serialisation.h"
#include "shape.h"
#include "newimage/newimage.h"
#include "profilemodel.h"
#include "profilefilters.h"
#include "transformation.h"
#include <boost/shared_ptr.hpp>

#include <unordered_map>
#include <vector>
#include <memory>
#include <random>

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

    enum class NormalisationMode
    {
        None,
        Additive,
        Multiplicative
    };

    ShapeModel(const Shape &shape, const std::vector<std::string> &modalitynames, int profilepoints, double profilespacing,
               bool usenormalisationmasks);
    virtual ~ShapeModel();

    std::vector<std::string> GetModalityNames() const;
    int GetProfilePoints() const;
    double GetProfileSpacing() const;
    boost::shared_ptr<const Shape> GetShape() const;
    boost::shared_ptr<ProfileModel> GetVertexModel(int vertex) const;
    std::unordered_map<std::string, double> GetMeanIntensities() const;
    void SetMaxIterations(int maxiter);

    bool GetUseNormalisationMasks() const;
    void SetNormalise(std::string modality, NormalisationMode mode);
    NormalisationMode GetNormalise(std::string modality) const;
    NEWMAT::ColumnVector NormaliseProfile(std::string modality, const NEWMAT::ColumnVector &profile, double mean) const;

    void SetFilter(std::string modality, boost::shared_ptr<ProfileFilter> filter);
    boost::shared_ptr<ProfileFilter> GetFilter(std::string modality) const;

    void CreateVertexModels(boost::shared_ptr<const ProfileModel> example);

    void Train(const std::unordered_map<std::string, std::vector<std::string> > &trainingdata,
               const std::vector<boost::shared_ptr<const Transformation> > &transformations,
               const std::vector<std::string> &normalisationmasks);

    NEWMAT::ColumnVector Fit(const std::unordered_map<std::string, std::string> &data,
                             const boost::shared_ptr<const Transformation> &transformation,
                             const std::string &normalisationmask,
                             NEWMAT::ColumnVector &cilower, NEWMAT::ColumnVector &ciupper,
                             double cilevel);

    std::vector<boost::shared_ptr<ProfileModel> > TrainVertices(const std::unordered_map<std::string, std::vector<std::string> > &trainingdata,
                                                                const std::vector<boost::shared_ptr<const Transformation> > &transformations,
                                                                const std::vector<std::string> &normalisationmasks,
                                                                std::vector<int> verts,
                                                                const Plotting::plotfunc &pfunc = nullptr);

    void LoadVertexModelsAndTrain(const std::vector<std::string> &archivefiles,
                                  const std::unordered_map<std::string, std::vector<std::string> > &trainingdata,
                                  const std::vector<boost::shared_ptr<const Transformation> > &transformations,
                                  const std::vector<std::string> &normalisationmasks);

    std::unordered_map<std::string, std::vector<std::vector<NEWMAT::ColumnVector> > > SampleAllVertices(
            const std::unordered_map<std::string, std::vector<std::string> > &vols,
            const std::vector<boost::shared_ptr<const Transformation> > &transformations,
            bool computemeans,
            const std::vector<std::string> &normalisationmasks);

    virtual void WritePlots(const Plotting::plotfunc &pfunc) const;

    // TODO: Make these protected again and sort out normalisation in PlotWindow
    NEWMAT::ColumnVector SampleProfile(int vertex, int profilepoints, double profilespacing,
                                       const NEWIMAGE::volume<double> &vol, const boost::shared_ptr<const Transformation> &transformation) const;

    double GetMeanIntensity(NEWIMAGE::volume<double> vol, NEWIMAGE::volume<double> *exclusionmask,
                            double xmin, double xmax, double ymin, double ymax, double zmin, double zmax) const;
    void GetShapeExtent(const boost::shared_ptr<const Transformation> &transformation,
                        double &xmin, double &xmax, double &ymin, double &ymax, double &zmin, double &zmax) const;

protected:
    boost::shared_ptr<Shape> m_shape;

    std::vector<std::string> m_modalityNames;

    int m_profilePoints = 0;
    double m_profileSpacing = 0;

    int m_maxiter = 1;

    std::vector<boost::shared_ptr<ProfileModel> > m_vertexModels;

    void CheckModality(const std::string &modalityname) const;
    void CheckFilenames(const std::unordered_map<std::string, std::vector<std::string> > &trainingdata,
                                    const std::vector<boost::shared_ptr<const Transformation> > &transformations,
                                    const std::vector<std::string> &normalisationmasks) const;

    std::vector<std::vector<double> > GetAllDeltaLikelihoods(const std::unordered_map<std::string, std::vector<NEWMAT::ColumnVector> > &data);

    std::vector<int> UpdateDeltas(const std::vector<int> &current,
                                  const std::vector<std::vector<double> > &profilelogprobs,
                                  int &diffs) const;

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

    virtual NEWMAT::ColumnVector FitImplementation(const std::unordered_map<std::string, std::vector<NEWMAT::ColumnVector> > &data,
                                                   NEWMAT::ColumnVector &cilower, NEWMAT::ColumnVector &ciupper, double cilevel);

    virtual boost::shared_ptr<ProfileModel> TrainVertexImplementation(int vert,
            const std::unordered_map<std::string, std::vector<std::vector<NEWMAT::ColumnVector> > > &trainingdata);

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

private:
    std::unordered_map<std::string, NormalisationMode> m_normalise;
    std::unordered_map<std::string, double> m_meanIntensities;
    bool m_useNormalisationMasks;

    std::unordered_map<std::string, boost::shared_ptr<ProfileFilter> > m_filters;

    ShapeModel(const ShapeModel &);
    ShapeModel& operator=(const ShapeModel &);

    friend class boost::serialization::access;

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

template<class Archive>
void ShapeModel::serialize(Archive &ar, const unsigned int version)
{
    BOOST_LOG_TRIVIAL(debug) << "(De)serialising ShapeModel";
    // NB: Constructor parameters are serialised in non-virtual derived classes!
    ar & m_vertexModels;
    ar & m_normalise;
    ar & m_meanIntensities;
    ar & m_filters;
    ar & m_maxiter;
}


#endif // SHAPEMODEL_H
