from pathlib import Path

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


ROOT = Path(__file__).resolve().parent
FRAMES = 64


def main() -> None:
    model = onnx.load(ROOT / "data.decoder.onnx")
    model.graph.ClearField("value_info")
    model.graph.input[0].CopyFrom(
        helper.make_tensor_value_info("z", TensorProto.FLOAT, [1, 192, FRAMES])
    )
    model.graph.output[0].CopyFrom(
        helper.make_tensor_value_info(model.graph.output[0].name, TensorProto.FLOAT, [1, 1, FRAMES * 256])
    )

    scale = 1
    nodes = []
    mask_scales = set()
    for node in model.graph.node:
        if node.op_type == "ConvTranspose":
            stride = next(attr.ints[0] for attr in node.attribute if attr.name == "strides")
            scale *= stride
        if node.op_type in {"Conv", "ConvTranspose"}:
            original = node.output[0]
            unmasked = f"{original}_unmasked"
            node.output[0] = unmasked
            nodes.append(node)
            mask_name = f"valid_mask_{scale}"
            nodes.append(helper.make_node("Mul", [unmasked, mask_name], [original], name=f"{node.name}/ValidMask"))
            mask_scales.add(scale)
        else:
            nodes.append(node)

    model.graph.ClearField("node")
    model.graph.node.extend(nodes)
    for current_scale in sorted(mask_scales):
        model.graph.input.append(
            helper.make_tensor_value_info(
                f"valid_mask_{current_scale}", TensorProto.FLOAT, [1, 1, FRAMES * current_scale]
            )
        )

    model = onnx.shape_inference.infer_shapes(model)
    onnx.checker.check_model(model)
    onnx.save(model, ROOT / "data.decoder.fixed64.onnx")
    print("mask scales", sorted(mask_scales))


if __name__ == "__main__":
    main()