<commit_before>import operator as op
from operator import itemgetter as ig
from functools import singledispatch, reduce

import funcy as fn
from pysmt.shortcuts import Symbol, Equals, get_model
from pysmt.typing import REAL

import stl

from blustl import game
from blustl.utils import Result

def _bounds(phi, g):
    vars_ = [stl.Var(1, s, t) for s, t in game.vars_in_phi(phi)]
    for i, op in enumerate((">=", "<=")):
        to_lineq = lambda v: stl.LinEq((v,), op, g.model.bounds[str(v.id)][i])
        yield from map(to_lineq, vars_)


def bounds(phi, g):
    return stl.andf(*_bounds(phi, g))


def sl_to_smt(phi:"SL", g):
    store = {(s, t): Symbol(f"{s}[{t}]", REAL) for s, t in game.vars_in_phi(phi)}
    psi = bounds(phi, g) & phi
    return encode(psi, store), store


GET_OP = {
    "=": lambda x, y: (x <= y) & (x >= y),
    ">=": op.ge,
    "<=": op.le,
    ">": op.gt,
    "<": op.lt,
}


@singledispatch
def encode(psi, s):
    raise NotImplementedError(psi)

@encode.register(stl.LinEq)
def _(psi, s:dict):
    x = sum(float(term.coeff)*(s[(term.id, term.time)]+0) for term in psi.terms)
    return GET_OP[psi.op](x, psi.const)


@encode.register(stl.AtomicPred)
def _(psi, s:dict):
    # TODO
    pass


@encode.register(stl.And)
def _(psi, s:dict):
    return reduce(op.and_, (encode(x, s) for x in psi.args))

@encode.register(stl.Or)
def _(psi, s:dict):
    return reduce(op.or_, (encode(x, s) for x in psi.args))

@encode.register(stl.Neg)
def _(psi, s:dict):
    return ~encode(psi.arg, s)
<commit_msg>added Atomic Pred support to smt encoding<commit_after>import operator as op
from operator import itemgetter as ig
from functools import singledispatch, reduce

import funcy as fn
from pysmt.shortcuts import Symbol, Equals, get_model
from pysmt.typing import REAL, BooleanType

import stl

from blustl import game
from blustl.utils import Result

def _bounds(phi, g):
    vars_ = [stl.Var(1, s, t) for s, t in game.vars_in_phi(phi)]
    for i, op in enumerate((">=", "<=")):
        to_lineq = lambda v: stl.LinEq((v,), op, g.model.bounds[str(v.id)][i])
        yield from map(to_lineq, vars_)


def bounds(phi, g):
    return stl.andf(*_bounds(phi, g))


def sl_to_smt(phi:"SL", g):
    store = {(s, t): Symbol(f"{s}[{t}]", REAL) for s, t in game.vars_in_phi(phi)}
    psi = bounds(phi, g) & phi
    return encode(psi, store), store


GET_OP = {
    "=": lambda x, y: (x <= y) & (x >= y),
    ">=": op.ge,
    "<=": op.le,
    ">": op.gt,
    "<": op.lt,
}


@singledispatch
def encode(psi, s):
    raise NotImplementedError(psi)

@encode.register(stl.LinEq)
def _(psi, s:dict):
    x = sum(float(term.coeff)*(s[(term.id, term.time)]+0) for term in psi.terms)
    return GET_OP[psi.op](x, psi.const)


@encode.register(stl.AtomicPred)
def _(psi, s:dict):
    return Symbol(s[(psi.id, psi.time)], BooleanType)


@encode.register(stl.And)
def _(psi, s:dict):
    return reduce(op.and_, (encode(x, s) for x in psi.args))

@encode.register(stl.Or)
def _(psi, s:dict):
    return reduce(op.or_, (encode(x, s) for x in psi.args))

@encode.register(stl.Neg)
def _(psi, s:dict):
    return ~encode(psi.arg, s)

{'fraction_non_alphanumeric': 0.17183008623442989, 'fraction_numerical': 0.0012775471095496647, 'mean_word_length': 3.723981900452489, 'pattern_counts': {'":': 10, '<': 11, '<?xml version=': 0, '>': 11, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3458255', 'n_tokens_mistral': 1283, 'n_tokens_neox': 1252, 'n_words': 333}