from collections.abc import Iterable, Iterator, Sequence
from contextlib import contextmanager
from dataclasses import dataclass, field, replace
from functools import cached_property
from typing import Literal, Protocol, cast

import lark
import z3

from common import config
from common.relations import Relation, relation
from common.simplifier import Simplifier, simplified_or
from common.z3_utils import check_with_timeout
from common.z3_utils.expr_utils import exactly_one
from fest.heuristic import find_sym_model_by_heuristic
from fest.symstructs import SymStruct
from fest.symstructs.assignment_finder import find_assignment
from fest.symstructs.model_checker import check_sym_struct
from fest.symstructs.pickling import load_struct
from fest.templates.stdlib.arith import TemplateChoice

from . import assets

config.add_default(assets)

grammar = r"""
    ?program        : data_defns pred_defns command
    ident           : IDENTIFIER 
    !integer        : INTEGER
    variable        : ident
    ?field_name     : ident
    ?data_name      : ident
    ?pred_name      : ident
    ?type           : data_name
    ?int_expr       : sum
    ?sum            : product
                    | sum "+" product                               -> add
                    | sum "-" product                               -> sub
    ?product        : atom
                    | product "*" atom                              -> mul
    ?atom           : integer
                    | variable
                    | "-" atom                                      -> neg
                    | "(" sum ")"
    ?heap_expr      : "null"
                    | "nil"
                    | variable                    
    ?expr           : int_expr | heap_expr
    _exprs          : expr | expr "," _exprs

    pure_formula   : _pure_conj
    _pure_conj      : pure_atom
                    | pure_atom "&" _pure_conj
    ?pure_atom      : "true"                                        -> true 
                    | "false"                                       -> false
                    | comp

    HEAP_COMP_OP    : "=" | "!="                
    INT_COMP_OP     : "=" | "!=" | ">=" | "<=" | ">" | "<"          
    comp            : heap_expr HEAP_COMP_OP heap_expr
                    | int_expr INT_COMP_OP int_expr                       


    ?point_expr     : heap_expr
                    | field_name ":" heap_expr                     -> explicit_point_expr
    _point_exprs    : point_expr | point_expr "," _point_exprs

    heap_formula   : _heap_conj
    _heap_conj      : heap_atom
                    | heap_atom "*" _heap_conj                      
    ?heap_atom      : "emp"                                         -> emp              
                    | variable "->" data_name "{" _point_exprs "}"  -> points
                    | pred_name "(" _exprs ")"                      -> pred

    ?formula        : heap_formula "&" pure_formula                 -> pair_to_formula
                    | heap_formula                                  -> heap_to_formula
                    | pure_formula                                  -> pure_to_formula
                    | "(" "exists" q_vars "." formula ")"           -> exists
                    | "(" formula ")"

    !field          : type field_name ";"
    data_defn       : "data" data_name "{" field* "}" ";"?
    !data_defns     : data_defn*

    _clause         : formula
    _clauses        : formula | formula "\/" _clauses
    !clauses        : _clauses
    _vars           : variable | variable "," _vars
    !q_vars         : _vars 
    pred_defn       : "pred" pred_name "(" q_vars? ")" ":=" clauses ";"
    !pred_defns     : pred_defn*
    command         : "checkentail" formula "|-" formula ";"        -> checkentail

    COMMENT: "//" /[^\n]/*

    %import common.CNAME -> IDENTIFIER
    %import common.NUMBER -> INTEGER
    %import common.WS

    %ignore WS
    %ignore COMMENT
"""


type DataFieldType = str


@dataclass(eq=True, frozen=True)
class DataField:
    name: str
    data_type: DataFieldType


@dataclass(eq=True, frozen=True)
class DataType:
    name: str
    fields: tuple[DataField, ...]


@dataclass
class Type:
    is_int: bool = False


