Assignment.py

import re, sys, copy, time, os
from dataclasses import dataclass, field
from enum import Enum, auto
from typing import List, Dict

# How deep the prover will search before giving up
MAX_DEPTH = 8



# lexer  –  breaks raw text into tokens 


# Every possible token type
class TT(Enum):
    LPAREN = auto(); RPAREN = auto(); COMMA = auto(); DOT = auto()
    AND    = auto(); OR     = auto(); NOT   = auto()
    IMPLIES= auto(); IFF    = auto()
    FORALL = auto(); EXISTS = auto(); EQUALS = auto()
    IDENT  = auto(); EOF    = auto()

@dataclass
class Token:
    type:  TT
    value: str
    pos:   int   # character position in the original string 

TOKEN_RULES = [
    (r"<->",              TT.IFF),
    (r"->",               TT.IMPLIES),
    (r"forall\b",         TT.FORALL),
    (r"exists\b",         TT.EXISTS),
    (r"&",                TT.AND),
    (r"\|",               TT.OR),
    (r"~",                TT.NOT),
    (r"\(",               TT.LPAREN),
    (r"\)",               TT.RPAREN),
    (r",",                TT.COMMA),
    (r"\.",               TT.DOT),
    (r"=",                TT.EQUALS),
    (r"[A-Za-z_][A-Za-z0-9_]*", TT.IDENT),
]

# Combine all patterns 
MASTER_PATTERN = re.compile(
    "|".join(f"(?P<T{i}>{pat})" for i, (pat, _) in enumerate(TOKEN_RULES))
)
# Map group name 
GROUP_TO_TYPE = {f"T{i}": tt for i, (_, tt) in enumerate(TOKEN_RULES)}

class LexError(Exception):  pass
class ParseError(Exception): pass

def tokenize(text):
    #Turn a formula string into a list of Token objects
    tokens = []
    pos = 0
    while pos < len(text):
        if text[pos].isspace():
            pos += 1
            continue
        match = MASTER_PATTERN.match(text, pos)
        if not match:
            raise LexError(f"Unexpected character {text[pos]!r} at position {pos}")
        tokens.append(Token(GROUP_TO_TYPE[match.lastgroup], match.group(), pos))
        pos = match.end()
    tokens.append(Token(TT.EOF, "", pos))  
    return tokens


# AST  –  the tree that represents a parsed formula

# Terms (things that can be arguments to predicates) 

class Term:
    def freeVars(self):           raise NotImplementedError
    def substitute(self, v, t):   raise NotImplementedError

@dataclass
class Variable(Term):
    #A bound variable introduced by a quantifier e.g x in forall x
    name: str
    def freeVars(self):          return {self.name}
    def substitute(self, v, t):  return t if self.name == v else self
    def __str__(self):           return self.name
    def __hash__(self):          return hash(("var", self.name))
    def __eq__(self, o):         return isinstance(o, Variable) and self.name == o.name

@dataclass
class Constant(Term):
    #fixed individual, e.g. a, b, or a fresh 
    name: str
    def freeVars(self):          return set()
    def substitute(self, v, t):  return self          # constants are never replaced
    def __str__(self):           return self.name
    def __hash__(self):          return hash(("con", self.name))
    def __eq__(self, o):         return isinstance(o, Constant) and self.name == o.name

@dataclass
class Function(Term):
    # function applied to arguments, e.g. f(x, a)
    name: str
    args: List[Term] = field(default_factory=list)
    def freeVars(self):          return set().union(*(a.freeVars() for a in self.args))
    def substitute(self, v, t):  return Function(self.name, [a.substitute(v, t) for a in self.args])
    def __str__(self):           return f"{self.name}({','.join(str(a) for a in self.args)})"
    def __hash__(self):          return hash(("fun", self.name, tuple(self.args)))
    def __eq__(self, o):         return isinstance(o, Function) and self.name == o.name and self.args == o.args


#Formulas

class Formula:
    def freeVars(self):          raise NotImplementedError
    def substitute(self, v, t):  raise NotImplementedError

