#pragma once

#include "storm/storage/SparseMatrix.h"
#include "storm/solver/LinearEquationSolver.h"
#include "storm/solver/MinMaxLinearEquationSolver.h"
#include "storm/solver/Multiplier.h"

namespace storm {
    class Environment;

    namespace storage {
        template <typename VT> class Scheduler;
    }

    namespace modelchecker {
        namespace helper {
            namespace internal {

                template <typename ValueType>
                class GameViHelper {
                public:
                    GameViHelper(storm::storage::SparseMatrix<ValueType> const& transitionMatrix, storm::storage::BitVector statesOfCoalition);

                    void prepareSolversAndMultipliers(const Environment& env);

                    /*!
                     * Perform value iteration until convergence
                     */
                    void performValueIteration(Environment const& env, std::vector<ValueType>& x, std::vector<ValueType> b, storm::solver::OptimizationDirection const dir);

                    /*!
                     * Perform value iteration until upper bound
                     */
                    void performValueIterationUpperBound(Environment const& env, std::vector<ValueType>& x, storm::solver::OptimizationDirection const dir, uint64_t upperBound, std::vector<ValueType>& constrainedChoiceValues);

                    /*!
                      * Fills the result vector to the original size with zeros for all states except the relevantStates
                      */
                    void fillResultVector(std::vector<ValueType>& result, storm::storage::BitVector relevantStates);

                    /*!
                     * Fills the result vector to the original size with ones for being psiStates, zeros for being not phiStates
                     */
                    void fillResultVector(std::vector<ValueType>& result, storm::storage::BitVector relevantStates, storm::storage::BitVector psiStates);

                    /*!
                     * Fills the choice values vector to the original size with zeros for ~psiState choices.
                     */
                    void fillChoiceValuesVector(std::vector<ValueType>& choiceValues, storm::storage::BitVector psiStates, std::vector<storm::storage::SparseMatrix<double>::index_type> rowGroupIndices);

                    /*!
                     * Sets whether an optimal scheduler shall be constructed during the computation
                     */
                    void setProduceScheduler(bool value);

                    /*!
                     * @return whether an optimal scheduler shall be constructed during the computation
                     */
                    bool isProduceSchedulerSet() const;

                    storm::storage::Scheduler<ValueType> extractScheduler() const;

                    void getChoiceValues(Environment const& env, std::vector<ValueType> const& x, std::vector<ValueType>& choiceValues);
                private:
                    /*!
                     * Performs one iteration step for value iteration
                     */
                    void performIterationStep(Environment const& env, storm::solver::OptimizationDirection const dir, std::vector<uint64_t>* choices = nullptr);

                    /*!
                     * Performs one iteration step for value iteration with upper bound
                     */
                    void performIterationStepUpperBound(Environment const& env, storm::solver::OptimizationDirection const dir, std::vector<uint64_t>* choices = nullptr);

                    /*!
                     * Checks whether the curently computed value achieves the desired precision
                     */
                    bool checkConvergence(ValueType precision) const;

                    std::vector<ValueType>& xNew();
                    std::vector<ValueType> const& xNew() const;

                    std::vector<ValueType>& xOld();
                    std::vector<ValueType> const& xOld() const;
                    bool _x1IsCurrent;

                    /*!
                     * @pre before calling this, a computation call should have been performed during which scheduler production was enabled.
                     * @return the produced scheduler of the most recent call.
                     */
                    std::vector<uint64_t> const& getProducedOptimalChoices() const;

                    /*!
                     * @pre before calling this, a computation call should have been performed during which scheduler production was enabled.
                     * @return the produced scheduler of the most recent call.
                     */
                    std::vector<uint64_t>& getProducedOptimalChoices();

                    storm::storage::SparseMatrix<ValueType> _transitionMatrix;
                    storm::storage::BitVector _statesOfCoalition;
                    std::vector<ValueType> _x, _x1, _x2, _b;
                    std::unique_ptr<storm::solver::Multiplier<ValueType>> _multiplier;

                    bool _produceScheduler = false;
                    boost::optional<std::vector<uint64_t>> _producedOptimalChoices;
                };
            }
        }
    }
}