"""Make a fixed-shape Piper encoder: 64 phoneme IDs in, 160 latent frames out.

Changes to data.encoder.onnx:
- input [1, N_IDS]; padding is masked by input_lengths
- total frame count (ReduceMax of predicted lengths) replaced by constant N_FRAMES;
  speech longer than that is cut off, shorter is masked
- both RandomNormalLike nodes become inputs (noise_dp, noise_z) so windows can be
  re-run with the same noise
- extra outputs: durations [1,1,N_IDS] (frames per ID) and frames [1] (real length)
"""
import sys
from pathlib import Path

import onnx
from onnx import TensorProto, helper

ROOT = Path(__file__).resolve().parent
N_IDS = int(sys.argv[1]) if len(sys.argv) > 1 else 64
N_FRAMES = int(sys.argv[2]) if len(sys.argv) > 2 else 160


def main() -> None:
    model = onnx.load(ROOT / "data.encoder.onnx")
    g = model.graph
    g.ClearField("value_info")

    nodes = []
    for n in g.node:
        if n.op_type == "Cast" and list(n.input) == ["/ReduceMax_output_0"]:
            to = next(a.i for a in n.attribute if a.name == "to")
            dtype = {TensorProto.INT64: TensorProto.INT64, TensorProto.INT32: TensorProto.INT32,
                     TensorProto.FLOAT: TensorProto.FLOAT}[to]
            nodes.append(helper.make_node("Constant", [], list(n.output), name="FixedFrames",
                                          value=helper.make_tensor("fixed_frames", dtype, [], [N_FRAMES])))
        elif n.op_type == "RandomNormalLike" and n.name == "/dp/RandomNormalLike":
            nodes.append(helper.make_node("Identity", ["noise_dp"], list(n.output), name="NoiseDp"))
        elif n.op_type == "RandomNormalLike" and n.name == "/RandomNormalLike":
            nodes.append(helper.make_node("Identity", ["noise_z"], list(n.output), name="NoiseZ"))
        else:
            nodes.append(n)
    assert sum(x.name in ("FixedFrames", "NoiseDp", "NoiseZ") for x in nodes) == 3
    nodes.append(helper.make_node("Identity", ["/Ceil_output_0"], ["durations"], name="DurationsOut"))
    nodes.append(helper.make_node("Identity", ["/Clip_output_0"], ["frames"], name="FramesOut"))
    g.ClearField("node")
    g.node.extend(nodes)

    g.input[0].CopyFrom(helper.make_tensor_value_info("input", TensorProto.INT64, [1, N_IDS]))
    g.input[1].CopyFrom(helper.make_tensor_value_info("input_lengths", TensorProto.INT64, [1]))
    g.input.append(helper.make_tensor_value_info("noise_dp", TensorProto.FLOAT, [1, 2, N_IDS]))
    g.input.append(helper.make_tensor_value_info("noise_z", TensorProto.FLOAT, [1, 192, N_FRAMES]))
    g.output[0].CopyFrom(helper.make_tensor_value_info("z", TensorProto.FLOAT, [1, 192, N_FRAMES]))
    g.output.append(helper.make_tensor_value_info("durations", TensorProto.FLOAT, [1, 1, N_IDS]))
    g.output.append(helper.make_tensor_value_info("frames", TensorProto.FLOAT, [1]))

    model = onnx.shape_inference.infer_shapes(model)
    onnx.checker.check_model(model)
    out = ROOT / f"data.encoder.fixed{N_IDS}x{N_FRAMES}.onnx"
    onnx.save(model, out)
    print("saved", out.name)


if __name__ == "__main__":
    main()