@dataclass(eq=True, unsafe_hash=True)
class Var:
    name: str
    type: Type = field(default_factory=Type, compare=False, hash=False)

    def is_heap(self) -> bool:
        return not self.type.is_int

    def is_int(self) -> bool:
        return self.type.is_int


type BinIntOp = Literal["+", "-", "*"]


@dataclass(eq=True, frozen=True)
class BinIntExpr:
    op: BinIntOp
    left: "Expr"
    right: "Expr"


type Expr = Var | BinIntExpr | int

type CompOp = Literal["=", "!=", ">", "<", ">=", "<="]


@dataclass(eq=True, frozen=True)
class Comp:
    op: CompOp
    left: Expr
    right: Expr


type PureAtomFormula = bool | Comp
type PureFormula = tuple[PureAtomFormula, ...]

type PointsTarget = tuple[str, Var]


def expr_contains_vars(expr: Expr, variables: set[str]) -> bool:
    match expr:
        case Var():
            return expr.name in variables
        case int():
            return False
        case BinIntExpr(_op, left, right):
            return expr_contains_vars(left, variables) or expr_contains_vars(
                right, variables
            )


def contains_vars(formula: PureAtomFormula, variables: set[str]) -> bool:
    match formula:
        case bool():
            return False
        case Comp(_op, left, right):
            return expr_contains_vars(left, variables) or expr_contains_vars(
                right, variables
            )


@dataclass(eq=True, frozen=True)
class Points:
    datatype: str
    source: Var
    targets: tuple[PointsTarget, ...]


@dataclass(eq=True, frozen=True)
class PredApp:
    pred: str
    args: tuple[Expr, ...]


type EmptyHeap = Literal["emp"]
type NonEmptyHeapAtom = Points | PredApp
type HeapAtomFormula = EmptyHeap | NonEmptyHeapAtom
type RawHeapFormula = list[HeapAtomFormula]


type FieldAccess = tuple[Var, str, str]  # source, datatype, field
type PointsConstraints = dict[Var, FieldAccess]


@dataclass(eq=True, frozen=True)
class ProcessedHeapFormula:
    points: tuple[Points, ...]
    preds: tuple[PredApp, ...]

    @classmethod
    def process(cls, formula: RawHeapFormula) -> "HeapFormula":
        points: list[Points] = []
        preds: list[PredApp] = []
        with_empty = False
        for atom in formula:
            if atom == "emp":
                with_empty = True
            elif isinstance(atom, Points):
                points.append(atom)
            else:
                assert isinstance(atom, PredApp)
                preds.append(atom)

        if not points and not preds and with_empty:
            return "emp"
        else:
            return cls(tuple(points), tuple(preds))

    def constraints(self, free_variables: set[Var]) -> PointsConstraints:
        return {
            target: (points.source, points.datatype, dt_field)
            for points in self.points
            for dt_field, target in points.targets
            if target in free_variables
        }


type HeapFormula = EmptyHeap | ProcessedHeapFormula


@dataclass(eq=True, frozen=True)
class Formula:
    q_vars: tuple[Var, ...] = field(kw_only=True, default=())
    heap_body: HeapFormula = field(kw_only=True, default="emp")
    pure_body: PureFormula = field(kw_only=True, default=(True,))


false = Formula(pure_body=(False,))


@dataclass(eq=True, frozen=True)
class Pred:
    name: str
    params: tuple[Var, ...]
    clauses: tuple[Formula, ...]

    def use_dc_optimization(self) -> bool:
        return (
            config.seplog.optimizations.distinct_case(bool)
            and len(self.clauses) == 2
            and not self.clauses[0].q_vars
        )

    def guards(self) -> Iterable[Formula]:
        if self.use_dc_optimization():
            return [false, self.clauses[0]]
        else:
            return [false for _ in self.clauses]

    def guarded_clauses(self) -> Iterable[tuple[Formula, Formula]]:
        return zip(self.guards(), self.clauses)


type Entailment = tuple[Formula, Formula]


class FormulaHeaplet(Protocol):
    def __call__(self, e: z3.ExprRef) -> z3.BoolRef: ...