@dataclass
class Predicate(Formula):
    #An atomic predicate
    name: str
    args: List[Term] = field(default_factory=list)
    def freeVars(self):          return set().union(*(a.freeVars() for a in self.args))
    def substitute(self, v, t):  return Predicate(self.name, [a.substitute(v, t) for a in self.args])
    def __str__(self):
        s = ",".join(str(a) for a in self.args)
        return f"{self.name}({s})" if s else self.name
    def __hash__(self):          return hash(("pred", self.name, tuple(self.args)))
    def __eq__(self, o):         return isinstance(o, Predicate) and self.name == o.name and self.args == o.args

@dataclass
class Equality(Formula):
    #Term equality: left = right
    left:  Term
    right: Term
    def freeVars(self):          return self.left.freeVars() | self.right.freeVars()
    def substitute(self, v, t):  return Equality(self.left.substitute(v, t), self.right.substitute(v, t))
    def __str__(self):           return f"{self.left}={self.right}"
    def __hash__(self):          return hash(("eq", self.left, self.right))
    def __eq__(self, o):         return isinstance(o, Equality) and self.left == o.left and self.right == o.right

@dataclass
class Negation(Formula):
    formula: Formula
    def freeVars(self):          return self.formula.freeVars()
    def substitute(self, v, t):  return Negation(self.formula.substitute(v, t))
    def __str__(self):           return f"~{self.formula}"


# Helper that builds a binary connective class (
def makeBinary(className, symbol):
    @dataclass
    class BinOp(Formula):
        left:  Formula
        right: Formula
        def freeVars(self):         return self.left.freeVars() | self.right.freeVars()
        def substitute(self, v, t): return BinOp(self.left.substitute(v, t), self.right.substitute(v, t))
        def __str__(self):          return f"({self.left}{symbol}{self.right})"
        def __hash__(self):         return hash((className, self.left, self.right))
        def __eq__(self, o):        return type(o) is type(self) and self.left == o.left and self.right == o.right
    BinOp.__name__ = BinOp.__qualname__ = className
    return BinOp

Conjunction = makeBinary("Conjunction", "&")   # A & B
Disjunction = makeBinary("Disjunction", "|")   # A | B
Implication = makeBinary("Implication", "->")  # A -> B


# Helper that builds a quantifier class
def makeQuantifier(className, symbol):
    @dataclass
    class Quant(Formula):
        variable: str     # the bound variable name
        formula:  Formula # the body
        def freeVars(self):
            return self.formula.freeVars() - {self.variable}
        def substitute(self, v, t):
            # Don't substitute inside a quantifier that re-binds the same variable
            return self if v == self.variable else Quant(self.variable, self.formula.substitute(v, t))
        def __str__(self):  return f"({symbol}{self.variable}.{self.formula})"
        def __hash__(self): return hash((className, self.variable, self.formula))
        def __eq__(self, o): return type(o) is type(self) and self.variable == o.variable and self.formula == o.formula
    Quant.__name__ = Quant.__qualname__ = className
    return Quant

Universal   = makeQuantifier("Universal",   "∀")  # forall x
Existential = makeQuantifier("Existential", "∃")  # exists x



#Sequent

@dataclass
class Sequent:
    ant: List[Formula] = field(default_factory=list)  # assumptions (left of ⊢)
    suc: List[Formula] = field(default_factory=list)  # goals      (right of ⊢)
    def __str__(self):
        left  = ", ".join(str(f) for f in self.ant) or "∅"
        right = ", ".join(str(f) for f in self.suc) or "∅"
        return f"{left} ⊢ {right}"


# Parser


