"""plan() on tiny synthetic graphs (no real weights)."""
from __future__ import annotations

import re
import unittest

import numpy as np
import onnx
from onnx import TensorProto, helper, numpy_helper

from tools.maya_export.quant_depth_frame import plan


def _graph(consumers):
    """One fp16 (4, 3) weight 'w' with one Cast per entry of consumers ('mm' = Transpose->MatMul, 'bad' = Add)."""
    w = numpy_helper.from_array(np.ones((4, 3), np.float16), "w")
    nodes, outs = [], []
    for i, kind in enumerate(consumers):
        nodes.append(helper.make_node("Cast", ["w"], [f"wf{i}"], name=f"cast{i}", to=TensorProto.FLOAT))
        if kind == "mm":
            nodes.append(helper.make_node("Transpose", [f"wf{i}"], [f"wt{i}"], name=f"tr{i}", perm=[1, 0]))
            nodes.append(helper.make_node("MatMul", ["x", f"wt{i}"], [f"y{i}"], name=f"mm{i}"))
        else:
            nodes.append(helper.make_node("Add", [f"wf{i}", f"wf{i}"], [f"y{i}"], name=f"add{i}"))
        outs.append(helper.make_tensor_value_info(f"y{i}", TensorProto.FLOAT, None))
    x = helper.make_tensor_value_info("x", TensorProto.FLOAT, [2, 3])
    return helper.make_model(helper.make_graph(nodes, "g", [x], outs, initializer=[w]))


class TestPlan(unittest.TestCase):
    def test_single_cast(self):
        targets, drop = plan(_graph(["mm"]), None)
        self.assertEqual([k for k, _ in targets["w"]], ["matmul"])
        self.assertEqual(drop, {"cast0", "tr0"})

    def test_shared_weight_accumulates_all_uses(self):
        targets, drop = plan(_graph(["mm", "mm"]), None)
        self.assertEqual([n.name for _, n in targets["w"]], ["mm0", "mm1"])
        self.assertEqual(drop, {"cast0", "tr0", "cast1", "tr1"})

    def test_mixed_convertible_and_not_raises(self):
        with self.assertRaises(ValueError):
            plan(_graph(["mm", "bad"]), None)

    def test_keep_regex_skips(self):
        targets, _ = plan(_graph(["mm"]), re.compile("^w$"))
        self.assertEqual(targets, {})


if __name__ == "__main__":
    unittest.main()