class SeparatedHeaplet(Protocol):
    def contains(self, item: z3.ExprRef) -> z3.BoolRef: ...
    def separated(self) -> z3.BoolRef: ...


class EmptySeparatedHeaplet(SeparatedHeaplet):
    def contains(self, item: z3.ExprRef) -> z3.BoolRef:
        return z3.BoolVal(False)

    def separated(self) -> z3.BoolRef:
        return z3.BoolVal(True)


EMPTY_SEPARATED_HEAPLET = EmptySeparatedHeaplet()


@dataclass(eq=True, frozen=True)
class AggFormulaHeaplet(SeparatedHeaplet):
    heaplets: tuple[FormulaHeaplet, ...]

    def contains(self, item: z3.ExprRef) -> z3.BoolRef:
        return simplified_or(*(heaplet(item) for heaplet in self.heaplets))

    def separated(self) -> z3.BoolRef:
        if len(self.heaplets) == 0:
            return z3.BoolVal(True)
        elif len(self.heaplets) == 1:
            return z3.BoolVal(True)

        e = z3.Const("e", Location)
        return z3.ForAll(e, self.at_most_one()(e))

    def exactly_one(self) -> FormulaHeaplet:
        return lambda e: exactly_one(*(heaplet(e) for heaplet in self.heaplets))

    def at_least_one(self) -> FormulaHeaplet:
        return lambda e: simplified_or(*(heaplet(e) for heaplet in self.heaplets))

    def at_most_one(self) -> FormulaHeaplet:
        if len(self.heaplets) == 0:
            return lambda e: z3.BoolVal(True)
        elif len(self.heaplets) == 1:
            return lambda e: z3.BoolVal(True)
        return lambda e: z3.And(
            *(
                z3.Not(z3.And(heaplet1(e), heaplet2(e)))
                for heaplet1, heaplet2 in pairs(self.heaplets)
            )
        )


def pairs[T](seq: Sequence[T]) -> Iterable[tuple[T, T]]:
    for i, e1 in enumerate(seq):
        for j in range(i + 1, len(seq)):
            e2 = seq[j]
            yield e1, e2


Location = z3.DeclareSort("Location")


def to_const(var: Var) -> z3.Const:
    if var.is_int():
        return z3.Int(var.name)
    return z3.Const(var.name, Location)


def get_signature(*params: Var) -> Iterable[z3.SortRef]:
    return (get_sort(param) for param in params)


def get_sort(var: Var) -> z3.SortRef:
    return Location if var.is_heap() else z3.IntSort()


_nil_const = to_const(Var("nil"))