class Parser:
    def __init__(self, tokens):
        self.tokens   = tokens
        self.pos      = 0
        self.boundVars = []   # stack of sets; each set holds vars bound by the current quantifier

    # Helpers 

    @property
    def current(self): return self.tokens[self.pos]

    def advance(self):
        t = self.current
        if t.type != TT.EOF: self.pos += 1
        return t

    def expect(self, tokenType):
        if self.current.type != tokenType:
            raise ParseError(f"Expected {tokenType.name} but got {self.current.value!r}")
        return self.advance()

    def check(self, *types): return self.current.type in types

    def isBound(self, name):
        #return True if 'name' was introduced by a surrounding quantifier
        return any(name in scope for scope in self.boundVars)

    #rules (from lowest to highest precedence) 

    def parse(self):
        f = self.parseBiconditional()
        if self.current.type != TT.EOF:
            raise ParseError(f"Unexpected token {self.current.value!r}")
        return f

    def parseBiconditional(self):
        left = self.parseImplication()
        while self.check(TT.IFF):
            self.advance()
            right = self.parseImplication()
            # A <-> B  is shorthand for  (A->B) & (B->A)
            left = Conjunction(Implication(left, right), Implication(right, left))
        return left

    def parseImplication(self):
        left = self.parseDisjunction()
        while self.check(TT.IMPLIES):
            self.advance()
            right = self.parseDisjunction()
            left  = Implication(left, right)
        return left

    def parseDisjunction(self):
        left = self.parseConjunction()
        while self.check(TT.OR):
            self.advance()
            right = self.parseConjunction()
            left  = Disjunction(left, right)
        return left

    def parseConjunction(self):
        left = self.parseUnary()
        while self.check(TT.AND):
            self.advance()
            right = self.parseUnary()
            left  = Conjunction(left, right)
        return left

    def parseUnary(self):
        if self.check(TT.NOT):
            self.advance()
            return Negation(self.parseUnary())
        if self.check(TT.FORALL, TT.EXISTS):
            return self.parseQuantifier()
        return self.parsePrimary()

    def parseQuantifier(self):
        isForAll = self.current.type == TT.FORALL
        self.advance()
        # Collect the variable names before the dot
        varNames = []
        while self.check(TT.IDENT):
            varNames.append(self.current.value)
            self.advance()
            if self.check(TT.DOT): break
        if not varNames:
            raise ParseError("Expected a variable name after quantifier")
        self.expect(TT.DOT)
        # Push a new scope so the body knows these names are bound
        self.boundVars.append(set(varNames))
        body = self.parseBiconditional()
        self.boundVars.pop()
        # Wrap body in nested quantifiers 
        Q = Universal if isForAll else Existential
        formula = body
        for v in reversed(varNames):
            formula = Q(v, formula)
        return formula

    def parsePrimary(self):
        if self.check(TT.LPAREN):
            self.advance()
            f = self.parseBiconditional()
            self.expect(TT.RPAREN)
            return f
        return self.parseAtom()

    def parseAtom(self):
        #parse a predicate or equality atom
        name = self.expect(TT.IDENT).value
        if self.check(TT.LPAREN):
            # Predicate with arguments
            self.advance()
            args = self.parseTermList()
            self.expect(TT.RPAREN)
            if self.check(TT.EQUALS):
                self.advance()
                return Equality(Function(name, args), self.parseTerm())
            return Predicate(name, args)
        if self.check(TT.EQUALS):
            # Bare equality: a = b
            lhs = self.makeTerm(name)
            self.advance()
            return Equality(lhs, self.parseTerm())
        return Predicate(name, [])   

    def parseTerm(self):
        name = self.expect(TT.IDENT).value
        if self.check(TT.LPAREN):
            self.advance()
            args = self.parseTermList()
            self.expect(TT.RPAREN)
            return Function(name, args)
        return self.makeTerm(name)

    def parseTermList(self):
        terms = [self.parseTerm()]
        while self.check(TT.COMMA):
            self.advance()
            terms.append(self.parseTerm())
        return terms

    def makeTerm(self, name):
        #decide whether a name is a Variable or Constant 
        return Variable(name) if self.isBound(name) else Constant(name)


def parseFormula(text):
    return Parser(tokenize(text)).parse()

def parseFile(path):
    #Read a file line by line and return 
    results = []
    with open(path) as fh:
        for line in fh:
            line = line.strip()
            if not line or line.startswith("#"):
                continue   # skip blank lines and comments
            try:
                formula = parseFormula(line)
                results.append((line, Sequent(ant=[], suc=[formula]), None))
            except (LexError, ParseError) as e:
                results.append((line, None, str(e)))
    return results


# proof treees

