#!/usr/bin/env python3
"""Keep only the actions output in a policy ONNX graph."""

import argparse
from pathlib import Path

import onnx


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--input", type=Path, required=True)
    parser.add_argument("--output", type=Path, required=True)
    parser.add_argument("--output-name", default="actions")
    args = parser.parse_args()

    model = onnx.load(str(args.input))
    outputs = list(model.graph.output)
    if not outputs:
        raise ValueError("ONNX graph has no outputs")

    selected = None
    for output in outputs:
        if output.name == args.output_name:
            selected = output
            break
    if selected is None:
        selected = outputs[0]
        print(f"[WARN] output {args.output_name!r} not found; keeping first output {selected.name!r}")

    del model.graph.output[:]
    model.graph.output.append(selected)

    args.output.parent.mkdir(parents=True, exist_ok=True)
    onnx.save(model, str(args.output))
    print(f"Wrote actions-only ONNX: {args.output}")


if __name__ == "__main__":
    main()
