#include "smt.hpp"

#include <brick-smtlib>
#include <algorithm>
#include <map>

#include <divine/smt/builder-smtlib.hpp>
#include <divine/smt/rpn.hpp>

template < typename key_t, typename val_t >
struct map_values_iterator
{
    using it_t = typename std::map< key_t, val_t >::const_iterator;

    it_t _it;

    map_values_iterator( it_t it ) : _it{ it } {}

    const val_t &operator*() const
    {
        return _it->second;
    }

    const val_t *operator->() const
    {
        return &_it->second;
    }

    map_values_iterator &operator++()
    {
        ++_it;
        return *this;
    }

    bool operator!=( const map_values_iterator &other )
    {
        return _it != other._it;
    }
};

namespace shoop
{

std::string formula::to_string() const
{
    auto ctx = brq::smtlib_context{};
    auto builder = divine::smt::builder::SMTLib2{ ctx, "", false };

    return brq::to_string( divine::smt::evaluate( builder, _representation ) );
}

// Transformation:
// A B & C & D & E & ~> B C & D & E &
static expr_t strip_conjunct_if_present( const expr_t &whole, const expr_t &conjunct )
{
    // Check that A is indeed a prefix.
    if ( conjunct.size() > whole.size() )
        return whole;

    for ( std::size_t i = 0; i < conjunct.size(); ++i )
        if ( whole[ i ] != conjunct[ i ] )
            return whole;

    // Strip A.
    auto res = std::vector< uint8_t >( bytes( whole ).begin() + bytes( conjunct ).size(), bytes( whole ).end() );

    // Find the first & and remove it.
    const auto and_opcode = static_cast< uint8_t >( brq::smt_op::bool_and );
    if ( auto it = std::find( res.cbegin(), res.cend(), and_opcode ); it != res.cend() )
        res.erase( it );

    if ( res.empty() )
        return formula::constant( true ).representation();

    return expr_t( res.cbegin(), res.cend() );
}

// TODO: Do we even need this? Don't we know which variable ids to expect?
static std::set< int > parse_variable_ids( const expr_t &expr )
{
    auto ids = std::set< int >{};

    for ( const auto &atom : expr )
    {
        if ( atom.is_var() )
            ids.insert( atom.varid() - 1 ); // Varids are 1-based.
    }

    return ids;
}

formula get_path_condition( const divine::mc::Trace &trace, const formula &pre )
{
    // Trace may be empty if we are running without a solver.
    if ( trace.assumes.empty() )
        return formula::constant( true );

    // Divine itself is building conjunctions, but only for formulas with the same
    // set of free variables.
    auto pc_fragments = std::map< std::set< int >, formula >{};

    for ( const auto &expr : trace.assumes )
    {
        // Don't repeat precondition, if present in the fragment
        const auto pure = strip_conjunct_if_present( expr, pre.representation() );
        const auto vars = parse_variable_ids( pure );

        pc_fragments.insert_or_assign( vars, formula{ pure, vars } );
    }

    return formula::conjunction( map_values_iterator< std::set< int >, formula >( pc_fragments.cbegin() ),
                                 map_values_iterator< std::set< int >, formula >( pc_fragments.cend() ) );
}

}