@dataclass
class ProofNode:
    sequent:  Sequent
    rule:     str  = ""
    children: list = field(default_factory=list)
    closed:   bool = False

    def isProved(self):
        # node is proved if it is closed & every child is proved
        return self.closed and all(c.isProved() for c in self.children)

    def printTree(self, indent=0):
        pad  = "  " * indent
        mark = "✓" if self.isProved() else "✗"
        print(f"{pad}{mark} {self.sequent}  [{self.rule}]")
        for child in self.children:
            child.printTree(indent + 1)


# shared utilities 

freshCounter = 0

def freshConstant():
    #Generate a brand new constant 
    global freshCounter
    c = Constant(f"_c{freshCounter}")
    freshCounter += 1
    return c

def collectGroundTerms(seq):
    #return all constants and function terms that appear anywhere in the sequent
    seen = set()
    result = []

    def visitTerm(t):
        key = str(t)
        if key in seen: return
        seen.add(key)
        if isinstance(t, (Constant, Function)):
            result.append(t)
        if isinstance(t, Function):
            for arg in t.args: visitTerm(arg)

    def visitFormula(f):
        if isinstance(f, Predicate):                           [visitTerm(a) for a in f.args]
        elif isinstance(f, Equality):                          visitTerm(f.left); visitTerm(f.right)
        elif isinstance(f, Negation):                          visitFormula(f.formula)
        elif isinstance(f, (Conjunction, Disjunction, Implication)): visitFormula(f.left); visitFormula(f.right)
        elif isinstance(f, (Universal, Existential)):          visitFormula(f.formula)

    for f in seq.ant + seq.suc: visitFormula(f)
    return result

def isAtom(f):
    return isinstance(f, (Predicate, Equality))

def isAxiom(seq):

    #A sequent is an axiom if false, bot, F appears in the antecedent or t = t appears in the succedent or same atomic formula appears on both sides
    for f in seq.ant:
        if isinstance(f, Predicate) and f.name in ("False", "bot", "F") and not f.args:
            return True
    for f in seq.suc:
        if isinstance(f, Equality) and f.left == f.right:
            return True
    antAtoms = {str(f) for f in seq.ant if isAtom(f)}
    return any(str(f) in antAtoms for f in seq.suc if isAtom(f))

def dropFirst(lst, item):
    #Return a copy of lst with the first occurrence of item removed
    result = list(lst)
    result.remove(item)
    return result


# baseline prover

