"""DeepTIAS 2.7 の検出モデル（TF1 物体検出の frozen graph `Detection_Normal.pb`）を ONNX にする.

TensorFlow と tf2onnx は本体の依存に入れず，実行するときだけ一時的に入れる（サーバで実行する想定）:

    uv run --with tensorflow --with tf2onnx python scripts/convert_detection_to_onnx.py \\
        --pb ~/models/TongueSegmentation/Detection_Normal.pb \\
        --out ~/models/TongueSegmentation/Detection_Normal.onnx

入出力（TF1 物体検出 API の既定．tf2onnx はテンソル名に `:0` を付けたまま ONNX の名前にする）:

- 入力 `image_tensor:0`: `(1, H, W, 3)` の uint8 画像（TF の物体検出 API の既定は RGB）
- 出力 `detection_boxes:0`: `(1, N, 4)` の float32．正規化 `[ymin, xmin, ymax, xmax]`
- 出力 `detection_scores:0`: `(1, N)` の float32（降順）
- 出力 `detection_classes:0`・`num_detections:0`: 変換はするが評価では使わない

`tias-seg eval` の `detection.input_name`・`boxes_name`・`scores_name` は `:0` の有無を問わずに
探すので，既定（`image_tensor`・`detection_boxes`・`detection_scores`）のままでよい．
名前が既定と違うグラフは `--list-nodes` で Placeholder と出力の候補を確かめ，`--inputs`・
`--outputs` で指定する．変換後に onnxruntime で黒い画像を 1 枚推論し，出力の形状を表示する．
変換に失敗した場合でも `tias-seg eval` は検出なし・正解の外接矩形の方式だけで完走する．
"""

from __future__ import annotations

import argparse
import logging
import sys
from pathlib import Path

import numpy as np

logger = logging.getLogger(__name__)

DEFAULT_INPUTS = ("image_tensor:0",)
DEFAULT_OUTPUTS = (
    "detection_boxes:0",
    "detection_scores:0",
    "detection_classes:0",
    "num_detections:0",
)
# TF の物体検出 API の NMS（ループ・制御フロー）を変換できる opset．
DEFAULT_OPSET = 13
# 変換後の確認に使う黒い画像の大きさ（高さ・幅）．
CHECK_IMAGE_SIZE = (480, 640)


def _load_graph_def(path: Path):  # noqa: ANN202  TensorFlow の型は一時的な依存のため書かない
    import tensorflow as tf

    graph_def = tf.compat.v1.GraphDef()
    graph_def.ParseFromString(path.read_bytes())
    return graph_def


def _list_nodes(graph_def) -> None:  # noqa: ANN001
    consumed = {i.split(":")[0].lstrip("^") for n in graph_def.node for i in n.input}
    placeholders = [n.name for n in graph_def.node if n.op == "Placeholder"]
    leaves = [n.name for n in graph_def.node if n.name not in consumed and n.op != "Const"]
    logger.info("Placeholder（入力の候補）: %s", placeholders)
    logger.info("どこからも使われないノード（出力の候補）: %s", leaves[:50])


def _check_onnx(path: Path, input_name: str) -> None:
    import onnxruntime as ort

    session = ort.InferenceSession(str(path), providers=["CPUExecutionProvider"])
    for i in session.get_inputs():
        logger.info("ONNX の入力: %s %s %s", i.name, i.shape, i.type)
    for o in session.get_outputs():
        logger.info("ONNX の出力: %s %s %s", o.name, o.shape, o.type)
    image = np.zeros((1, *CHECK_IMAGE_SIZE, 3), dtype=np.uint8)
    outputs = session.run(None, {input_name: image})
    for o, value in zip(session.get_outputs(), outputs, strict=True):
        logger.info("黒い画像での出力 %s: 形状 %s", o.name, np.asarray(value).shape)


def _parse_args(argv: list[str] | None) -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="TF1 物体検出の frozen graph を ONNX に変換する")
    parser.add_argument("--pb", type=Path, required=True, help="Detection_Normal.pb のパス")
    parser.add_argument("--out", type=Path, help="書き出す ONNX のパス（--list-nodes 以外で必須）")
    parser.add_argument(
        "--inputs", default=",".join(DEFAULT_INPUTS), help="入力テンソル名（カンマ区切り）"
    )
    parser.add_argument(
        "--outputs", default=",".join(DEFAULT_OUTPUTS), help="出力テンソル名（カンマ区切り）"
    )
    parser.add_argument("--opset", type=int, default=DEFAULT_OPSET, help="ONNX の opset")
    parser.add_argument(
        "--list-nodes", action="store_true", help="入力・出力の候補を表示するだけで変換しない"
    )
    args = parser.parse_args(argv)
    if not args.list_nodes and args.out is None:
        parser.error("--out を指定する")
    return args


def main(argv: list[str] | None = None) -> int:
    """エントリポイント.

    Args:
        argv: コマンドライン引数（None なら `sys.argv[1:]`）．

    Returns:
        終了コード（成功 0・失敗 1）．
    """
    logging.basicConfig(
        level=logging.INFO, format="%(asctime)s %(levelname)s %(name)s: %(message)s"
    )
    args = _parse_args(argv)
    pb = args.pb.expanduser()
    if not pb.is_file():
        logger.error("frozen graph が無い: %s", pb)
        return 1
    try:
        graph_def = _load_graph_def(pb)
    except ImportError:
        logger.error(
            "tensorflow が無い．uv run --with tensorflow --with tf2onnx python %s ... で実行する",
            Path(__file__).as_posix(),
        )
        return 1
    if args.list_nodes:
        _list_nodes(graph_def)
        return 0

    import tf2onnx

    out = args.out.expanduser()
    out.parent.mkdir(parents=True, exist_ok=True)
    inputs = [s.strip() for s in args.inputs.split(",") if s.strip()]
    outputs = [s.strip() for s in args.outputs.split(",") if s.strip()]
    logger.info("変換する: 入力 %s・出力 %s・opset %d", inputs, outputs, args.opset)
    tf2onnx.convert.from_graph_def(
        graph_def,
        input_names=inputs,
        output_names=outputs,
        opset=args.opset,
        output_path=str(out),
    )
    logger.info("書き出した: %s", out)
    _check_onnx(out, inputs[0])
    return 0


if __name__ == "__main__":
    sys.exit(main())