Transition.hpp

#pragma once

#include "definitions.hpp"

#include<vector>
#include<string>
#include <algorithm>
#include <stdexcept>
#include <numeric>
#include <cmath>

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

namespace craam {

using namespace std;
using namespace util::lang;

const prec_t tolerance = 1e-5;

class Transition {

public:
    Transition() : indices(0), probabilities(0), rewards(0) {};

    Transition(const indvec& indices,
                const numvec& probabilities,
                const numvec& rewards) : Transition() {

        if(indices.size() != probabilities.size() || indices.size() != rewards.size())
            throw invalid_argument("All parameters for the constructor of Transition must have the same size.");
        auto sorted = sort_indexes(indices);
        for(auto k : sorted)
            add_sample(indices[k],probabilities[k],rewards[k]);
    }

    Transition(const indvec& indices,
                const numvec& probabilities) : Transition() {

        if(indices.size() != probabilities.size())
            throw invalid_argument("All parameters for the constructor of Transition must have the same size.");
        auto sorted = sort_indexes(indices);
        for(auto k : sorted)
            add_sample(indices[k],probabilities[k],0.0);
    }

    Transition(const numvec& probabilities) : Transition() {
        for(auto k : util::lang::indices(probabilities))
            add_sample(k, probabilities[k], 0.0);
    }

    void add_sample(long stateid, prec_t probability, prec_t reward){

        if(probability < -0.001) throw invalid_argument("probabilities must be non-negative.");
        if(stateid < 0) throw invalid_argument("State id must be non-negative.");
        // if the probability is 0 or negative, just do not add the sample
        if(probability <= 0) return;

        // test for the last index; the index is not in the transition yet and belong to the end
        if(indices.size() == 0 || this->indices.back() < stateid){
            indices.push_back(stateid);
            probabilities.push_back(probability);
            rewards.push_back(reward);
        }
        // the index is already in the transitions, or belongs in the middle
        else{
            size_t findex;  // lower bound on the index of the element
            bool present;   // whether the index was found

            // test the last element for efficiency sake
            if(stateid == indices.back()){
                findex = indices.size() - 1;
                present = true;
            }
            else{
                // find the closest existing index to the new one
                auto fiter = lower_bound(indices.cbegin(),indices.cend(),stateid);
                findex = fiter - indices.cbegin();
                present = (*fiter == stateid);
            }
            // there is a transition to this element already
            if(present){
                auto p_old = probabilities[findex];
                probabilities[findex] += probability;
                auto r_old = rewards[findex];
                auto new_reward = (p_old * r_old + probability * reward) /
                                    (probabilities[findex]);
                rewards[findex] = new_reward;
            // the transition is not there, the element needs to be inserted
            }else{
                indices.insert(indices.cbegin()+findex,stateid);
                probabilities.insert(probabilities.cbegin()+findex,probability);
                rewards.insert(rewards.cbegin()+findex,reward);
            }
        }

    }

    prec_t sum_probabilities() const{
        return accumulate(probabilities.cbegin(),probabilities.cend(),0.0);
    }

    void normalize(){
        // nothing to do if there are no transitions
        if(probabilities.empty())
            return;

        prec_t sp = sum_probabilities();
        if(sp > tolerance){
            for(auto& p : probabilities)
                p /= sp;
        }else{
            throw invalid_argument("Probabilities sum to 0 (or close) and cannot be normalized.");
        }
    }

    bool is_normalized() const{
        if(indices.empty()) return true;
        else return abs(1.0 - sum_probabilities()) < tolerance;
    }

    prec_t value(numvec const& valuefunction, prec_t discount, numvec probabilities) const{
        assert(valuefunction.size() >= probabilities.size());
        assert(rewards.size() == probabilities.size());
        assert(probabilities.size() == indices.size());

        if(indices.empty())
            throw range_error("No transitions defined for the state action-pair. Cannot compute value.");
        prec_t value = 0.0;

        //Note: in simple benchmarks, the simd statement seems to speed up the computation
        // by a factor of 2-4 with -march=native on a computer with AVX support
        #pragma omp simd reduction(+:value)
        for(size_t c = 0; c < size(); c++){
            value +=  probabilities[c] * (rewards[c] + discount * valuefunction[indices[c]]);
        }
        return value;
    }

    prec_t value(numvec const& valuefunction, prec_t discount = 1.0) const{

        return value(valuefunction, discount, probabilities);
    }

    prec_t mean_reward(const numvec& probabilities) const{
        assert(probabilities.size() == size());
        if(indices.empty())
            throw range_error("No transitions defined. Cannot compute mean reward.");

        return inner_product(cbegin(probabilities), end(probabilities), cbegin(rewards), 0.0);
    }


    prec_t mean_reward() const{
        return mean_reward(probabilities);
    }

    size_t size() const {
        return indices.size();
    };

    bool empty() const {
        return indices.empty();
    };

    long max_index() const {
        return indices.empty() ? -1 : indices.back();
    };

    void probabilities_addto(prec_t scale, numvec& transition) const{
        for(size_t i : util::lang::indices(*this))
            transition[indices[i]] += scale*probabilities[i];
    }

    void probabilities_addto(prec_t scale, Transition& transition) const{
        for(size_t i : util::lang::indices(*this))
            transition.add_sample(indices[i], scale*probabilities[i], scale*rewards[i]);
    }

    numvec probabilities_vector(size_t size) const{

        if(max_index() >= 0 && static_cast<long>(size) <= max_index())
            throw range_error("Size must be greater than the maximal index");
        numvec result(size, 0.0);
        for(size_t i : util::lang::indices(indices))
            result[indices[i]] = probabilities[i];
        return result;
    }

    numvec rewards_vector(size_t size) const{

        if(max_index() >= 0 && static_cast<long>(size) <= max_index())
            throw range_error("Size must be greater than the maximal index");
        numvec result(size, 0.0);
        for(size_t i : util::lang::indices(indices))
            result[indices[i]] = rewards[i];
        return result;
    }

    const indvec& get_indices() const {return indices;};

    long get_index(long k){assert(k>=0 && k < long(size())); return indices[k];}

    const numvec& get_probabilities() const {return probabilities;};
    const numvec& get_rewards() const {return rewards;};

    void set_reward(long sampleid, prec_t reward) {rewards[sampleid] = reward;};

    prec_t get_reward(long sampleid) const {
        assert(sampleid >= 0 && sampleid < long(size())); 
        return rewards[sampleid];
    };

    string to_json(long outcomeid = -1) const{
        string result{"{"};
        result += "\"outcomeid\" : ";
        result += std::to_string(outcomeid);
        result += ",\"stateids\" : [";
        for(auto i : indices){
            result += std::to_string(i);
            result += ",";
        }
        if(!indices.empty()) result.pop_back();// remove last comma
        result += "],\"probabilities\" : [";
        for(auto p : probabilities){
            result += std::to_string(p);
            result += ",";
        }
        if(!probabilities.empty()) result.pop_back();// remove last comma
        result += "],\"rewards\" : [" ;
        for(auto r : rewards){
            result += std::to_string(r);
            result += ",";
        }
        if(!rewards.empty()) result.pop_back();// remove last comma
        result += "]}";
        return result;
    }

protected:

    indvec indices;
    numvec probabilities;
    numvec rewards;
};

}