def proveBaseline(seq, depth=0, used=None):
    if used is None: used = {}
    node = ProofNode(sequent=seq)

    # Group one - close the branch immediately if possible
    if isAxiom(seq):
        node.rule, node.closed = "id/⊥L", True
        return node

    # Group 2 - Single-premise rules (no branching) 
    # ¬R
    for f in seq.suc:
        if isinstance(f, Negation):
            ch = proveBaseline(Sequent(seq.ant + [f.formula], dropFirst(seq.suc, f)), depth, copy.deepcopy(used))
            node.rule, node.children, node.closed = "¬R", [ch], ch.isProved()
            return node
    # ¬L
    for f in seq.ant:
        if isinstance(f, Negation):
            ch = proveBaseline(Sequent(dropFirst(seq.ant, f), seq.suc + [f.formula]), depth, copy.deepcopy(used))
            node.rule, node.children, node.closed = "¬L", [ch], ch.isProved()
            return node
    # →R
    for f in seq.suc:
        if isinstance(f, Implication):
            newAnt = seq.ant + [f.left]
            newSuc = dropFirst(seq.suc, f) + [f.right]
            ch = proveBaseline(Sequent(newAnt, newSuc), depth, copy.deepcopy(used))
            node.rule, node.children, node.closed = "→R", [ch], ch.isProved()
            return node
    # ∧L
    for f in seq.ant:
        if isinstance(f, Conjunction):
            ch = proveBaseline(Sequent(dropFirst(seq.ant, f) + [f.left, f.right], list(seq.suc)), depth, copy.deepcopy(used))
            node.rule, node.children, node.closed = "∧L", [ch], ch.isProved()
            return node
    # ∨R
    for f in seq.suc:
        if isinstance(f, Disjunction):
            ch = proveBaseline(Sequent(list(seq.ant), dropFirst(seq.suc, f) + [f.left, f.right]), depth, copy.deepcopy(used))
            node.rule, node.children, node.closed = "∨R", [ch], ch.isProved()
            return node
    # ∀R
    for f in seq.suc:
        if isinstance(f, Universal):
            c  = freshConstant()
            ch = proveBaseline(Sequent(list(seq.ant), dropFirst(seq.suc, f) + [f.formula.substitute(f.variable, c)]), depth, copy.deepcopy(used))
            node.rule, node.children, node.closed = f"∀R[{c}]", [ch], ch.isProved()
            return node
    # ∃L
    for f in seq.ant:
        if isinstance(f, Existential):
            c  = freshConstant()
            ch = proveBaseline(Sequent(dropFirst(seq.ant, f) + [f.formula.substitute(f.variable, c)], list(seq.suc)), depth, copy.deepcopy(used))
            node.rule, node.children, node.closed = f"∃L[{c}]", [ch], ch.isProved()
            return node

    #  Group 3: Branching rules (create two sub-goals) 
    # ∧R
    for f in seq.suc:
        if isinstance(f, Conjunction):
            rest = dropFirst(seq.suc, f)
            lc = proveBaseline(Sequent(list(seq.ant), rest + [f.left]),  depth, copy.deepcopy(used))
            rc = proveBaseline(Sequent(list(seq.ant), rest + [f.right]), depth, copy.deepcopy(used))
            node.rule, node.children, node.closed = "∧R", [lc, rc], lc.isProved() and rc.isProved()
            return node
    # ∨L
    for f in seq.ant:
        if isinstance(f, Disjunction):
            rest = dropFirst(seq.ant, f)
            lc = proveBaseline(Sequent(rest + [f.left],  list(seq.suc)), depth, copy.deepcopy(used))
            rc = proveBaseline(Sequent(rest + [f.right], list(seq.suc)), depth, copy.deepcopy(used))
            node.rule, node.children, node.closed = "∨L", [lc, rc], lc.isProved() and rc.isProved()
            return node
    # →L
    for f in seq.ant:
        if isinstance(f, Implication):
            rest = dropFirst(seq.ant, f)
            lc = proveBaseline(Sequent(list(rest),          list(seq.suc) + [f.left]), depth, copy.deepcopy(used))
            rc = proveBaseline(Sequent(rest + [f.right],    list(seq.suc)),            depth, copy.deepcopy(used))
            node.rule, node.children, node.closed = "→L", [lc, rc], lc.isProved() and rc.isProved()
            return node

    # Group 4: Quantifier instantiation (bounded repetition)
    if depth >= MAX_DEPTH:
        node.rule, node.closed = "DEPTH_LIMIT", False
        return node

    terms = collectGroundTerms(seq) or [freshConstant()]

    # ∃R
    for f in seq.suc:
        if isinstance(f, Existential):
            key  = str(f)
            done = used.get(key, set())
            t    = next((x for x in terms if str(x) not in done), freshConstant())
            nu   = copy.deepcopy(used); nu.setdefault(key, set()).add(str(t))
            ch   = proveBaseline(Sequent(list(seq.ant), dropFirst(seq.suc, f) + [f.formula.substitute(f.variable, t)]), depth + 1, nu)
            node.rule, node.children, node.closed = f"∃R[{t}]", [ch], ch.isProved()
            return node
    # ∀L
    for f in seq.ant:
        if isinstance(f, Universal):
            key  = str(f)
            done = used.get(key, set())
            t    = next((x for x in terms if str(x) not in done), freshConstant())
            nu   = copy.deepcopy(used); nu.setdefault(key, set()).add(str(t))
            ch   = proveBaseline(Sequent(list(seq.ant) + [f.formula.substitute(f.variable, t)], list(seq.suc)), depth + 1, nu)
            node.rule, node.children, node.closed = f"∀L[{t}]", [ch], ch.isProved()
            return node

    node.rule, node.closed = "STUCK", False
    return node


# Improved prover
# Improvement 1 – Loop detectionm
# Improvement 2 – Complexity ordering
# Improvement 3 – fast axiom check
# improvement 4 - Branch ordering
complexityCache: Dict[int, int] = {}

