deepcad-canonical-equivalence

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

Builds a canonical signature for each CAD command program and compares the two
programs in a pair. Two programs are equivalent iff their canonical signatures
match.

Semantics implemented (per task spec):
  * param expressions are evaluated to numeric values (tolerance 1e-6 via
    rounding); parameter names are not compared, only values.
  * sketch planes are case-insensitive literal tokens; axes are not reordered.
  * rect/circle/slot profiles preserve their geometric fields + plane.
    Rectangle width/height are ordered. Slot angle is taken modulo 180.
    Unused profiles do not affect the solid.
  * extrude features preserve operation/extent/direction/depth (defaults
    new/one_side/normal), profile geometry, and body/channel partition topology.
    Depth sign and direction are kept separate. Profile list order is not
    semantic. Extrude command order is not semantic (features form a multiset).
  * body/channel labels only induce partition topology; label spelling is not
    preserved and body vs channel is not a semantic namespace. Features with no
    body/channel share one unlabeled default partition; an explicit label does
    not alias it.
  * construction:true commands and annotation kinds (constraint, dimension,
    note, metadata, view, comment) do not affect the solid.
  * any other non-construction command is an opaque solid-affecting payload,
    compared as an unordered multiset after stripping only its top-level id.
    Opaque fields are compared literally (no rename/eval/tolerance/sort).
