deepcad-canonical-equivalence

#!/usr/bin/env python3
"""DeepCAD canonical-equivalence solver.

Reads a JSON file of program pairs and decides, for each pair, whether the two
programs describe the same canonical solid, following the benchmark semantics.
"""

from __future__ import annotations

import json
import re
import sys
from pathlib import Path


# ---------------------------------------------------------------------------
# Expression evaluation (params + numeric fields)
# ---------------------------------------------------------------------------

_TOKEN_RE = re.compile(
    r"\s*(?:(?P<num>\d+\.\d+|\.\d+|\d+\.?|\d+)|(?P<ident>[A-Za-z_][A-Za-z0-9_]*)|(?P<op>[()+\-*/]))"
)


def _tokenize(expr):
    tokens = []
    pos = 0
    n = len(expr)
    while pos < n:
        if expr[pos].isspace():
            pos += 1
            continue
        m = _TOKEN_RE.match(expr, pos)
        if not m or m.start() != pos and m.start() > pos:
            # fall back: try match at pos
            m = _TOKEN_RE.match(expr, pos)
        if not m:
            raise ValueError("cannot tokenize expression: %r at %d" % (expr, pos))
        if m.group("num") is not None:
            tokens.append(("num", float(m.group("num"))))
        elif m.group("ident") is not None:
            tokens.append(("ident", m.group("ident")))
        else:
            tokens.append(("op", m.group("op")))
        pos = m.end()
    return tokens


class _Parser:
    def __init__(self, tokens, resolve):
        self.tokens = tokens
        self.i = 0
        self.resolve = resolve

    def peek(self):
        return self.tokens[self.i] if self.i < len(self.tokens) else (None, None)

    def advance(self):
        tok = self.tokens[self.i]
        self.i += 1
        return tok

    def parse(self):
        val = self.expr()
        if self.i != len(self.tokens):
            raise ValueError("trailing tokens in expression")
        return val

    def expr(self):
        val = self.term()
        while True:
            typ, t = self.peek()
            if typ == "op" and t in ("+", "-"):
                self.advance()
                rhs = self.term()
                val = val + rhs if t == "+" else val - rhs
            else:
                break
        return val

    def term(self):
        val = self.factor()
        while True:
            typ, t = self.peek()
            if typ == "op" and t in ("*", "/"):
                self.advance()
                rhs = self.factor()
                val = val * rhs if t == "*" else val / rhs
            else:
                break
        return val

    def factor(self):
        typ, t = self.peek()
        if typ == "op" and t in ("+", "-"):
            self.advance()
            v = self.factor()
            return v if t == "+" else -v
        if typ == "op" and t == "(":
            self.advance()
            v = self.expr()
            typ2, t2 = self.peek()
            if not (typ2 == "op" and t2 == ")"):
                raise ValueError("expected )")
            self.advance()
            return v
        if typ == "num":
            self.advance()
            return t
        if typ == "ident":
            self.advance()
            return self.resolve(t)
        raise ValueError("unexpected token in expression: %r" % (t,))


def make_evaluator(params):
    """params: dict name -> raw expression (str/number). Returns eval_expr(str)->float."""
    cache = {}
    in_progress = set()

    def resolve(name):
        if name in cache:
            return cache[name]
        if name in in_progress:
            raise ValueError("cyclic param reference: %s" % name)
        if name not in params:
            raise ValueError("unknown param: %s" % name)
        in_progress.add(name)
        val = eval_expr(params[name])
        in_progress.discard(name)
        cache[name] = val
        return val

    def eval_expr(raw):
        if isinstance(raw, bool):
            return float(raw)
        if isinstance(raw, (int, float)):
            return float(raw)
        tokens = _tokenize(str(raw))
        return _Parser(tokens, resolve).parse()

    return eval_expr


def eval_field(val, eval_expr):
    """Evaluate a numeric field: number, expression string, or (nested) list."""
    if isinstance(val, bool):
        return val
    if isinstance(val, (int, float)):
        return float(val)
    if isinstance(val, str):
        return eval_expr(val)
    if isinstance(val, list):
        return [eval_field(x, eval_expr) for x in val]
    return val


def round_val(v, ndigits=6):
    if isinstance(v, bool):
        return v
    if isinstance(v, float):
        r = round(v, ndigits)
        if r == 0.0:
            r = 0.0  # normalise -0.0
        return r
    if isinstance(v, list):
        return [round_val(x, ndigits) for x in v]
    return v


# ---------------------------------------------------------------------------
# Program canonicalisation
# ---------------------------------------------------------------------------

METADATA_KINDS = {"constraint", "dimension", "note", "metadata", "view", "comment"}
SUPPORTED_KINDS = {"param", "sketch", "rect", "circle", "slot", "extrude"}
PROFILE_KINDS = {"rect", "circle", "slot"}


def is_construction(cmd):
    return bool(cmd.get("construction"))


