deepcad-canonical-equivalence

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

Reads a pairs JSON, decides for each pair whether program_a and program_b
describe the same canonical solid, and writes predictions.json.
"""

import json
import sys


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

TOL_DECIMALS = 6


def q(x):
    """Quantize a numeric value to the comparison tolerance grid."""
    v = round(float(x), TOL_DECIMALS)
    if v == 0.0:
        v = 0.0  # normalize -0.0 -> 0.0
    return v


# ---------------------------------------------------------------------------
# Arithmetic expression evaluation
# ---------------------------------------------------------------------------
class ExprError(Exception):
    pass


def tokenize(s):
    tokens = []
    i = 0
    n = len(s)
    while i < n:
        c = s[i]
        if c.isspace():
            i += 1
            continue
        if c in "+-*/()":
            tokens.append((c, c))
            i += 1
            continue
        if c.isdigit() or c == ".":
            j = i
            while j < n and (s[j].isdigit() or s[j] in ".eE" or
                             (s[j] in "+-" and j > i and s[j - 1] in "eE")):
                j += 1
            tokens.append(("num", float(s[i:j])))
            i = j
            continue
        if c.isalpha() or c == "_":
            j = i
            while j < n and (s[j].isalnum() or s[j] == "_"):
                j += 1
            tokens.append(("id", s[i:j]))
            i = j
            continue
        raise ExprError("bad char %r in %r" % (c, s))
    return tokens


class Parser:
    def __init__(self, tokens, resolver):
        self.tokens = tokens
        self.pos = 0
        self.resolver = resolver

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

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

    def parse(self):
        val = self.expr()
        if self.pos != len(self.tokens):
            raise ExprError("trailing tokens")
        return val

    def expr(self):
        val = self.term()
        while self.peek()[0] in ("+", "-"):
            op = self.advance()[0]
            rhs = self.term()
            val = val + rhs if op == "+" else val - rhs
        return val

    def term(self):
        val = self.factor()
        while self.peek()[0] in ("*", "/"):
            op = self.advance()[0]
            rhs = self.factor()
            val = val * rhs if op == "*" else val / rhs
        return val

    def factor(self):
        typ, v = self.peek()
        if typ == "+":
            self.advance()
            return self.factor()
        if typ == "-":
            self.advance()
            return -self.factor()
        if typ == "num":
            self.advance()
            return v
        if typ == "id":
            self.advance()
            return self.resolver(v)
        if typ == "(":
            self.advance()
            val = self.expr()
            if self.peek()[0] != ")":
                raise ExprError("missing )")
            self.advance()
            return val
        raise ExprError("unexpected token %r" % (typ,))


class ParamEnv:
    """Resolves parameter names to numeric values with memoization."""

    def __init__(self, params):
        # params: name -> raw expr (string or number)
        self.params = params
        self.cache = {}
        self.stack = set()

    def resolve(self, name):
        if name in self.cache:
            return self.cache[name]
        if name in self.stack:
            raise ExprError("cyclic param %r" % name)
        if name not in self.params:
            raise ExprError("unknown param %r" % name)
        self.stack.add(name)
        val = self.eval(self.params[name])
        self.stack.discard(name)
        self.cache[name] = val
        return val

    def eval(self, raw):
        if isinstance(raw, bool):
            return float(raw)
        if isinstance(raw, (int, float)):
            return float(raw)
        if isinstance(raw, str):
            toks = tokenize(raw)
            return Parser(toks, self.resolve).parse()
        raise ExprError("cannot eval %r" % (raw,))


def eval_value(v, env):
    """Evaluate a geometric field value into a canonical form."""
    if isinstance(v, bool):
        return v
    if isinstance(v, (int, float)):
        return q(v)
    if isinstance(v, str):
        return q(env.eval(v))
    if isinstance(v, list):
        return tuple(eval_value(e, env) for e in v)
    return v


# ---------------------------------------------------------------------------
# Program canonicalization
# ---------------------------------------------------------------------------
PROFILE_SKIP_FIELDS = {"cmd", "id", "sketch", "construction"}


def profile_signature(prof, sketches, env):
    kind = prof.get("cmd")
    plane = sketches.get(prof.get("sketch"))
    fields = []
    for key in sorted(prof.keys()):
        if key in PROFILE_SKIP_FIELDS:
            continue
        val = eval_value(prof[key], env)
        if key == "angle":
            val = angle_mod(val)
        fields.append((key, val))
    return (kind, plane, tuple(fields))


def angle_mod(val):
    if isinstance(val, tuple):
        return tuple(angle_mod(v) for v in val)
    if isinstance(val, (int, float)):
        return q(((float(val) % 180.0) + 180.0) % 180.0)
    return val


def canon_direction(v, env):
    if isinstance(v, str):
        return v.lower()
    if isinstance(v, list):
        return tuple(eval_value(e, env) for e in v)
    if isinstance(v, (int, float)):
        return q(v)
    return v


DEFAULT_PARTITION = ("__default__",)


def partition_key(cmd):
    if "body" in cmd and cmd["body"] is not None:
        return ("label", str(cmd["body"]))
    if "channel" in cmd and cmd["channel"] is not None:
        return ("label", str(cmd["channel"]))
    return DEFAULT_PARTITION


def extrude_content_sig(cmd, profiles, sketches, env):
    operation = str(cmd.get("operation", "new")).lower()
    extent = str(cmd.get("extent", "one_side")).lower()
    direction = canon_direction(cmd.get("direction", "normal"), env)
    depth_raw = cmd.get("depth", None)
    if depth_raw is None:
        depth = None
    else:
        depth = eval_value(depth_raw, env)

    # collect referenced profile ids
    ref_ids = []
    if "profiles" in cmd and cmd["profiles"] is not None:
        ref_ids.extend(cmd["profiles"])
    if "profile" in cmd and cmd["profile"] is not None:
        ref_ids.append(cmd["profile"])

    prof_sigs = []
    for pid in ref_ids:
        prof = profiles.get(pid)
        if prof is None:
            continue
        prof_sigs.append(profile_signature(prof, sketches, env))
    prof_sigs.sort()

    return ("extrude", operation, extent, direction, depth, tuple(prof_sigs))


def opaque_signature(cmd):
    payload = {k: v for k, v in cmd.items() if k != "id"}
    return json.dumps(payload, sort_keys=True, ensure_ascii=True)


def canonicalize(program):
    params = {}
    sketches = {}
    profiles = {}
    extrudes = []
    opaque = []

    for cmd in program:
        if not isinstance(cmd, dict):
            continue
        if cmd.get("construction") is True:
            continue
        kind = cmd.get("cmd")
        if kind in METADATA_KINDS:
            continue
        if kind == "param":
            raw = cmd.get("expr", cmd.get("value"))
            params[cmd.get("name")] = raw
        elif kind == "sketch":
            plane = cmd.get("plane")
            if isinstance(plane, str):
                plane = plane.upper()
            sketches[cmd.get("id")] = plane
        elif kind in PROFILE_KINDS:
            profiles[cmd.get("id")] = cmd
        elif kind == "extrude":
            extrudes.append(cmd)
        else:
            opaque.append(cmd)

    env = ParamEnv(params)

    # Build extrude features with partition grouping.
    groups = {}
    for cmd in extrudes:
        content = extrude_content_sig(cmd, profiles, sketches, env)
        key = partition_key(cmd)
        groups.setdefault(key, []).append(content)

    group_sigs = []
    for key, members in groups.items():
        members.sort()
        group_sigs.append(tuple(members))
    group_sigs.sort()

    opaque_sigs = sorted(opaque_signature(c) for c in opaque)

    return (tuple(group_sigs), tuple(opaque_sigs))


def signature_repr(sig):
    return repr(sig)


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------
def main():
    if len(sys.argv) != 3:
        sys.stderr.write("usage: solve.py <input_json> <output_json>\n")
        sys.exit(2)

    in_path, out_path = sys.argv[1], sys.argv[2]
    with open(in_path, "r") as f:
        data = json.load(f)

    predictions = []
    for pair in data.get("pairs", []):
        pid = pair.get("pair_id")
        try:
            sig_a = canonicalize(pair.get("program_a", []))
            sig_b = canonicalize(pair.get("program_b", []))
            equivalent = sig_a == sig_b
        except Exception as e:  # noqa: BLE001
            sys.stderr.write("error on %s: %s\n" % (pid, e))
            sig_a = sig_b = None
            equivalent = False
        predictions.append({
            "pair_id": pid,
            "equivalent": bool(equivalent),
            "signature_a": signature_repr(sig_a),
            "signature_b": signature_repr(sig_b),
        })

    with open(out_path, "w") as f:
        json.dump({"predictions": predictions}, f, indent=2)


if __name__ == "__main__":
    main()

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

Reads a pairs JSON, decides for each pair whether program_a and program_b
describe the same canonical solid, and writes predictions.json.
"""