"""

from __future__ import annotations

import json
import sys
from pathlib import Path


TOL_DECIMALS = 6

PROFILE_KINDS = {"rect", "circle", "slot"}
ANNOTATION_KINDS = {"constraint", "dimension", "note", "metadata", "view", "comment"}
SUPPORTED_KINDS = {"param", "sketch", "extrude"} | PROFILE_KINDS
PROFILE_META_FIELDS = {"cmd", "id", "sketch", "construction"}

DEFAULT_PARTITION = "\x00default"


# --------------------------------------------------------------------------- #
# Expression evaluation
# --------------------------------------------------------------------------- #

def _tokenize(expr):
    tokens = []
    i = 0
    n = len(expr)
    while i < n:
        c = expr[i]
        if c.isspace():
            i += 1
            continue
        if c in "+-*/()":
            tokens.append(c)
            i += 1
            continue
        if c.isdigit() or c == ".":
            j = i
            while j < n and (expr[j].isdigit() or expr[j] == "."):
                j += 1
            tokens.append(("num", float(expr[i:j])))
            i = j
            continue
        if c.isalpha() or c == "_":
            j = i
            while j < n and (expr[j].isalnum() or expr[j] == "_"):
                j += 1
            tokens.append(("id", expr[i:j]))
            i = j
            continue
        raise ValueError(f"bad char {c!r} in expr {expr!r}")
    return tokens


def _eval_expr(expr, get):
    """Evaluate an arithmetic expression string. `get(name)` resolves params."""
    tokens = _tokenize(expr)
    pos = 0

    def peek():
        return tokens[pos] if pos < len(tokens) else None

    def advance():
        nonlocal pos
        t = tokens[pos]
        pos += 1
        return t

    def parse_expr():
        v = parse_term()
        while peek() in ("+", "-"):
            op = advance()
            rhs = parse_term()
            v = v + rhs if op == "+" else v - rhs
        return v

    def parse_term():
        v = parse_factor()
        while peek() in ("*", "/"):
            op = advance()
            rhs = parse_factor()
            v = v * rhs if op == "*" else v / rhs
        return v

    def parse_factor():
        t = peek()
        if t == "+":
            advance()
            return parse_factor()
        if t == "-":
            advance()
            return -parse_factor()
        if t == "(":
            advance()
            v = parse_expr()
            if peek() != ")":
                raise ValueError("missing )")
            advance()
            return v
        if isinstance(t, tuple):
            advance()
            kind, val = t
            if kind == "num":
                return val
            return get(val)
        raise ValueError(f"unexpected token {t!r}")

    v = parse_expr()
    if pos != len(tokens):
        raise ValueError("trailing tokens")
    return v


def _make_env(param_exprs):
    resolved = {}
    resolving = set()

    def get(name):
        if name in resolved:
            return resolved[name]
        if name in resolving:
            raise ValueError(f"cyclic param {name}")
        if name not in param_exprs:
            raise ValueError(f"unknown param {name}")
        resolving.add(name)
        val = _eval_value(param_exprs[name], get)
        resolving.discard(name)
        resolved[name] = val
        return val

    return get


def _eval_value(value, get):
    """Evaluate a scalar/list value to a numeric (rounded) or literal form."""
    if isinstance(value, bool):
        return value
    if isinstance(value, (int, float)):
        return round(float(value), TOL_DECIMALS)
    if isinstance(value, str):
        try:
            return round(_eval_expr(value, get), TOL_DECIMALS)
        except Exception:
            return value
    if isinstance(value, list):
        return [_eval_value(v, get) for v in value]
    if isinstance(value, dict):
        return {k: _eval_value(v, get) for k, v in value.items()}
    return value


# --------------------------------------------------------------------------- #
# Canonicalization
# --------------------------------------------------------------------------- #

def _jkey(obj):
    return json.dumps(obj, sort_keys=True, separators=(",", ":"))


def _is_construction(cmd):
    return bool(cmd.get("construction", False))


def _profile_canon(cmd, sketches, get):
    kind = cmd.get("cmd")
    sketch_id = cmd.get("sketch")
    plane = ""
    if sketch_id is not None and sketch_id in sketches:
        plane = str(sketches[sketch_id]).upper()
    fields = {}
    for k, v in cmd.items():
        if k in PROFILE_META_FIELDS:
            continue
        val = _eval_value(v, get)
        if kind == "slot" and k == "angle" and isinstance(val, (int, float)) and not isinstance(val, bool):
            val = round(((float(val) % 180.0) + 180.0) % 180.0, TOL_DECIMALS)
        fields[k] = val
    return {"kind": kind, "plane": plane, "fields": fields}


def _extrude_feature(cmd, profile_map, get):
    op = str(cmd.get("operation", "new")).lower()
    extent = str(cmd.get("extent", "one_side")).lower()
    direction = str(cmd.get("direction", "normal")).lower()
    depth = _eval_value(cmd.get("depth", 0), get)

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

    prof_canons = [profile_map[pid] for pid in prof_ids if pid in profile_map]
    prof_canons.sort(key=_jkey)

    feature = {
        "op": op,
        "extent": extent,
        "dir": direction,
        "depth": depth,
        "profiles": prof_canons,
    }

    if "body" in cmd:
        label = ("label", _jkey(cmd["body"]))
    elif "channel" in cmd:
        label = ("label", _jkey(cmd["channel"]))
    else:
        label = DEFAULT_PARTITION

    return feature, label


def _program_signature(program):
    # Pass 1: collect params, sketches (ignore construction commands).
    param_exprs = {}
    sketches = {}
    active = []  # non-construction commands
    for cmd in program:
        if not isinstance(cmd, dict):
            continue
        if _is_construction(cmd):
            continue
        active.append(cmd)
        kind = cmd.get("cmd")
        if kind == "param":
            name = cmd.get("name")
            if name is not None:
                param_exprs[name] = cmd.get("expr", 0)
        elif kind == "sketch":
            sid = cmd.get("id")
            if sid is not None:
                sketches[sid] = cmd.get("plane", "")

    get = _make_env(param_exprs)

    # Pass 2: profile canons.
    profile_map = {}
    for cmd in active:
        if cmd.get("cmd") in PROFILE_KINDS:
            pid = cmd.get("id")
            if pid is not None:
                profile_map[pid] = _profile_canon(cmd, sketches, get)

    # Pass 3: extrudes and opaque commands.
    groups = {}
    opaque = []
    for cmd in active:
        kind = cmd.get("cmd")
        if kind == "extrude":
            feature, label = _extrude_feature(cmd, profile_map, get)
            groups.setdefault(label, []).append(feature)
        elif kind in SUPPORTED_KINDS or kind in ANNOTATION_KINDS:
            # param / sketch / profiles already handled or ignored; annotations
            # do not affect the solid.
            continue
        else:
            payload = {k: v for k, v in cmd.items() if k != "id"}
            opaque.append(payload)

    partition = []
    for feats in groups.values():
        feats_sorted = sorted(feats, key=_jkey)
        partition.append(feats_sorted)
    partition.sort(key=_jkey)

    opaque.sort(key=_jkey)

    return {"partition": partition, "opaque": opaque}


# --------------------------------------------------------------------------- #
# Driver
# --------------------------------------------------------------------------- #

def main() -> int:
    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 = _program_signature(pair.get("program_a", []))
            sig_b = _program_signature(pair.get("program_b", []))
            key_a = _jkey(sig_a)
            key_b = _jkey(sig_b)
            equivalent = key_a == key_b
        except Exception:
            key_a = key_b = ""
            equivalent = False
        predictions.append({
            "pair_id": pid,
            "equivalent": bool(equivalent),
            "signature_a": key_a,
            "signature_b": key_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())