import yaml
import onnx
from onnx import helper
from pathlib import Path

CONFIG_PATH = Path(__file__).parent.parent / "config.yaml"

# Fixed — graph-specific constants, not runtime inputs
TARGET_TENSOR_NAME = "/model.23/Concat_3_output_0"
OUTPUT_SHAPE       = [1, 15, 8400]


def load_config() -> dict:
    with open(CONFIG_PATH, "r") as f:
        return yaml.safe_load(f)["yolo26_modification"]


def amputate_model(model):
    new_output = helper.make_tensor_value_info(
        TARGET_TENSOR_NAME,
        onnx.TensorProto.FLOAT,
        OUTPUT_SHAPE
    )
    while len(model.graph.output) > 0:
        model.graph.output.pop()
    model.graph.output.extend([new_output])
    return model


def main():
    cfg = load_config()

    print(f"Loading '{cfg['input']}' for surgery...")
    model = onnx.load(cfg["input"])
    model = amputate_model(model)

    onnx.save(model, cfg["output"])
    print(f"\n[SUCCESS] Surgery complete! Amputated model saved to: '{cfg['output']}'")


if __name__ == "__main__":
    main()