@dataclass(eq=True, frozen=True)
class Program:
    datatypes: dict[str, DataType]
    preds: dict[str, Pred]
    antecedent: Formula
    consequent: Formula
    uninterpreted: bool

    @cached_property
    def simplier(self) -> Simplifier:
        return Simplifier(True)

    @cached_property
    def funs(self) -> dict[str, dict[str, z3.FuncDeclRef]]:
        return {
            datatype.name: {
                dt_field.name: z3.Function(
                    f"{datatype.name}_{dt_field.name}", Location, Location
                )
                for dt_field in datatype.fields
            }
            for datatype in self.datatypes.values()
        }

    @cached_property
    def h_rels(self) -> dict[str, Relation]:
        return {
            pred.name: relation(
                f"h_{pred.name}", *get_signature(*pred.params), Location
            )
            for pred in self.preds.values()
        }

    @cached_property
    def rels(self) -> dict[str, Relation]:
        return {
            pred.name: relation(f"p_{pred.name}", *get_signature(*pred.params))
            for pred in self.preds.values()
        }

    @cached_property
    def heap_element(self) -> z3.Const:
        return to_const(Var("e"))

    @cached_property
    def h_empty(self) -> FormulaHeaplet:
        return lambda _e: self.simplier.false

    def h_points(
        self, points: Points, constraints: PointsConstraints
    ) -> FormulaHeaplet:
        source = self.to_const_with_constraints(points.source, constraints)
        return lambda e: e == source

    def h_pred(self, pred: Pred, *args: z3.ExprRef) -> FormulaHeaplet:
        assert len(args) == len(pred.params)
        h_rel = self.h_rels[pred.name]
        return lambda e: h_rel(*args, e)

    def translate_pure(
        self, formula: PureFormula, constraints: PointsConstraints
    ) -> z3.BoolRef:
        return self.simplier.simplified_and(
            *[self._translate_pure_atom(atom, constraints) for atom in formula]
        )

    def _translate_pure_atom(
        self, atom: PureAtomFormula, constraints: PointsConstraints
    ) -> z3.BoolRef:
        if isinstance(atom, bool):
            return z3.BoolVal(atom)
        elif isinstance(atom.left, Var) and atom.left.is_heap():
            left = self.to_const_with_constraints(atom.left, constraints)
            assert isinstance(atom.right, Var) and atom.right.is_heap()
            right = self.to_const_with_constraints(atom.right, constraints)
            match atom.op:
                case "=":
                    return left == right
                case "!=":
                    return left != right
                case _:
                    raise RuntimeError("Unreachable")
        else:
            left = self.to_int_with_constraints(atom.left, constraints)
            right = self.to_int_with_constraints(atom.right, constraints)
            match atom.op:
                case "=":
                    return left == right
                case "!=":
                    return left != right
                case _:
                    raise RuntimeError("Unreachable")

    def to_expr(self, expr: Expr, constraints: PointsConstraints) -> z3.ExprRef:
        if isinstance(expr, Var) and expr.is_heap():
            return self.to_const_with_constraints(expr, constraints)
        else:
            return self.to_int_with_constraints(expr, constraints)

    def to_const_with_constraints(
        self, name: Var, constraints: PointsConstraints
    ) -> z3.ExprRef:
        if name not in constraints:
            return to_const(name)
        source, dt, dt_field = constraints[name]
        source_location = self.to_const_with_constraints(source, constraints)
        return self.funs[dt][dt_field](source_location)

    def to_int_with_constraints(
        self, expr: Expr, constraints: PointsConstraints
    ) -> z3.ArithRef:
        if isinstance(expr, int):
            return z3.IntVal(expr)
        elif isinstance(expr, Var):
            return cast(z3.ArithRef, self.to_const_with_constraints(expr, constraints))
        else:
            left = self.to_int_with_constraints(expr.left, constraints)
            right = self.to_int_with_constraints(expr.right, constraints)
            match expr.op:
                case "+":
                    return left + right
                case "-":
                    return left - right
                case "*":
                    return left * right
                case _:
                    raise RuntimeError("Unreachable")

    def translate_heap(
        self,
        formula: HeapFormula,
        free_variables: set[Var],
        /,
        *,
        drop_pred: bool = False,
    ) -> tuple[z3.BoolRef, SeparatedHeaplet]:
        if formula == "emp":
            return z3.BoolVal(True), EMPTY_SEPARATED_HEAPLET

        formulas: list[z3.BoolRef] = []
        heaplets: list[FormulaHeaplet] = []

        if config.seplog.optimizations.points_constraints:
            constraints = formula.constraints(free_variables)
        else:
            constraints = {}

        for points in formula.points:
            source = self.to_const_with_constraints(points.source, constraints)
            formulas.append(source != _nil_const)
            dt = points.datatype
            for dt_field, target in points.targets:
                if target not in constraints:
                    formulas.append(self.funs[dt][dt_field](source) == to_const(target))
            heaplets.append(self.h_points(points, constraints))

        for pred in formula.preds:
            atom_formula, heaplet = self._translate_pred(
                pred, constraints, drop_pred=drop_pred
            )
            formulas.append(atom_formula)
            heaplets.append(heaplet)

        return self.simplier.simplified_and(*formulas), AggFormulaHeaplet(
            tuple(heaplets)
        )

    def _translate_pred(
        self,
        pred_app: PredApp,
        constraints: PointsConstraints,
        *,
        drop_pred: bool = False,
    ) -> tuple[z3.BoolRef, FormulaHeaplet]:
        pred_name = pred_app.pred
        str_args = pred_app.args
        pred = self.preds[pred_name]
        args = list(
            map(
                lambda arg: self.to_expr(arg, constraints),
                str_args,
            )
        )
        if drop_pred:
            return self.simplier.true, self.h_pred(pred, *args)
        return self.simplier.simplified_and(
            self.rels[pred_name](*args),
        ), self.h_pred(pred, *args)

    def wrap_quantifier(
        self, q_vars: list[z3.Const], *formulas: z3.BoolRef, exists: bool = True
    ) -> z3.BoolRef:
        body = self.simplier.simplified_and(*formulas)
        if not q_vars:
            return body
        if exists:
            return z3.Exists(q_vars, body)
        else:
            return z3.ForAll(q_vars, body)

    def translate_formula(
        self, formula: Formula, /, *, drop_pred: bool = False
    ) -> tuple[list[z3.Const], z3.BoolRef, z3.BoolRef, SeparatedHeaplet]:
        free_variables = set(formula.q_vars)
        heap_formula, heaplet = self.translate_heap(
            formula.heap_body, free_variables, drop_pred=drop_pred
        )

        constraints: PointsConstraints = {}
        if (
            formula.heap_body != "emp"
            and config.seplog.optimizations.points_constraints
        ):
            constraints = formula.heap_body.constraints(free_variables)

        unconstrained_vars = free_variables - constraints.keys()
        q_vars: list[z3.Const] = list(map(to_const, unconstrained_vars))
        return (
            q_vars,
            heap_formula,
            self.translate_pure(formula.pure_body, constraints),
            heaplet,
        )

    def translate_antecedent(
        self,
    ) -> tuple[z3.BoolRef, SeparatedHeaplet]:
        antecedent = self.antecedent
        _ante_q_vars, ante_heap_formula, ante_pure_formula, ante_heaplet = (
            self.translate_formula(antecedent)
        )
        formula = self.simplier.simplified_and(
            ante_heap_formula,
            ante_pure_formula,
            ante_heaplet.separated(),
        )
        return formula, ante_heaplet

    def translate_consequent_and_negate(
        self,
    ) -> tuple[list[z3.Const], z3.BoolRef, SeparatedHeaplet]:
        consequent = self.consequent
        cons_q_vars, cons_heap_formula, cons_pure_formula, cons_heaplet = (
            self.translate_formula(consequent)
        )
        formulas = [cons_heap_formula, cons_pure_formula]
        if not config.seplog.optimizations.valid_consequent_heap:
            formulas.append(cons_heaplet.separated())

        formula = self.simplier.simplified_not(self.simplier.simplified_and(*formulas))

        return cons_q_vars, formula, cons_heaplet

    def translate_guard(self, guard: Formula) -> z3.BoolRef:
        _, guard_heap, guard_pure, _ = self.translate_formula(guard)
        guard_formula = self.simplier.simplified_not(
            self.simplier.simplified_and(guard_heap, guard_pure)
        )
        return guard_formula

    def define_pred_heaplet(self, pred_name: str) -> z3.BoolRef:
        fold_formulas: list[z3.BoolRef] = []

        pred = self.preds[pred_name]
        p_vars = [to_const(var) for var in pred.params] + [self.heap_element]

        for guard, clause in pred.guarded_clauses():
            q_vars, heap_formula, pure_formula, heaplet = self.translate_formula(
                clause, drop_pred=True
            )
            clause_fo = self.wrap_quantifier(
                q_vars,
                self.translate_guard(guard),
                heap_formula,
                pure_formula,
                heaplet.separated(),
            )
            fold_formulas.append(
                self.simplier.simplified_implies(
                    clause_fo,
                    self.h_rels[pred_name](*p_vars)
                    == heaplet.contains(self.heap_element),
                )
            )

        return self.wrap_quantifier(
            p_vars,
            *fold_formulas,
            exists=False,
        )

    def define_pred_heaplets(self) -> z3.BoolRef:
        return self.simplier.simplified_and(
            *(self.define_pred_heaplet(pred_name) for pred_name in self.preds.keys())
        )

    def equal_heaps(
        self, heaplet1: SeparatedHeaplet, heaplet2: FormulaHeaplet
    ) -> z3.BoolRef:
        return z3.ForAll(
            self.heap_element,
            heaplet2(self.heap_element) == heaplet1.contains(self.heap_element),
        )

    def define_pred(self, pred_name: str) -> z3.BoolRef:
        unfold_formulas: list[z3.BoolRef] = []
        fold_formulas: list[z3.BoolRef] = []

        pred = self.preds[pred_name]
        p_vars = [to_const(var) for var in pred.params]

        for guard, clause in pred.guarded_clauses():
            q_vars, heap_formula, pure_formula, heaplet = self.translate_formula(clause)
            clause_formulas = [
                self.translate_guard(guard),
                heap_formula,
                pure_formula,
                heaplet.separated(),
            ]
            clause_fo = self.wrap_quantifier(
                q_vars,
                *clause_formulas,
            )
            unfold_formulas.append(clause_fo)
            fold_parts = [self.rels[pred_name](*p_vars)]
            if not config.seplog.optimizations.distinct_case:
                fold_parts.append(self.equal_heaps(heaplet, self.h_pred(pred, *p_vars)))

            fold_formulas.append(
                self.wrap_quantifier(
                    q_vars,
                    self.simplier.simplified_implies(
                        self.simplier.simplified_and(*clause_formulas),
                        self.simplier.simplified_and(*fold_parts),
                    ),
                    exists=False,
                )
            )

        return self.wrap_quantifier(
            p_vars,
            self.simplier.simplified_implies(
                self.rels[pred_name](*p_vars),
                self.simplier.simplified_or(*unfold_formulas),
            ),
            *fold_formulas,
            exists=False,
        )

    def define_preds(self) -> z3.BoolRef:
        return self.simplier.simplified_and(
            *(self.define_pred(pred_name) for pred_name in self.preds.keys())
        )

    def encode_without_optimizations(self) -> z3.BoolRef:
        with optimizations(False):
            return self.encode()

    def encode(self) -> z3.BoolRef:
        ante_formula, ante_heaplet = self.translate_antecedent()
        cons_q_vars, cons_formula, cons_heaplet = self.translate_consequent_and_negate()

        formulas = [
            self.define_preds(),
            ante_formula,
        ]

        if config.seplog.optimizations.same_heap:
            formulas.append(
                self.wrap_quantifier(
                    cons_q_vars,
                    z3.ForAll(
                        self.heap_element,
                        ante_heaplet.contains(self.heap_element)
                        == cons_heaplet.contains(self.heap_element),
                    ),
                    exists=False,
                )
            )
            formulas.append(
                self.wrap_quantifier(
                    cons_q_vars,
                    cons_formula,
                    exists=False,
                )
            )
        else:
            formulas.append(
                self.wrap_quantifier(
                    cons_q_vars,
                    z3.Implies(
                        z3.ForAll(
                            self.heap_element,
                            ante_heaplet.contains(self.heap_element)
                            == cons_heaplet.contains(self.heap_element),
                        ),
                        cons_formula,
                    ),
                    exists=False,
                )
            )

        if config.seplog.optimizations.distinct_case:
            formulas.append(self.define_pred_heaplets())

        return self.simplier.simplified_and(*formulas)


