#!/usr/bin/env python3
"""Remove the unsupported Log branch from this YOLO depth ONNX export.

For the depth head exported by this project:

  Log(Clip(Exp(Clip(x, lo, hi)), floor, +inf))

is equivalent to Clip(x, max(lo, log(floor)), hi).  The generated AX637
model then contains no AxQuantizedLog CPU fallback.
"""

from __future__ import annotations

import argparse
import math
from pathlib import Path

import numpy as np
import onnx
from onnx import numpy_helper


def scalar_initializer(model: onnx.ModelProto, name: str) -> float:
    """Return a scalar initializer and reject models with a different topology."""
    for initializer in model.graph.initializer:
        if initializer.name == name:
            value = numpy_helper.to_array(initializer)
            if value.size != 1:
                raise ValueError(f"{name} must be scalar, got shape {value.shape}")
            return float(value.reshape(-1)[0])
    raise ValueError(f"missing initializer: {name}")


def fold_depth_log(model: onnx.ModelProto) -> None:
    """Fold the Clip -> Exp -> Clip -> Log sequence used by the depth head."""
    nodes = list(model.graph.node)
    producers = {output: node for node in nodes for output in node.output if output}
    log_node = next(
        (node for node in nodes if node.name == "/model.23/Log" and node.op_type == "Log"),
        None,
    )
    if log_node is None or len(log_node.input) != 1 or len(log_node.output) != 1:
        raise ValueError("expected a single-input /model.23/Log node")

    outer_clip = producers.get(log_node.input[0])
    exp_node = producers.get(outer_clip.input[0]) if outer_clip is not None and outer_clip.input else None
    inner_clip = producers.get(exp_node.input[0]) if exp_node is not None and exp_node.input else None
    if (
        outer_clip is None
        or outer_clip.op_type != "Clip"
        or exp_node is None
        or exp_node.op_type != "Exp"
        or inner_clip is None
        or inner_clip.op_type != "Clip"
        or len(inner_clip.input) != 3
        or len(outer_clip.input) < 2
    ):
        raise ValueError("unexpected depth calibration topology around /model.23/Log")

    inner_min = scalar_initializer(model, inner_clip.input[1])
    floor = scalar_initializer(model, outer_clip.input[1])
    if floor <= 0.0:
        raise ValueError(f"Clip floor must be positive, got {floor}")

    # Do not mutate a potentially shared initializer from elsewhere in the graph.
    folded_min = max(inner_min, math.log(floor))
    min_name = "depth_log_folded_min"
    model.graph.initializer.append(numpy_helper.from_array(np.array(folded_min, dtype=np.float32), min_name))
    inner_clip.input[1] = min_name
    inner_clip.name = "/model.23/Clip_log_folded"
    inner_clip.output[0] = log_node.output[0]

    remove_ids = {id(outer_clip), id(exp_node), id(log_node)}
    del model.graph.node[:]
    model.graph.node.extend(node for node in nodes if id(node) not in remove_ids)


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("input", type=Path, help="source ONNX model")
    parser.add_argument("output", type=Path, help="rewritten ONNX model")
    args = parser.parse_args()

    model = onnx.load(args.input)
    fold_depth_log(model)
    onnx.checker.check_model(model)
    args.output.parent.mkdir(parents=True, exist_ok=True)
    onnx.save(model, args.output)
    print(f"wrote {args.output}")


if __name__ == "__main__":
    main()