def complexity(f) -> int:
    #Score how complex a formula is ( simpler formulas get lower scores)
    fid = id(f)
    if fid in complexityCache:
        return complexityCache[fid]
    if isinstance(f, (Predicate, Equality)):
        score = 0
    elif isinstance(f, Negation):
        score = 1 + complexity(f.formula)
    elif isinstance(f, (Conjunction, Disjunction, Implication)):
        score = 1 + complexity(f.left) + complexity(f.right)
    elif isinstance(f, (Universal, Existential)):
        score = 2 + complexity(f.formula)
    else:
        score = 0
    complexityCache[fid] = score
    return score

def seqSignature(ant, suc):
    # signature for sequent used by loop detection"""
    return (frozenset(str(f) for f in ant), frozenset(str(f) for f in suc))

def fastAxiom(ant, suc):
    #axiom check 
    for f in ant:
        if isinstance(f, Predicate) and f.name in ("False", "bot", "F") and not f.args:
            return True
    #t = t is always true
    for f in suc:
        if isinstance(f, Equality) and f.left == f.right:
            return True
    antAtoms = {str(f) for f in ant if isAtom(f)}
    return any(str(f) in antAtoms for f in suc if isAtom(f))


def proveImproved(seq, depth=0, used=None, seen=None):
    if used is None: used = {}
    if seen is None: seen = set()   # sequent signatures on the current branch

    ant, suc = seq.ant, seq.suc

    #fast axiom check
    if fastAxiom(ant, suc):
        return ProofNode(sequent=seq, rule="id/⊥L", closed=True)

    #loop detection – only needed when quantifiers are present
    hasQuantifier = any(isinstance(f, (Universal, Existential)) for f in ant + suc)
    sig = None
    if hasQuantifier:
        sig = seqSignature(ant, suc)
        if sig in seen:
            return ProofNode(sequent=seq, rule="LOOP", closed=False)
        seen.add(sig)

    node = ProofNode(sequent=seq)

    def cleanup():
        # Remove our signature when we leave this call 
        if sig is not None: seen.discard(sig)

    # sort by complexity so we pick the simplest formula first
    antSorted = sorted(ant, key=complexity) if len(ant) > 1 else list(ant)
    sucSorted = sorted(suc, key=complexity) if len(suc) > 1 else list(suc)

    def recurse(newSeq, d=depth, u=None):
        return proveImproved(newSeq, d, copy.deepcopy(u) if u is not None else copy.deepcopy(used), seen)

    # Group 2
    for f in sucSorted:
        if isinstance(f, Negation):
            ch = recurse(Sequent(ant + [f.formula], dropFirst(suc, f)))
            node.rule, node.children, node.closed = "¬R", [ch], ch.isProved()
            cleanup(); return node
    for f in antSorted:
        if isinstance(f, Negation):
            ch = recurse(Sequent(dropFirst(ant, f), suc + [f.formula]))
            node.rule, node.children, node.closed = "¬L", [ch], ch.isProved()
            cleanup(); return node
    for f in sucSorted:
        if isinstance(f, Implication):
            ch = recurse(Sequent(ant + [f.left], dropFirst(suc, f) + [f.right]))
            node.rule, node.children, node.closed = "→R", [ch], ch.isProved()
            cleanup(); return node
    for f in antSorted:
        if isinstance(f, Conjunction):
            ch = recurse(Sequent(dropFirst(ant, f) + [f.left, f.right], list(suc)))
            node.rule, node.children, node.closed = "∧L", [ch], ch.isProved()
            cleanup(); return node
    for f in sucSorted:
        if isinstance(f, Disjunction):
            ch = recurse(Sequent(list(ant), dropFirst(suc, f) + [f.left, f.right]))
            node.rule, node.children, node.closed = "∨R", [ch], ch.isProved()
            cleanup(); return node
    for f in sucSorted:
        if isinstance(f, Universal):
            c  = freshConstant()
            ch = recurse(Sequent(list(ant), dropFirst(suc, f) + [f.formula.substitute(f.variable, c)]))
            node.rule, node.children, node.closed = f"∀R[{c}]", [ch], ch.isProved()
            cleanup(); return node
    for f in antSorted:
        if isinstance(f, Existential):
            c  = freshConstant()
            ch = recurse(Sequent(dropFirst(ant, f) + [f.formula.substitute(f.variable, c)], list(suc)))
            node.rule, node.children, node.closed = f"∃L[{c}]", [ch], ch.isProved()
            cleanup(); return node

    #Group 3: branching rules 
    for f in sucSorted:
        if isinstance(f, Conjunction):
            rest = dropFirst(suc, f)
            lc = recurse(Sequent(list(ant), rest + [f.left]))
            rc = recurse(Sequent(list(ant), rest + [f.right]))
            node.rule, node.children, node.closed = "∧R", [lc, rc], lc.isProved() and rc.isProved()
            cleanup(); return node
    for f in antSorted:
        if isinstance(f, Disjunction):
            rest = dropFirst(ant, f)
            lc = recurse(Sequent(rest + [f.left],  list(suc)))
            rc = recurse(Sequent(rest + [f.right], list(suc)))
            node.rule, node.children, node.closed = "∨L", [lc, rc], lc.isProved() and rc.isProved()
            cleanup(); return node
    for f in antSorted:
        if isinstance(f, Implication):
            rest = dropFirst(ant, f)
            lc = recurse(Sequent(list(rest),         list(suc) + [f.left]))
            rc = recurse(Sequent(rest + [f.right],   list(suc)))
            node.rule, node.children, node.closed = "→L", [lc, rc], lc.isProved() and rc.isProved()
            cleanup(); return node

    # Group 4
    if depth >= MAX_DEPTH:
        node.rule, node.closed = "DEPTH_LIMIT", False
        cleanup(); return node

    terms = collectGroundTerms(seq) or [freshConstant()]

    # Pick the term that has been used the least across all quantifiers (freshest)
    def termScore(t):
        ts = str(t)
        return sum(1 for f in ant + suc
                   if isinstance(f, (Universal, Existential))
                   and ts in used.get(str(f), set()))

    bestTerm   = min(terms, key=termScore)
    bestTermStr = str(bestTerm)

    newAnt = list(ant)
    newSuc = list(suc)
    nu     = copy.deepcopy(used)
    applied = []   # labels for the rule annotation

    for f in antSorted:
        if isinstance(f, Universal):
            key = str(f)
            if bestTermStr not in nu.get(key, set()):
                newAnt.append(f.formula.substitute(f.variable, bestTerm))
                nu.setdefault(key, set()).add(bestTermStr)
                applied.append(f"∀L[{bestTerm}]")

    for f in sucSorted:
        if isinstance(f, Existential):
            key = str(f)
            if bestTermStr not in nu.get(key, set()):
                newSuc.append(f.formula.substitute(f.variable, bestTerm))
                nu.setdefault(key, set()).add(bestTermStr)
                applied.append(f"∃R[{bestTerm}]")

    # If every known term is exhausted create a fresh one and try again
    if not applied:
        fresh    = freshConstant()
        freshStr = str(fresh)
        for f in antSorted:
            if isinstance(f, Universal):
                key = str(f)
                newAnt.append(f.formula.substitute(f.variable, fresh))
                nu.setdefault(key, set()).add(freshStr)
                applied.append(f"∀L[{fresh}]")
        for f in sucSorted:
            if isinstance(f, Existential):
                key = str(f)
                newSuc.append(f.formula.substitute(f.variable, fresh))
                nu.setdefault(key, set()).add(freshStr)
                applied.append(f"∃R[{fresh}]")

    if not applied:
        node.rule, node.closed = "STUCK", False
        cleanup(); return node

    ruleLabel = ", ".join(applied)
    ch = proveImproved(Sequent(newAnt, newSuc), depth + 1, nu, seen)
    node.rule, node.children, node.closed = ruleLabel, [ch], ch.isProved()
    cleanup(); return node