@dataclass
class Scope:
    current: dict[str, Var] = field(default_factory=dict)

    def add(self, var: str) -> Var:
        if var not in self.current:
            self.current[var] = Var(var)
        return self.current[var]

    def remove(self, var: Var) -> None:
        del self.current[var.name]


@dataclass
class ProgramTransformer(lark.Transformer[lark.Token, Program]):
    datatypes: dict[str, DataType] = field(default_factory=dict)
    predicates: dict[str, Pred] = field(default_factory=dict)
    current_scope: Scope = field(default_factory=Scope)
    uninterpreted: bool = True

    def ident(self, tree: list[lark.Token]) -> str:
        return str(tree[0].value)

    def integer(self, tree: list[lark.Token]) -> int:
        return int(tree[0].value)

    def variable(self, tree: tuple[str]) -> Var:
        return self.current_scope.add(tree[0])

    def field(self, tree: list[lark.Token]) -> DataField:
        data_type, name, _semicolon = tree
        return DataField(name, data_type)

    def data_defn(self, tree: tuple[str, *tuple[DataField, ...]]) -> DataType:
        name, *fields = tree
        data_type = DataType(name, tuple(fields))
        self.datatypes[name] = data_type
        return data_type

    def data_defns(self, tree: list[DataType]) -> list[DataType]:
        return tree

    def comp(self, tree: tuple[Expr, lark.Token, Expr]) -> Comp:
        left, op_token, right = tree
        op = str(op_token.value)

        if op not in {"=", "!="}:
            self.uninterpreted = False
            if isinstance(left, Var):
                left.type.is_int = True

            if isinstance(right, Var):
                right.type.is_int = True
        else:
            if not isinstance(left, Var) and isinstance(right, Var):
                right.type.is_int = True
                self.uninterpreted = False
            elif not isinstance(right, Var) and isinstance(left, Var):
                left.type.is_int = True
                self.uninterpreted = False
            elif isinstance(left, Var) and isinstance(right, Var):  # unify
                if left.type.is_int:
                    right.type = left.type
                else:
                    left.type = right.type

        return Comp(cast(CompOp, op), left, right)

    def sub(self, tree: tuple[Expr, Expr]) -> BinIntExpr:
        self.uninterpreted = False
        left, right = tree
        if isinstance(left, Var):
            left.type.is_int = True

        if isinstance(right, Var):
            right.type.is_int = True

        return BinIntExpr("-", left, right)

    def add(self, tree: tuple[Expr, Expr]) -> BinIntExpr:
        self.uninterpreted = False
        left, right = tree
        if isinstance(left, Var):
            left.type.is_int = True

        if isinstance(right, Var):
            right.type.is_int = True

        return BinIntExpr("+", left, right)

    def true(self, _tree: list[lark.Token]) -> Literal[True]:
        return True

    def false(self, _tree: list[lark.Token]) -> Literal[False]:
        return False

    def pure_formula(self, tree: list[PureAtomFormula]) -> PureFormula:
        return tuple(tree)

    def pure_to_formula(self, tree: list[PureFormula]) -> Formula:
        (pure,) = tree
        return Formula(pure_body=pure)

    def emp(self, _tree: list[lark.Token]) -> Literal["emp"]:
        return "emp"

    def pred(self, tree: tuple[str, *tuple[Expr, ...]]) -> PredApp:
        pred, *args = tree
        if pred in self.predicates:
            pred_defn = self.predicates[pred]
            for arg, param in zip(args, pred_defn.params):
                if isinstance(arg, Var):
                    arg.type = param.type
        return PredApp(pred, tuple(args))

    def explicit_point_expr(self, tree: list[str]) -> tuple[str, str]:
        field_name, target = tree
        return (field_name, target)

    def points(self, tree: tuple[Var, str, *tuple[PointsTarget | Var, ...]]) -> Points:
        source, datatype_name, *raw_targets = tree
        datatype = self.datatypes[datatype_name]

        targets = []
        for i, target in enumerate(raw_targets):
            if isinstance(target, tuple):
                targets.append(target)
            else:
                field = datatype.fields[i].name
                targets.append((field, target))

        return Points(datatype_name, source, tuple(targets))

    def heap_formula(self, tree: list[HeapAtomFormula]) -> HeapFormula:
        return ProcessedHeapFormula.process(tree)

    def heap_to_formula(self, tree: tuple[HeapFormula]) -> Formula:
        (heap,) = tree
        return Formula(heap_body=heap)

    def pair_to_formula(self, tree: tuple[HeapFormula, PureFormula]) -> Formula:
        heap, pure = tree
        return Formula(heap_body=heap, pure_body=pure)

    def q_vars(self, tree: list[str]) -> tuple[str, ...]:
        return tuple(tree)

    def exists(self, tree: tuple[tuple[Var, ...], Formula]) -> Formula:
        q_vars, formula = tree
        for var in q_vars:
            self.current_scope.remove(var)
        return replace(formula, q_vars=formula.q_vars + tuple(q_vars))

    def clauses(self, tree: list[Formula]) -> tuple[Formula, ...]:
        return tuple(tree)

    def pred_defn(self, tree: tuple[str, tuple[Var, ...], tuple[Formula, ...]]) -> Pred:
        name, p_vars, clauses = tree
        for var in p_vars:
            self.current_scope.remove(var)
        self.predicates[name] = Pred(name, tuple(p_vars), clauses)
        return self.predicates[name]

    def pred_defns(self, tree: list[Pred]) -> list[Pred]:
        return tree

    def checkentail(self, tree: Entailment) -> Entailment:
        antecedent, consequent = tree
        return antecedent, consequent

    def program(self, tree: tuple[list[DataType], list[Pred], Entailment]) -> Program:
        datatypes, preds, (antecedent, consequent) = tree
        assert (
            not antecedent.q_vars
        ), "No existential quantifiers allowed in antecedent (skolemize instead)"
        return Program(
            datatypes={datatype.name: datatype for datatype in datatypes},
            preds={pred.name: pred for pred in preds},
            antecedent=antecedent,
            consequent=consequent,
            uninterpreted=self.uninterpreted,
        )