def profile_signature(cmd, plane, eval_expr):
    kind = cmd.get("cmd")
    fields = {}
    for k, v in cmd.items():
        if k in ("cmd", "id", "sketch", "construction"):
            continue
        fields[k] = eval_field(v, eval_expr)
    if kind == "slot" and "angle" in fields:
        try:
            fields["angle"] = float(fields["angle"]) % 180.0
        except (TypeError, ValueError):
            pass
    fields = {k: round_val(v) for k, v in fields.items()}
    sig = {"kind": kind, "plane": plane, "fields": fields}
    return json.dumps(sig, sort_keys=True)


def canon_direction(val, eval_expr):
    if isinstance(val, str):
        return val.lower()
    return round_val(eval_field(val, eval_expr))


def canon_program(prog):
    # Pass 1: params
    params = {}
    for c in prog:
        if not isinstance(c, dict):
            continue
        if c.get("cmd") == "param" and "name" in c:
            params[c["name"]] = c.get("expr", 0)
    eval_expr = make_evaluator(params)

    # Pass 2: sketches (id -> plane, uppercased literal token)
    sketches = {}
    for c in prog:
        if isinstance(c, dict) and c.get("cmd") == "sketch":
            plane = c.get("plane")
            if isinstance(plane, str):
                plane = plane.upper()
            sketches[c.get("id")] = plane

    # Pass 3: profiles (id -> canonical signature)
    profiles = {}
    for c in prog:
        if not isinstance(c, dict):
            continue
        if c.get("cmd") in PROFILE_KINDS:
            plane = sketches.get(c.get("sketch"))
            try:
                profiles[c.get("id")] = profile_signature(c, plane, eval_expr)
            except Exception:
                profiles[c.get("id")] = json.dumps({"kind": c.get("cmd"), "raw": True})

    # Pass 4: extrude features + Pass 5: opaque commands
    features = []  # list of (content_sig_str, partition_key)
    opaque = []
    for c in prog:
        if not isinstance(c, dict):
            continue
        if is_construction(c):
            continue
        kind = c.get("cmd")
        if kind in METADATA_KINDS:
            continue
        if kind == "param" or kind == "sketch" or kind in PROFILE_KINDS:
            continue
        if kind == "extrude":
            features.append(feature_signature(c, profiles, eval_expr))
            continue
        # opaque solid-affecting payload: strip only top-level id
        payload = {k: v for k, v in c.items() if k != "id"}
        opaque.append(json.dumps(payload, sort_keys=True))

    # Partition topology: group features by label (unlabeled -> shared default)
    groups = {}
    for content, label in features:
        key = label if label is not None else ("__DEFAULT__",)
        groups.setdefault(key, []).append(content)
    blocks = [json.dumps(sorted(v)) for v in groups.values()]
    blocks.sort()

    opaque.sort()

    signature = json.dumps({"features": blocks, "opaque": opaque}, sort_keys=True)
    return signature


def feature_signature(cmd, profiles, eval_expr):
    op = str(cmd.get("operation", "new")).lower()
    extent = str(cmd.get("extent", "one_side")).lower()
    direction = canon_direction(cmd.get("direction", "normal"), eval_expr)

    depth_raw = cmd.get("depth", 0)
    try:
        depth = round_val(eval_field(depth_raw, eval_expr))
    except Exception:
        depth = depth_raw

    prof_ids = []
    if "profiles" in cmd and isinstance(cmd["profiles"], list):
        prof_ids.extend(cmd["profiles"])
    if "profile" in cmd:
        p = cmd["profile"]
        if isinstance(p, list):
            prof_ids.extend(p)
        else:
            prof_ids.append(p)

    prof_sigs = []
    for pid in prof_ids:
        prof_sigs.append(profiles.get(pid, json.dumps({"missing": pid})))
    prof_sigs.sort()

    content = json.dumps(
        {
            "operation": op,
            "extent": extent,
            "direction": direction,
            "depth": depth,
            "profiles": prof_sigs,
        },
        sort_keys=True,
    )

    label = None
    if "body" in cmd:
        label = ("label", str(cmd["body"]))
    elif "channel" in cmd:
        label = ("label", str(cmd["channel"]))

    return content, label


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------

def main():
    input_path = Path(sys.argv[1]) if len(sys.argv) > 1 else Path("/workspace/data/public_pairs.json")
    output_path = Path(sys.argv[2]) if len(sys.argv) > 2 else Path("/workspace/predictions.json")

    payload = json.loads(input_path.read_text(encoding="utf-8"))

    predictions = []
    for pair in payload.get("pairs", []):
        pid = pair.get("pair_id")
        try:
            sig_a = canon_program(pair.get("program_a", []))
            sig_b = canon_program(pair.get("program_b", []))
            equivalent = sig_a == sig_b
        except Exception:
            sig_a = sig_b = None
            equivalent = False
        predictions.append(
            {
                "pair_id": pid,
                "equivalent": bool(equivalent),
                "signature_a": sig_a,
                "signature_b": sig_b,
            }
        )

    output_path.parent.mkdir(parents=True, exist_ok=True)
    output_path.write_text(
        json.dumps({"predictions": predictions}, indent=2) + "\n",
        encoding="utf-8",
    )
    return 0


if __name__ == "__main__":
    raise SystemExit(main())