"""Remove data-dependent shapes from the duration predictor's spline flows.

Each spline flow (dp/flows.3/5/7) uses NonZero(mask) to pick the elements inside
the spline range (+-5), runs the spline on just those, then scatters results back.
That makes tensor sizes depend on the data.

Dense rewrite, same output:
- every NonZero is replaced by a constant listing all positions, so gathers and
  scatters always see all N elements
- the spline input is clamped to the spline range so out-of-range elements can't
  index past the bins (their results are discarded below)
- the final "inside" scatter becomes Where(inside_mask, spline_result, previous),
  so outside elements keep the identity value written by the "outside" scatter
"""
import sys
from pathlib import Path

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

ROOT = Path(__file__).resolve().parent
SRC = sys.argv[1] if len(sys.argv) > 1 else "data.encoder.fixed64x160.onnx"
DST = sys.argv[2] if len(sys.argv) > 2 else SRC.replace(".onnx", ".static.onnx")


def main() -> None:
    m = onnx.load(ROOT / SRC)
    g = m.graph
    m = onnx.shape_inference.infer_shapes(m)
    g = m.graph
    shapes = {}
    for vi in list(g.value_info) + list(g.input) + list(g.output):
        shapes[vi.name] = [d.dim_value if d.HasField("dim_value") else None for d in vi.type.tensor_type.shape.dim]
    producer = {o: n for n in g.node for o in n.output}
    consumers = {}
    for n in g.node:
        for x in n.input:
            consumers.setdefault(x, []).append(n)

    def mask_root(name):
        # Follow Expand back to the And / Not node that made the mask.
        n = producer[name]
        while n.op_type == "Expand":
            n = producer[n.input[0]]
        return n

    new_nodes, extra = [], []
    nz_kind = {}  # NonZero output -> "in" / "out"
    for n in g.node:
        if n.op_type != "NonZero":
            continue
        root = mask_root(n.input[0])
        shp = shapes.get(n.input[0])
        if not (shp and all(shp)):
            # Expanded masks: batch and channel are 1 here, same as the source mask.
            shp = shapes.get(root.output[0])
        # dp flows act on one channel of a [1, 2, N] tensor: masks are [1, 1, N].
        shp = [d if d else 1 for d in (shp or [None, None, None])]
        assert len(shp) == 3 and shp[-1] > 1, f"unexpected mask shape for {n.name}: {shp}"
        nz_kind[n.output[0]] = "in" if root.op_type == "And" else "out"
        idx = np.stack(np.meshgrid(*[np.arange(s) for s in shp], indexing="ij")).reshape(len(shp), -1).astype(np.int64)
        n.op_type = "Constant"
        del n.input[:]
        del n.attribute[:]
        n.attribute.append(helper.make_attribute("value", numpy_helper.from_array(idx, n.output[0] + "_all")))

    def idx_kind(idx_name):
        # GatherND/ScatterND indices = Transpose(NonZero) (possibly via Cast/Reshape).
        n = producer.get(idx_name)
        while n is not None and n.output[0] not in nz_kind:
            n = producer.get(n.input[0]) if n.input else None
        return nz_kind.get(n.output[0]) if n is not None else None

    clamped = where_fixed = 0
    nodes = list(g.node)
    out_nodes = []
    for n in nodes:
        if n.op_type == "GatherND" and idx_kind(n.input[1]) == "in" and len(shapes.get(n.input[0], [])) == 3:
            # Spline input x (shape [1,1,N]): clamp into the spline range.
            lo = n.name + "/lo"; hi = n.name + "/hi"; cl = n.input[0] + "_clamped_" + str(clamped)
            out_nodes.append(helper.make_node("Constant", [], [lo], value=numpy_helper.from_array(np.array(-5.0, np.float32))))
            out_nodes.append(helper.make_node("Constant", [], [hi], value=numpy_helper.from_array(np.array(5.0, np.float32))))
            out_nodes.append(helper.make_node("Clip", [n.input[0], lo, hi], [cl], name=n.name + "/Clamp"))
            n.input[0] = cl
            clamped += 1
            out_nodes.append(n)
        elif n.op_type == "ScatterND" and idx_kind(n.input[1]) == "in":
            data, idx, upd = n.input
            shp = shapes.get(data)
            # The mask fed to the (now constant) NonZero: find it via the original Expand node.
            nz = producer[idx]
            while nz.output[0] not in nz_kind:
                nz = producer[nz.input[0]]
            # Rebuild the boolean mask: inside = And(...) expanded to data's shape.
            and_node = None
            for cand in g.node:
                if cand.op_type == "And" and cand.name.rsplit("/", 1)[0] == n.name.rsplit("/", 1)[0]:
                    and_node = cand
            assert and_node is not None, n.name
            if not (shp and all(shp)):
                shp = shapes[and_node.output[0]]
            shp = [d if d else 1 for d in shp]
            shp_name = n.name + "/shape"
            out_nodes.append(helper.make_node("Constant", [], [shp_name], value=numpy_helper.from_array(np.array(shp, np.int64))))
            mask = n.name + "/inside_mask"
            out_nodes.append(helper.make_node("Expand", [and_node.output[0], shp_name], [mask], name=n.name + "/MaskExpand"))
            upd_r = n.name + "/upd_dense"
            out_nodes.append(helper.make_node("Reshape", [upd, shp_name], [upd_r], name=n.name + "/UpdReshape"))
            out_nodes.append(helper.make_node("Where", [mask, upd_r, data], list(n.output), name=n.name + "/Where"))
            where_fixed += 1
        else:
            out_nodes.append(n)
    g.ClearField("node")
    g.node.extend(out_nodes)
    g.ClearField("value_info")
    m = onnx.shape_inference.infer_shapes(m)
    onnx.checker.check_model(m)
    onnx.save(m, ROOT / DST)
    left = sum(1 for n in m.graph.node if n.op_type == "NonZero")
    print(f"NonZero replaced: {len(nz_kind)}; clamps: {clamped}; inside scatters -> Where: {where_fixed}; NonZero left: {left}")
    print("saved", DST)


if __name__ == "__main__":
    main()