program_parser = lark.Lark(grammar, start="program")


def raw_parse_program(code: str) -> lark.Tree[lark.Token]:
    return program_parser.parse(code)


def parse_program(code: str) -> Program:
    raw_program = raw_parse_program(code)
    return ProgramTransformer().transform(raw_program)


def find_model_for_program(
    code: str, *, template: TemplateChoice = TemplateChoice.EXTENDED
) -> SymStruct | None:
    program = parse_program(code)
    formula = program.encode()
    struct = find_sym_model_by_heuristic(formula, template=template)
    if struct is not None:
        if (
            config.seplog.optimizations.distinct_case
            and check_sym_struct(program.encode_without_optimizations(), struct)
            != z3.sat
        ):
            return None
        return struct
    return None


def find_model_from_catalog_for_program(
    code: str, struct: SymStruct | None = None
) -> SymStruct | None:
    from pathlib import Path

    program = parse_program(code)
    formula = program.encode()
    if struct is None:
        struct = load_struct(
            Path("/Users/docinf/Developer/infinite-models/code/tests/lsevenodd.fest")
        )
    assignment = find_assignment(formula, struct)
    if assignment is None:
        return None
    for const, element in assignment.items():
        struct.constants[z3.Const(const, element.sort)] = element

    return struct


@config(
    timeout=config.seplog.timeouts.validity[float],
)
def check_if_program_valid(code: str, /, *, timeout: float = 0) -> bool:
    program = parse_program(code)
    formula = program.encode_without_optimizations()
    result = check_with_timeout(formula, int(1000 * timeout))
    return result == z3.unsat


def check_if_program_uninterpreted(code: str) -> bool:
    return parse_program(code).uninterpreted


@contextmanager
def optimizations(enabled: bool = True) -> Iterator[None]:
    with config.seplog.optimizations.override(
        distinct_case=enabled, same_heap=enabled, valid_consequent_heap=enabled
    ):
        yield