import json
import sys


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

TOL_DECIMALS = 6


def q(x):
    """Quantize a numeric value to the comparison tolerance grid."""
    v = round(float(x), TOL_DECIMALS)
    if v == 0.0:
        v = 0.0  # normalize -0.0 -> 0.0
    return v


# ---------------------------------------------------------------------------
# Arithmetic expression evaluation
# ---------------------------------------------------------------------------
class ExprError(Exception):
    pass


def tokenize(s):
    tokens = []
    i = 0
    n = len(s)
    while i < n:
        c = s[i]
        if c.isspace():
            i += 1
            continue
        if c in "+-*/()":
            tokens.append((c, c))
            i += 1
            continue
        if c.isdigit() or c == ".":
            j = i
            while j < n and (s[j].isdigit() or s[j] in ".eE" or
                             (s[j] in "+-" and j > i and s[j - 1] in "eE")):
                j += 1
            tokens.append(("num", float(s[i:j])))
            i = j
            continue
        if c.isalpha() or c == "_":
            j = i
            while j < n and (s[j].isalnum() or s[j] == "_"):
                j += 1
            tokens.append(("id", s[i:j]))
            i = j
            continue
        raise ExprError("bad char %r in %r" % (c, s))
    return tokens


class Parser:
    def __init__(self, tokens, resolver):
        self.tokens = tokens
        self.pos = 0
        self.resolver = resolver

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

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

    def parse(self):
        val = self.expr()
        if self.pos != len(self.tokens):
            raise ExprError("trailing tokens")
        return val

    def expr(self):
        val = self.term()
        while self.peek()[0] in ("+", "-"):
            op = self.advance()[0]
            rhs = self.term()
            val = val + rhs if op == "+" else val - rhs
        return val

    def term(self):
        val = self.factor()
        while self.peek()[0] in ("*", "/"):
            op = self.advance()[0]
            rhs = self.factor()
            val = val * rhs if op == "*" else val / rhs
        return val

    def factor(self):
        typ, v = self.peek()
        if typ == "+":
            self.advance()
            return self.factor()
        if typ == "-":
            self.advance()
            return -self.factor()
        if typ == "num":
            self.advance()
            return v
        if typ == "id":
            self.advance()
            return self.resolver(v)
        if typ == "(":
            self.advance()
            val = self.expr()
            if self.peek()[0] != ")":
                raise ExprError("missing )")
            self.advance()
            return val
        raise ExprError("unexpected token %r" % (typ,))


