import collections
import json
import onnx
from onnx import helper, numpy_helper

model = onnx.load('/data/contextra/models/rfdetr-medium.onnx', load_external_data=False)
graph = model.graph
external = [t for t in graph.initializer if t.data_location == onnx.TensorProto.EXTERNAL]
if external:
    raise RuntimeError('external data detected; inspection stopped before access')
def shape(v):
    t = v.type.tensor_type
    return {'name': v.name, 'dtype': onnx.TensorProto.DataType.Name(t.elem_type),
            'shape': [d.dim_param if d.dim_param else d.dim_value for d in t.shape.dim]}
producers = {o: n for n in graph.node for o in n.output}
consumers = collections.defaultdict(list)
for n in graph.node:
    for i in n.input:
        consumers[i].append(n)
def node(n):
    attrs = {}
    for a in n.attribute:
        if a.type == onnx.AttributeProto.TENSOR:
            if a.t.data_location == onnx.TensorProto.EXTERNAL:
                attrs[a.name] = 'EXTERNAL_NOT_READ'
            elif len(a.t.dims) <= 1 and (not a.t.dims or a.t.dims[0] <= 8):
                attrs[a.name] = numpy_helper.to_array(a.t).tolist()
            else:
                attrs[a.name] = {'dims': list(a.t.dims)}
        else:
            value = helper.get_attribute_value(a)
            attrs[a.name] = value.decode() if isinstance(value, bytes) else value
    return {'name': n.name, 'op': n.op_type, 'inputs': list(n.input), 'outputs': list(n.output), 'attrs': attrs}
def trace(value, depth=5):
    result, seen = [], set()
    def walk(v, d):
        n = producers.get(v)
        if n is None or n.name in seen or d < 0:
            return
        seen.add(n.name)
        result.append(node(n))
        for i in n.input:
            walk(i, d-1)
    walk(value, depth)
    return result
report = {'producer': model.producer_name, 'producer_version': model.producer_version,
          'ir_version': model.ir_version, 'opsets': [(x.domain,x.version) for x in model.opset_import],
          'inputs': [shape(v) for v in graph.input], 'outputs': [shape(v) for v in graph.output],
          'nodes': len(graph.node), 'external_initializers': len(external),
          'ops': dict(collections.Counter(n.op_type for n in graph.node)),
          'input_consumers': {v.name: [node(n) for n in consumers[v.name]] for v in graph.input},
          'output_paths': {v.name: trace(v.name) for v in graph.output}}
cast = next(n for n in consumers['input'] if n.op_type == 'Cast')
report['after_input_cast'] = [node(n) for n in consumers[cast.output[0]]]
report['topk'] = [node(n) for n in graph.node if n.op_type == 'TopK']
report['head_initializers'] = [{'name': t.name, 'dims': list(t.dims)} for t in graph.initializer
                              if t.name in ('class_embed.bias', 'query_feat.weight', 'refpoint_embed.weight')]
onnx.checker.check_model(model)
report['checker'] = 'PASS'
inferred = onnx.shape_inference.infer_shapes(model)
report['inferred_outputs'] = [shape(v) for v in inferred.graph.output]
from pathlib import Path
Path('/tmp/codex/ltb1021/20260909-onnx-graph/graph_summary.json').write_text(json.dumps(report, indent=2))
print(json.dumps({k: v for k,v in report.items() if k not in ('output_paths','ops','input_consumers')}, indent=2))