"""Build one function of the gate's tensor program with a second, independent
tensor compiler (XLA, through its Python package), run it on one device, and
compare every output with the exact one.

    python -I xla_run.py <platform> <module.mlir> <function> <input>... -- <expected>...

<platform> is the package's name for the device kind (`cpu`, `cuda`, `tpu`).
This compiler builds a module's `main`, so the chosen function is renamed
`main` and every other public one made private (and unused). An input is a
`.npy` file, or `<shape>x<type>=@<raw file>` (`16x112xbf16=@w.bin`); an
expected output is a raw file of little-endian 64-bit integers, or
`<shape>x<type>=@<raw file>`. Every output must equal its expected values
exactly, floats included.

Exit status: 0 every output exact; 1 an output differs; 2 a usage error.
"""

import sys

import numpy as np


def main(argv):
    if len(argv) < 4 or "--" not in argv:
        print(__doc__, file=sys.stderr)
        return 2
    platform, path, fn = argv[:3]
    rest = argv[3:]
    k = rest.index("--")
    inputs, expected = rest[:k], rest[k + 1:]

    import jax

    jax.config.update("jax_enable_x64", True)
    import re

    text = open(path).read()
    if text.count(f"func.func @{fn}(") != 1:
        print(f"{path}: expected one public function @{fn}", file=sys.stderr)
        return 2
    text = re.sub(r"func\.func @(\w+)\(", lambda g: "func.func @main(" if g.group(1) == fn else f"func.func private @{g.group(1)}(", text)

    import ml_dtypes

    types = {"bf16": ml_dtypes.bfloat16, "f32": "<f4", "i32": "<i4", "i64": "<i8"}

    def typed(spec):
        """`<shape>x<type>=@<file>` as an array, or None for a plain path."""
        if "=@" not in spec:
            return None
        head, file = spec.split("=@", 1)
        *dims, ty = head.split("x")
        return np.fromfile(file, dtype=types[ty]).reshape([int(d) for d in dims])

    device = jax.devices(platform)[0]
    exe = device.client.compile_and_load(text, [device])
    args = [typed(p) if typed(p) is not None else np.load(p) for p in inputs]
    results = exe.execute([jax.device_put(x, device) for x in args])
    if len(results) != len(expected):
        print(f"{fn}: {len(results)} outputs, {len(expected)} expected", file=sys.stderr)
        return 1
    bad = 0
    for got, name in zip(results, expected):
        want = typed(name)
        if want is None:
            want = np.fromfile(name, dtype="<i8")
        got = np.asarray(got)
        if got.dtype != want.dtype or got.size != want.size:
            print(f"{name}: {got.dtype}{list(got.shape)}, not {want.dtype}{list(want.shape)}")
            bad += 1
            continue
        got = got.reshape(want.shape)
        if not np.array_equal(got, want):
            print(f"{name}: {int(np.sum(got != want))} of {want.size} values differ")
            bad += 1
    print(f"{fn} on {device.platform} ({device.device_kind}): {'exact' if bad == 0 else 'differs'}")
    return 1 if bad else 0


if __name__ == "__main__":
    sys.exit(main(sys.argv[1:]))