#!/usr/bin/env python3
import subprocess


def parse_mapping(output):
    for line in output.splitlines():
        fields = line.split()
        if (
            len(fields) >= 4
            and fields[0].isdigit()
            and fields[1].isdigit()
            and (fields[2].isdigit() or fields[2] == "-")
            and fields[3].startswith("Ascend")
        ):
            return int(fields[0])
    raise RuntimeError("unable to find an Ascend NPU mapping in npu-smi output")


def detect_npu_index():
    completed = subprocess.run(
        ["npu-smi", "info", "-m"],
        check=True,
        capture_output=True,
        text=True,
    )
    return parse_mapping(completed.stdout)


if __name__ == "__main__":
    print(detect_npu_index())
