Document.hpp

#ifndef _LDAPLUSPLUS_DOCUMENT_HPP_
#define _LDAPLUSPLUS_DOCUMENT_HPP_


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

#include <Eigen/Core>

namespace ldaplusplus {
namespace corpus {


// Forward declaration for the compiler
class Corpus;


class Document
{
    public:
        virtual const std::shared_ptr<const Corpus> get_corpus() const = 0;
        virtual const Eigen::VectorXi & get_words() const = 0;

        template <typename T>
        const std::shared_ptr<const T> get_corpus() const {
            return std::static_pointer_cast<const T>(get_corpus());
        }

        virtual ~Document(){};
};


class ClassificationDocument : public Document
{
    public:
        virtual int get_class() const = 0;
};


class Corpus
{
    public:
        virtual size_t size() const = 0;
        virtual const std::shared_ptr<Document> at(size_t index) const = 0;
        virtual void shuffle() = 0;

        virtual ~Corpus(){};
};


class ClassificationCorpus : public Corpus
{
    public:
        virtual float get_prior(int y) const = 0;
};


class EigenDocument : public Document
{
    public:
        EigenDocument(Eigen::VectorXi X) : EigenDocument(X, nullptr) {}
        EigenDocument(Eigen::VectorXi X, std::shared_ptr<const Corpus> corpus);

        const std::shared_ptr<const Corpus> get_corpus() const override;
        const Eigen::VectorXi & get_words() const override;

    private:
        Eigen::VectorXi X_;
        std::shared_ptr<const Corpus> corpus_;
};


class ClassificationDecorator : public ClassificationDocument
{
    public:
        ClassificationDecorator(std::shared_ptr<Document> doc, int y);

        const std::shared_ptr<const Corpus> get_corpus() const override;
        const Eigen::VectorXi & get_words() const override;
        int get_class() const override;

    private:
        std::shared_ptr<Document> document_;
        int y_;
};


class CorpusIndexes
{
    public:
        CorpusIndexes(int N, int random_state=0);

        int get_index(int index) const { return indices_[index]; }

        void shuffle();

    private:
        std::vector<int> indices_;

        std::mt19937 prng_;
};


class EigenCorpus : public Corpus
{
    public:
        EigenCorpus(const Eigen::MatrixXi & X, int random_state=0);

        size_t size() const override;
        virtual const std::shared_ptr<Document> at(size_t index) const override;
        void shuffle() override;

    protected:
        CorpusIndexes indices_;

        const Eigen::MatrixXi & X_;

};


class EigenClassificationCorpus : public ClassificationCorpus
{
    public:
        EigenClassificationCorpus(
            const Eigen::MatrixXi &X,
            const Eigen::VectorXi &y,
            int random_state = 0
        );

        size_t size() const override;
        virtual const std::shared_ptr<Document> at(size_t index) const override;
        void shuffle() override;
        float get_prior(int y) const override;

    private:
        CorpusIndexes indices_;

        // The data
        const Eigen::MatrixXi & X_;
        const Eigen::VectorXi & y_;

        // The class priors
        Eigen::VectorXf priors_;
};

}  // namespace corpus
}  // namespace ldaplusplus

#endif  // _LDAPLUSPLUS_DOCUMENT_HPP_