#include "src/abstraction/prism/AbstractModule.h"

#include "src/abstraction/AbstractionInformation.h"

#include "src/storage/dd/DdManager.h"
#include "src/storage/dd/Add.h"

#include "src/storage/prism/Module.h"

namespace storm {
    namespace abstraction {
        namespace prism {
            
            template <storm::dd::DdType DdType, typename ValueType>
            AbstractModule<DdType, ValueType>::AbstractModule(storm::prism::Module const& module, AbstractionInformation<DdType>& abstractionInformation, storm::utility::solver::SmtSolverFactory const& smtSolverFactory) : smtSolverFactory(smtSolverFactory), abstractionInformation(abstractionInformation), commands(), module(module) {
                
                // For each concrete command, we create an abstract counterpart.
                for (auto const& command : module.getCommands()) {
                    commands.emplace_back(command, abstractionInformation, smtSolverFactory);
                }
            }
            
            template <storm::dd::DdType DdType, typename ValueType>
            void AbstractModule<DdType, ValueType>::refine(std::vector<uint_fast64_t> const& predicates) {
                for (auto& command : commands) {
                    command.refine(predicates);
                }
            }
            
            template <storm::dd::DdType DdType, typename ValueType>
            AbstractModule<DdType, ValueType>::~AbstractModule() {
                std::cout << "destructing abs module at " << this << std::endl;
            }
            
            template <storm::dd::DdType DdType, typename ValueType>
            std::pair<storm::dd::Bdd<DdType>, uint_fast64_t> AbstractModule<DdType, ValueType>::getAbstractBdd() {
                // First, we retrieve the abstractions of all commands.
                std::vector<std::pair<storm::dd::Bdd<DdType>, uint_fast64_t>> commandDdsAndUsedOptionVariableCounts;
                uint_fast64_t maximalNumberOfUsedOptionVariables = 0;
                for (auto& command : commands) {
                    commandDdsAndUsedOptionVariableCounts.push_back(command.getAbstractBdd());
                    maximalNumberOfUsedOptionVariables = std::max(maximalNumberOfUsedOptionVariables, commandDdsAndUsedOptionVariableCounts.back().second);
                }
                
                // Then, we build the module BDD by adding the single command DDs. We need to make sure that all command
                // DDs use the same amount DD variable encoding the choices of player 2.
                storm::dd::Bdd<DdType> result = this->getAbstractionInformation().getDdManager().getBddZero();
                for (auto const& commandDd : commandDdsAndUsedOptionVariableCounts) {
                    result |= commandDd.first && this->getAbstractionInformation().getPlayer2ZeroCube(maximalNumberOfUsedOptionVariables, commandDd.second);
                }
                return std::make_pair(result, maximalNumberOfUsedOptionVariables);
            }
            
            template <storm::dd::DdType DdType, typename ValueType>
            storm::dd::Add<DdType, ValueType> AbstractModule<DdType, ValueType>::getCommandUpdateProbabilitiesAdd() const {
                storm::dd::Add<DdType, ValueType> result = this->getAbstractionInformation().getDdManager().template getAddZero<ValueType>();
                for (auto const& command : commands) {
                    result += command.getCommandUpdateProbabilitiesAdd();
                }
                return result;
            }
            
            template <storm::dd::DdType DdType, typename ValueType>
            AbstractionInformation<DdType> const& AbstractModule<DdType, ValueType>::getAbstractionInformation() const {
                return abstractionInformation.get();
            }
            
            template class AbstractModule<storm::dd::DdType::CUDD, double>;
            template class AbstractModule<storm::dd::DdType::Sylvan, double>;
        }
    }
}