Treatment.hpp

#pragma once
#include <creativity/data/CSVParser.hpp>
#include <creativity/data/Variable.hpp>
#include <Eigen/Core>
#include <memory>

namespace creativity { namespace data {

class Treatment {
    public:
        Treatment() = default;

        explicit Treatment(CSVParser &&csv);

        explicit Treatment(const std::string &filename);

        void readCSV(CSVParser &&csv);

        void requirePiracy(bool require = true);

        void requirePolicy(bool require = true);

        void requirePre(bool require = true);

        void requireSR(bool require = true);

        const bool& hasPre() const { return has_pre_; }

        const bool& hasPiracy() const { return has_piracy_; }

        const bool& hasPolicy() const { return has_policy_; }

        const bool& hasPiracySR() const { return has_piracy_sr_; }

        const bool& hasPolicySR() const { return has_policy_sr_; }

        bool hasShortrun() const { return (hasPolicySR() or not hasPolicy()) and (hasPiracySR() or not hasPiracy()); }

        const unsigned int& rowsPerSimulation() const { return rows_per_sim_; }

        unsigned int simulations() const { return data().rows() / rowsPerSimulation(); }

        const Eigen::MatrixXd& data() const { return data_; }

        const std::vector<std::string>& sourceFiles() const { return source_; }

        const std::string& sourceFile(unsigned r) const { return sourceFiles().at(r / rowsPerSimulation()); }

        const std::unordered_map<std::string, unsigned>& columns() const;

        std::shared_ptr<const SimpleVariable> variable(const std::string &name) const;

        std::shared_ptr<const SimpleVariable> operator[](const std::string &name) const { return variable(name); }

    protected:
        Eigen::MatrixXd data_; 
        std::vector<std::string> source_; 

        bool have_data_{false}, 
             has_pre_{false}, 
             has_piracy_{false}, 
             has_piracy_sr_{false}, 
             has_policy_{false}, 
             has_policy_sr_{false}, 
             require_pre_{false}, 
             require_piracy_{false}, 
             require_policy_{false}, 
             require_sr_{false}; 
        unsigned int rows_per_sim_{0};

        std::unordered_map<std::string, unsigned> data_column_;

        mutable std::unordered_map<std::string, std::shared_ptr<SimpleVariable>> var_cache_;

        void generateRow(const CSVParser &csv,
                Eigen::Ref<Eigen::RowVectorXd, 0, Eigen::InnerStride<>> newrow,
                const std::string &prefix) const;

};

}}