#  runner –  reads a dataset file and benchmarks both provers


def runDataset(path):
    global freshCounter
    results = parseFile(path)
    if not results:
        print(f"  No formulae found in {path}"); return None

    W = 72
    print()
    print("=" * W)
    print(f"  Dataset: {os.path.basename(path)}  |  Formulae: {len(results)}")
    print("=" * W)
    print(f"  {'Formula':<38} {'Baseline':>12} {'Improved':>12} {'Speedup':>8}")
    print(f"  {'-'*38} {'-'*12} {'-'*12} {'-'*8}")

    baselineProved = improvedProved = errors = 0
    baselineTime   = improvedTime   = 0.0

    for text, seq, err in results:
        if err:
            print(f"  [PARSE ERROR] {text}: {err}"); errors += 1; continue

        label = (text[:35] + "...") if len(text) > 38 else text

        # Run baseline
        freshCounter = 0; complexityCache.clear()
        t0 = time.perf_counter()
        baselineNode = proveBaseline(seq)
        bt = time.perf_counter() - t0
        baselineTime += bt
        if baselineNode.isProved(): baselineProved += 1

        # Run improved
        freshCounter = 0; complexityCache.clear()
        t0 = time.perf_counter()
        improvedNode = proveImproved(Sequent(list(seq.ant), list(seq.suc)))
        it = time.perf_counter() - t0
        improvedTime += it
        if improvedNode.isProved(): improvedProved += 1

        bResult = ("✓" if baselineNode.isProved() else "✗") + f" {bt*1000:5.1f}ms"
        iResult = ("✓" if improvedNode.isProved() else "✗") + f" {it*1000:5.1f}ms"
        speedup = f"{bt/it:.1f}x" if it > 1e-6 else "—"
        print(f"  {label:<38} {bResult:>12} {iResult:>12} {speedup:>8}")

    total = len(results) - errors
    print(f"  {'='*70}")
    print(f"  {'Proved':<38} {str(baselineProved)+'/'+str(total):>12} {str(improvedProved)+'/'+str(total):>12}")
    print(f"  {'Total time':<38} {f'{baselineTime*1000:.1f}ms':>12} {f'{improvedTime*1000:.1f}ms':>12}")
    print("=" * W)

    return dict(
        name          = os.path.basename(path),
        n             = total,
        baselineProved = baselineProved,
        improvedProved = improvedProved,
        baselineMs    = baselineTime * 1000,
        improvedMs    = improvedTime * 1000,
    )