class ParamEnv:
    """Resolves parameter names to numeric values with memoization."""

    def __init__(self, params):
        # params: name -> raw expr (string or number)
        self.params = params
        self.cache = {}
        self.stack = set()

    def resolve(self, name):
        if name in self.cache:
            return self.cache[name]
        if name in self.stack:
            raise ExprError("cyclic param %r" % name)
        if name not in self.params:
            raise ExprError("unknown param %r" % name)
        self.stack.add(name)
        val = self.eval(self.params[name])
        self.stack.discard(name)
        self.cache[name] = val
        return val

    def eval(self, raw):
        if isinstance(raw, bool):
            return float(raw)
        if isinstance(raw, (int, float)):
            return float(raw)
        if isinstance(raw, str):
            toks = tokenize(raw)
            return Parser(toks, self.resolve).parse()
        raise ExprError("cannot eval %r" % (raw,))


def eval_value(v, env):
    """Evaluate a geometric field value into a canonical form."""
    if isinstance(v, bool):
        return v
    if isinstance(v, (int, float)):
        return q(v)
    if isinstance(v, str):
        return q(env.eval(v))
    if isinstance(v, list):
        return tuple(eval_value(e, env) for e in v)
    return v


# ---------------------------------------------------------------------------
# Program canonicalization
# ---------------------------------------------------------------------------
PROFILE_SKIP_FIELDS = {"cmd", "id", "sketch", "construction"}


def angle_mod(val):
    if isinstance(val, tuple):
        return tuple(angle_mod(v) for v in val)
    if isinstance(val, (int, float)):
        return q(((float(val) % 180.0) + 180.0) % 180.0)
    return val


