State.hpp

#pragma once

#include "Action.hpp"

#include <utility>
#include <tuple>
#include <vector>
#include <stdexcept>
#include <limits>
#include <string>

#include "cpp11-range-master/range.hpp"


namespace craam {

using namespace std;



// **************************************************************************************
//  SA State (SA rectangular, also used for a regular MDP)
// **************************************************************************************

template<class AType>
class SAState{
protected:
    vector<AType> actions;
    vector<bool> valid;
    
public:

    SAState() : actions(0), valid(0) {};

    SAState(const vector<AType>& actions) : actions(actions), valid(actions.size(),true) { };

    size_t action_count() const { return actions.size();};

    size_t size() const { return action_count();};

    AType& create_action(long actionid){
        assert(actionid >= 0);
        
        // assumes that the default constructor makes the actions invalid
        if(actionid >= (long) actions.size()){
            actions.resize(actionid+1);
            valid.resize(actionid+1, false);
        }

        // set only the action that is being added as valid
        valid[actionid] = true;
        return actions[actionid];
    }

    AType& create_action() {return create_action(actions.size());};

    const AType& get_action(long actionid) const
                {assert(actionid >= 0 && size_t(actionid) < action_count());
                 return actions[actionid];};

    const AType& operator[](long actionid) const {return get_action(actionid);}

    AType& get_action(long actionid)
                {assert(actionid >= 0 && size_t(actionid) < action_count());
                 return actions[actionid];};

    AType& operator[](long actionid) {return get_action(actionid);}

    bool is_valid(long actionid) const {
        assert(actionid < long(valid.size()) && actionid >= 0);
        return valid[actionid];
    };

    void set_valid(long actionid, bool value = true){
        assert(actionid < long(valid.size()) && actionid >= 0);
        valid[actionid] = value;
    };


    const vector<AType>& get_actions() const {return actions;};

    bool is_terminal() const {return actions.empty();};

    void normalize(){
        for(AType& a : actions)
            a.normalize();
    }

    bool is_action_correct(long aid, numvec nataction) const{
        if( (aid < 0) || ((size_t)aid >= actions.size()))
            return false;

        return actions[aid].is_nature_correct(nataction);
    }

    bool is_action_correct(long aid) const{
        if( (aid < 0) || ((size_t)aid >= actions.size()))
            return false;
        else
            return true;
    }

    prec_t mean_reward(long actionid, numvec nataction) const{
        if(is_terminal()) return 0;
        else return get_action(actionid).mean_reward(nataction);
    }

    prec_t mean_reward(long actionid) const{
        if(is_terminal()) return 0;
        else return get_action(actionid).mean_reward(); 
    }

    Transition mean_transition(long action, numvec nataction) const{
        if(is_terminal()) return Transition();
        else return get_action(action).mean_transition(nataction);
    }

    Transition mean_transition(long action) const{
        if(is_terminal()) return Transition();
        else return get_action(action).mean_transition();
    }

    string to_json(long stateid = -1) const{
        string result{"{"};
        result += "\"stateid\" : ";
        result += std::to_string(stateid);
        result += ",\"actions\" : [";
        for(auto ai : indices(actions)){
            const auto& a = actions[ai];
            result += a.to_json(ai);
            result += ",";
        }
        if(!actions.empty()) result.pop_back(); // remove last comma
        result += ("]}");
        return result;
    }
};

// **********************************************************************
// *********************    SPECIFIC STATE DEFINITIONS    ***************
// **********************************************************************

typedef SAState<RegularAction> RegularState;
typedef SAState<WeightedOutcomeAction> WeightedRobustState;
}


namespace internal{
    using namespace craam;

    template<class SType>
    bool is_action_correct(const SType& state, long stateid, const std::pair<indvec,vector<numvec>>& policies){
        return state.is_action_correct(policies.first[stateid], policies.second[stateid]);
    }

    template<class SType>
    bool is_action_correct(const SType& state, long stateid, const indvec& policy){
        return state.is_action_correct(policy[stateid]);
    }
}