#include "tracker.h"
#include "src/helpers.h"
#include <storm-pomdp/analysis/MemlessStrategySearchQualitative.h>
#include <storm-pomdp/analysis/QualitativeAnalysisOnGraphs.h>
#include <storm-pomdp/analysis/WinningRegionQueryInterface.h>
#include <storm/logic/Formula.h>

template<typename ValueType> using SparsePomdp = storm::models::sparse::Pomdp<ValueType>;

template<typename ValueType>
std::shared_ptr<storm::pomdp::MemlessStrategySearchQualitative<ValueType>> createWinningRegionSolver(SparsePomdp<ValueType> const& pomdp, storm::logic::Formula const& formula, storm::pomdp::MemlessSearchOptions const& options) {

    STORM_LOG_TRACE("Run qualitative preprocessing...");
    storm::analysis::QualitativeAnalysisOnGraphs<ValueType> qualitativeAnalysis(pomdp);
    // After preprocessing, this might be done cheaper.
    storm::storage::BitVector targetStates = qualitativeAnalysis.analyseProb1(formula.asProbabilityOperatorFormula());
    storm::storage::BitVector surelyNotAlmostSurelyReachTarget = qualitativeAnalysis.analyseProbSmaller1(formula.asProbabilityOperatorFormula());

    storm::expressions::ExpressionManager expressionManager;
    std::shared_ptr<storm::utility::solver::SmtSolverFactory> smtSolverFactory = std::make_shared<storm::utility::solver::Z3SmtSolverFactory>();

    return std::make_shared<storm::pomdp::MemlessStrategySearchQualitative<ValueType>>(pomdp, targetStates, surelyNotAlmostSurelyReachTarget, smtSolverFactory, options);

}

template<typename ValueType>
void define_qualitative_policy_search(py::module& m, std::string const& vtSuffix) {
    m.def(("create_iterative_qualitative_search_solver_" + vtSuffix).c_str(), &createWinningRegionSolver<ValueType>, "Create solver " ,py::arg("pomdp"), py::arg("formula"), py::arg("options"));
    py::class_<storm::pomdp::MemlessStrategySearchQualitative<ValueType>, std::shared_ptr<storm::pomdp::MemlessStrategySearchQualitative<ValueType>>> mssq(m, ("IterativeQualitativeSearchSolver" + vtSuffix).c_str(), "Solver for POMDPs that solves qualitative queries");
    mssq.def("compute_winning_region", &storm::pomdp::MemlessStrategySearchQualitative<ValueType>::computeWinningRegion, py::arg("lookahead"));
    mssq.def_property_readonly("last_winning_region", &storm::pomdp::MemlessStrategySearchQualitative<ValueType>::getLastWinningRegion, "get the last computed winning region");

    py::class_<storm::pomdp::WinningRegionQueryInterface<ValueType>> wrqi(m, ("BeliefSupportWinningRegionQueryInterface" + vtSuffix).c_str());
    wrqi.def(py::init<SparsePomdp <ValueType> const&, storm::pomdp::WinningRegion const&>(), py::arg("pomdp"), py::arg("BeliefSupportWinningRegion"));
    wrqi.def("query_current_belief", &storm::pomdp::WinningRegionQueryInterface<ValueType>::isInWinningRegion, py::arg("current_belief"));
    wrqi.def("query_action",  &storm::pomdp::WinningRegionQueryInterface<ValueType>::staysInWinningRegion, py::arg("current_belief"), py::arg("action"));
}

template void define_qualitative_policy_search<double>(py::module& m, std::string const& vtSuffix);

void define_qualitative_policy_search_nt(py::module& m) {
    py::class_<storm::pomdp::MemlessSearchOptions> mssqopts(m, "IterativeQualitativeSearchOptions", "Options for the IterativeQualitativeSearch");
    mssqopts.def(py::init<>());

    py::class_<storm::pomdp::WinningRegion> winningRegion(m, "BeliefSupportWinningRegion");
}