def profile_signature(prof, sketches, env):
    kind = prof.get("cmd")
    plane = sketches.get(prof.get("sketch"))
    fields = []
    for key in sorted(prof.keys()):
        if key in PROFILE_SKIP_FIELDS:
            continue
        val = eval_value(prof[key], env)
        if key == "angle":
            val = angle_mod(val)
        fields.append((key, val))
    return (kind, plane, tuple(fields))


def canon_direction(v, env):
    if isinstance(v, str):
        return v.lower()
    if isinstance(v, list):
        return tuple(eval_value(e, env) for e in v)
    if isinstance(v, (int, float)):
        return q(v)
    return v


DEFAULT_PARTITION = ("__default__",)


def partition_key(cmd):
    if "body" in cmd and cmd["body"] is not None:
        return ("label", str(cmd["body"]))
    if "channel" in cmd and cmd["channel"] is not None:
        return ("label", str(cmd["channel"]))
    return DEFAULT_PARTITION


def extrude_content_sig(cmd, profiles, sketches, env):
    operation = str(cmd.get("operation", "new")).lower()
    extent = str(cmd.get("extent", "one_side")).lower()
    direction = canon_direction(cmd.get("direction", "normal"), env)
    depth_raw = cmd.get("depth", None)
    depth = None if depth_raw is None else eval_value(depth_raw, env)

    ref_ids = []
    if cmd.get("profiles") is not None:
        ref_ids.extend(cmd["profiles"])
    if cmd.get("profile") is not None:
        ref_ids.append(cmd["profile"])

    prof_sigs = []
    for pid in ref_ids:
        prof = profiles.get(pid)
        if prof is None:
            continue
        prof_sigs.append(profile_signature(prof, sketches, env))
    prof_sigs.sort()

    return ("extrude", operation, extent, direction, depth, tuple(prof_sigs))


def opaque_signature(cmd):
    payload = {k: v for k, v in cmd.items() if k != "id"}
    return json.dumps(payload, sort_keys=True, ensure_ascii=True)


def canonicalize(program):
    params = {}
    sketches = {}
    profiles = {}
    extrudes = []
    opaque = []

    for cmd in program:
        if not isinstance(cmd, dict):
            continue
        if cmd.get("construction") is True:
            continue
        kind = cmd.get("cmd")
        if kind in METADATA_KINDS:
            continue
        if kind == "param":
            raw = cmd.get("expr", cmd.get("value"))
            params[cmd.get("name")] = raw
        elif kind == "sketch":
            plane = cmd.get("plane")
            if isinstance(plane, str):
                plane = plane.upper()
            sketches[cmd.get("id")] = plane
        elif kind in PROFILE_KINDS:
            profiles[cmd.get("id")] = cmd
        elif kind == "extrude":
            extrudes.append(cmd)
        else:
            opaque.append(cmd)

    env = ParamEnv(params)

    # Build extrude features grouped by body/channel partition topology.
    groups = {}
    for cmd in extrudes:
        content = extrude_content_sig(cmd, profiles, sketches, env)
        key = partition_key(cmd)
        groups.setdefault(key, []).append(content)

    group_sigs = []
    for members in groups.values():
        members.sort()
        group_sigs.append(tuple(members))
    group_sigs.sort()

    opaque_sigs = sorted(opaque_signature(c) for c in opaque)

    return (tuple(group_sigs), tuple(opaque_sigs))


def signature_repr(sig):
    return repr(sig)


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------
def main():
    if len(sys.argv) != 3:
        sys.stderr.write("usage: solve.py <input_json> <output_json>\n")
        sys.exit(2)

    in_path, out_path = sys.argv[1], sys.argv[2]
    with open(in_path, "r") as f:
        data = json.load(f)

    predictions = []
    for pair in data.get("pairs", []):
        pid = pair.get("pair_id")
        try:
            sig_a = canonicalize(pair.get("program_a", []))
            sig_b = canonicalize(pair.get("program_b", []))
            equivalent = sig_a == sig_b
        except Exception as e:  # noqa: BLE001
            sys.stderr.write("error on %s: %s\n" % (pid, e))
            sig_a = sig_b = None
            equivalent = False
        predictions.append({
            "pair_id": pid,
            "equivalent": bool(equivalent),
            "signature_a": signature_repr(sig_a),
            "signature_b": signature_repr(sig_b),
        })

    with open(out_path, "w") as f:
        json.dump({"predictions": predictions}, f, indent=2)


if __name__ == "__main__":
    main()