def printSummary(statsList):
    """Print a combined table when multiple datasets were run."""
    W = 100
    print("\n" + "=" * W)
    print("  OVERALL SUMMARY ACROSS ALL DATASETS")
    print("=" * W)
    print(f"  {'Dataset':<28} {'N':>3} {'Base':>8} {'Improv':>9} {'Base time':>12} {'Impr time':>12} {'Speedup':>10}")
    print(f"  {'-'*28} {'-'*6} {'-'*8} {'-'*8} {'-'*11} {'-'*11} {'-'*10}")

    totN = totBase = totImpr = totBms = totIms = 0
    for s in statsList:
        if s is None: continue
        spd = f"{s['baselineMs']/s['improvedMs']:.1f}x" if s['improvedMs'] > 0.01 else "—"
        print(f"  {s['name']:<28} {s['n']:>4} {s['baselineProved']:>6} {s['improvedProved']:>8} "
              f"{s['baselineMs']:>10.1f} {s['improvedMs']:>12.1f} {spd:>13}")
        totN   += s['n']
        totBase += s['baselineProved']
        totImpr += s['improvedProved']
        totBms  += s['baselineMs']
        totIms  += s['improvedMs']

    print("  " + "=" * 98)
    spd = f"{totBms/totIms:.1f}x" if totIms > 0.01 else "—"
    print(f"  {'TOTAL':<28} {totN:>4} {totBase:>6} {totImpr:>8} "
          f"{totBms:>10.1f} {totIms:>12.1f} {spd:>13}")
    print("=" * W)


if __name__ == "__main__":
    files = sys.argv[1:]

    # If no files given on the command line, look for dataset_*.txt in this folder
    if not files:
        here  = os.path.dirname(os.path.abspath(sys.argv[0]))
        files = sorted(
            os.path.join(here, name) for name in os.listdir(here)
            if name.endswith(".txt") and (name.startswith("dataset_") or name == "formulas.txt")
        )

    stats = [runDataset(f) for f in files]
    if len(stats) > 1:
        printSummary(stats)