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

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

/*  CCOPYRIGHT  */

#ifndef SERIALISATION_H
#define SERIALISATION_H

#include <boost/log/trivial.hpp>
#include <boost/serialization/serialization.hpp>
#include <boost/serialization/base_object.hpp>
#include <boost/serialization/export.hpp>
#include <boost/serialization/utility.hpp>
#include <boost/serialization/shared_ptr.hpp>
#include <boost/serialization/collections_save_imp.hpp>
#include <boost/serialization/collections_load_imp.hpp>
#include <boost/serialization/string.hpp>
#include <boost/serialization/vector.hpp>
#include <boost/serialization/unordered_map.hpp>
#include <boost/serialization/split_free.hpp>
#include <boost/archive/text_iarchive.hpp>
#include <boost/archive/text_oarchive.hpp>
#include "armawrap/newmat.h"
#include <unordered_map>

// NOTE: According to string.hpp, strings are treated as primitive types
//       (important when creating temporary strings in load/save etc)

// NB: Removed code for serialising std::unordered_map, as the latest version of boost includes support.
//     Serialisation format may have changed...

BOOST_CLASS_EXPORT_KEY2(NEWMAT::Matrix, "Matrix")
BOOST_CLASS_EXPORT_KEY2(NEWMAT::SymmetricMatrix, "SymmetricMatrix")
BOOST_CLASS_EXPORT_KEY2(NEWMAT::DiagonalMatrix, "DiagonalMatrix")
BOOST_CLASS_EXPORT_KEY2(NEWMAT::ColumnVector, "ColumnVector")
BOOST_CLASS_EXPORT_KEY2(NEWMAT::RowVector, "RowVector")

BOOST_CLASS_TRACKING(NEWMAT::Matrix, boost::serialization::track_never)
BOOST_CLASS_TRACKING(NEWMAT::SymmetricMatrix, boost::serialization::track_never)
BOOST_CLASS_TRACKING(NEWMAT::DiagonalMatrix, boost::serialization::track_never)
BOOST_CLASS_TRACKING(NEWMAT::ColumnVector, boost::serialization::track_never)
BOOST_CLASS_TRACKING(NEWMAT::RowVector, boost::serialization::track_never)

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

namespace boost {
namespace serialization {
    template<class Archive>
    void save(Archive &ar, const arma::Mat<NEWMAT::Real> &m, const unsigned int version)
    {
        std::size_t storage = m.n_elem;
        ar << storage;
        for (std::size_t i = 0; i < storage; i++) {
            ar << m(i);
        }
    }

    template<class Archive>
    void load(Archive &ar, arma::Mat<NEWMAT::Real> &m, const unsigned int version)
    {
        std::size_t storage;

        ar >> storage;

        if (storage != m.n_elem)
            throw SerialisationException("The matrix that is being constructed has a different storage size from the one that was serialised");

        for (std::size_t i = 0; i < storage; i++)
            ar >> m(i);
    }

    template<class Archive>
    void save(Archive &ar, const NEWMAT::Matrix &m, const unsigned int version)
    {
        std::size_t rows = m.Nrows();
        ar << rows;
        std::size_t cols = m.Ncols();
        ar << cols;
        ar << boost::serialization::base_object<arma::Mat<NEWMAT::Real>>(m);
    }

    template<class Archive>
    void load(Archive &ar, NEWMAT::Matrix &m, const unsigned int version)
    {
        std::size_t rows, cols;
        ar >> rows;
        ar >> cols;
        m.ReSize(rows, cols);
        ar >> boost::serialization::base_object<arma::Mat<NEWMAT::Real>>(m);
    }

    template<class Archive>
    void save(Archive &ar, const NEWMAT::SymmetricMatrix &m, const unsigned int version)
    {
        std::size_t rows = m.Nrows();
        ar << rows;
        ar << boost::serialization::base_object<arma::Mat<NEWMAT::Real>>(m);
    }

    template<class Archive>
    void load(Archive &ar, NEWMAT::SymmetricMatrix &m, const unsigned int version)
    {
        std::size_t rows;
        ar >> rows;

        m.ReSize(rows);
        ar >> boost::serialization::base_object<arma::Mat<NEWMAT::Real>>(m);
    }

    template<class Archive>
    void save(Archive &ar, const NEWMAT::DiagonalMatrix &m, const unsigned int version)
    {
        std::size_t rows = m.Nrows();
        ar << rows;
        ar << boost::serialization::base_object<arma::Mat<NEWMAT::Real>>(m);
    }

    template<class Archive>
    void load(Archive &ar, NEWMAT::DiagonalMatrix &m, const unsigned int version)
    {
        std::size_t rows;
        ar >> rows;
        m.ReSize(rows);
        ar >> boost::serialization::base_object<arma::Mat<NEWMAT::Real>>(m);
    }

    template<class Archive>
    void save(Archive &ar, const NEWMAT::ColumnVector &m, const unsigned int version)
    {
        std::size_t rows = m.Nrows();
        ar << rows;
        ar << boost::serialization::base_object<arma::Mat<NEWMAT::Real>>(m);
    }

    template<class Archive>
    void load(Archive &ar, NEWMAT::ColumnVector &m, const unsigned int version)
    {
        std::size_t rows;
        ar >> rows;
        m.ReSize(rows);
        ar >> boost::serialization::base_object<arma::Mat<NEWMAT::Real>>(m);
    }

    template<class Archive>
    void save(Archive &ar, const NEWMAT::RowVector &m, const unsigned int version)
    {
        std::size_t cols = m.Ncols();
        ar << cols;
        ar << boost::serialization::base_object<arma::Mat<NEWMAT::Real>>(m);
    }

    template<class Archive>
    void load(Archive &ar, NEWMAT::RowVector &m, const unsigned int version)
    {
        std::size_t cols;
        ar >> cols;
        m.ReSize(cols);
        ar >> boost::serialization::base_object<arma::Mat<NEWMAT::Real>>(m);
    }
}
}

BOOST_SERIALIZATION_SPLIT_FREE(arma::Mat<NEWMAT::Real>)
BOOST_SERIALIZATION_SPLIT_FREE(NEWMAT::Matrix)
BOOST_SERIALIZATION_SPLIT_FREE(NEWMAT::SymmetricMatrix)
BOOST_SERIALIZATION_SPLIT_FREE(NEWMAT::DiagonalMatrix)
BOOST_SERIALIZATION_SPLIT_FREE(NEWMAT::ColumnVector)
BOOST_SERIALIZATION_SPLIT_FREE(NEWMAT::RowVector)

#endif // SERIALISATION_H
