#!/usr/bin/env python3
"""A corpus of results to carry from one machine to another.

    tools/golden_corpus generate DIR     write the corpus with this build
    tools/golden_corpus verify DIR       recompute it with this build
    tools/golden_corpus flat DIR FILE    the corpus as one text file for
                                         tests/native/corpus_check.c

``generate`` writes test-vector files (verifloat.vectors, JSON lines) for
every operation on a set of formats that covers every format option, with
random, biased and edge-case operands, and a manifest with each file's
SHA-256. ``verify`` recomputes every vector with the build at hand, through
the scalar operations and, where an operation has them, through the array
kernels at each level this CPU offers. A build on another compiler, C
library, operating system or CPU architecture must reproduce every result
bit and every flag.

Exit status of ``verify``: 0 when everything matches, 1 otherwise.

``flat`` is for places Python does not reach: tests/native/corpus_check.c
is a C program that reads the flat file and checks every vector through the
C library alone (for instance built with MinGW and run on Windows).
"""

import hashlib
import json
import platform
import struct
import sys
import warnings
from pathlib import Path

from verifloat import (BF16, E2M1, E3M2, E4M3, E5M2, FP16, FP32, FP64, UE4M3, UE8M0, FPArray, FPFormat, NaNMode,
                       Rounding, _core, vectors)

FORMATS = [
    FP16, FP32, FP64, BF16, E4M3, E5M2, E3M2, E2M1, UE4M3, UE8M0,
    FPFormat(15, 112),                                                      # binary128
    FPFormat(15, 64, rounding=Rounding.RUP),
    FPFormat(12, 150, rounding=Rounding.RNA),                               # wider than any standard
    FP32.replace(rounding=Rounding.RTZ, nan_mode=NaNMode.X86, tininess="before"),
    FP64.replace(rounding=Rounding.RDN, nan_mode=NaNMode.ARM),
    FP16.replace(rounding=Rounding.RUP, nan_mode=NaNMode.PROPAGATE, ftz=True),
    FPFormat(4, 3, 10, rounding=Rounding.RTZ, ftz=True),
    FPFormat(5, 6, wrap=True, rounding=Rounding.RNA),
    FPFormat(6, 9, saturate=True, inf_nan="fn"),
    FPFormat(3, 4, inf_nan=False, rounding=Rounding.RDN),
    FPFormat(4, 0, inf_nan="fn"),                                           # no mantissa bits
    E4M3.replace(rounding=Rounding.SR, sr_bits=6),
    FP16.replace(rounding=Rounding.SR, sr_bits=3),
]
COUNT = 1500            # biased random vectors per file
ARRAY_OPS = {"add": lambda a, b: a + b, "sub": lambda a, b: a - b, "mul": lambda a, b: a * b,
             "div": lambda a, b: a / b}


def specs():
    """(file name, generate keyword arguments) of every file of the corpus."""
    for i, f in enumerate(FORMATS):
        for op in vectors.OPS:
            kw = {}
            if op == "convert":
                for j, to in enumerate((FP16, FP32, E4M3, FPFormat(15, 112))):
                    yield f"f{i:02d}_convert_{j}", dict(op=op, fmt=f, to=to)
                continue
            if op in ("to_int", "round_to_integral") and f.rounding is Rounding.SR:
                kw["rounding"] = Rounding.RNE
            if op in ("to_int", "from_int"):
                for bits, signed in ((32, True), (64, False), (8, True)):
                    yield f"f{i:02d}_{op}_{'i' if signed else 'u'}{bits}", dict(op=op, fmt=f, bits=bits, signed=signed, **kw)
                continue
            yield f"f{i:02d}_{op}", dict(op=op, fmt=f, **kw)


def sha256(path: Path) -> str:
    return hashlib.sha256(path.read_bytes()).hexdigest()


def generate(root: Path) -> int:
    root.mkdir(parents=True, exist_ok=True)
    manifest, total = {}, 0
    with warnings.catch_warnings():
        warnings.simplefilter("ignore")
        for name, kw in specs():
            kw = dict(kw)
            op, fmt = kw.pop("op"), kw.pop("fmt")
            vecs = list(vectors.generate(op, fmt, COUNT, seed=20261002, style="mixed", **kw))
            spec = vectors.VectorSpec(op, fmt, **kw)
            edges = vectors.generate(op, fmt, 400 if spec.arity > 1 else None, seed=7, style="edges", **kw)
            vecs += list(edges)
            path = root / f"{name}.jsonl"
            vectors.write(path, vecs, spec)
            manifest[path.name] = {"sha256": sha256(path), "vectors": len(vecs)}
            total += len(vecs)
    info = {"made_on": describe(), "files": manifest, "vectors": total}
    (root / "MANIFEST.json").write_text(json.dumps(info, indent=1) + "\n")
    print(f"wrote {total} vectors in {len(manifest)} files to {root}")
    return 0


def describe() -> str:
    return (f"{platform.machine()} {sys.byteorder}-endian, {platform.system()} "
            f"{' '.join(platform.libc_ver()) or 'libc ?'}, Python {platform.python_version()} "
            f"({platform.python_compiler()}), SIMD {_core.simd_available() or 'none'}")


def array_mismatches(path: Path) -> tuple[int, int]:
    """(elements checked, mismatches) of a file's vectors through FPArray at every level."""
    meta = vectors.read_meta(path)
    spec = vectors.VectorSpec.from_meta(meta)
    fn = ARRAY_OPS.get(spec.op)
    if (fn is None and spec.op != "convert") or spec.fmt.size > 64 or spec.uses_sr:
        return 0, 0
    if spec.op == "convert" and spec.to.size > 64:
        return 0, 0
    vecs = list(vectors.read(path))
    want = [(v.result, v.flags) for v in vecs]
    checked = bad = 0
    was = (_core.set_fast(True), _core.simd())
    try:
        for level in ("general", "scalar", *_core.simd_available()):
            _core.set_fast(level != "general")
            _core.set_simd("none" if level in ("general", "scalar") else level)
            a = FPArray.from_raw([v.operands[0] for v in vecs], spec.fmt)
            if fn is None:
                r = a.convert(spec.to)
            else:
                r = fn(a, FPArray.from_raw([v.operands[1] for v in vecs], spec.fmt))
            flags = r._bytes(1)
            got = list(zip(struct.unpack(f"={len(flags)}Q", r._bytes(0)), flags))
            checked += len(got)
            bad += sum(g != w for g, w in zip(got, want))
    finally:
        _core.set_fast(was[0])
        _core.set_simd(was[1])
    return checked, bad


def verify(root: Path) -> int:
    info = json.loads((root / "MANIFEST.json").read_text())
    print(f"corpus made on: {info['made_on']}")
    print(f"verified on:    {describe()}")
    checked = bad = arr = arr_bad = damaged = 0
    with warnings.catch_warnings():
        warnings.simplefilter("ignore")
        for name, entry in sorted(info["files"].items()):
            path = root / name
            if sha256(path) != entry["sha256"]:
                print(f"{name}: the file is not the one the manifest lists")
                damaged += 1
                continue
            report = vectors.verify(path, max_report=3)
            checked += report.checked
            bad += report.mismatch_count
            n, b = array_mismatches(path)
            arr += n
            arr_bad += b
            if report.mismatch_count or b:
                print(f"{name}: {report.mismatch_count} scalar and {b} array mismatches")
                for m in report.mismatches:
                    print("   ", m.describe(report.spec))
    ok = not (bad or arr_bad or damaged) and checked == info["vectors"]
    print(f"scalar operations: {checked} vectors, {bad} mismatches")
    print(f"array kernels:     {arr} elements, {arr_bad} mismatches")
    print("IDENTICAL" if ok else "DIFFERENT")
    return 0 if ok else 1


# Operation numbers of the C library (VF_OP_* in verifloat.h) and its rounding numbers.
C_OPS = {"add": 0, "sub": 1, "mul": 2, "div": 3, "remainder": 4, "fmod": 5, "minimum": 6, "maximum": 7,
         "minimum_number": 8, "maximum_number": 9, "copysign": 10, "fma": 13, "sqrt": 14, "next_up": 17,
         "next_down": 18, "logb": 19}
C_COMPARE = {"compare": (0, 0), "compare_signaling": (0, 1), "eq": (1, 0), "eq_signaling": (1, 1),
             "lt_quiet": (2, 0), "lt": (2, 1), "le_quiet": (3, 0), "le": (3, 1)}
C_ROUND = {Rounding.RNE: 0, Rounding.RNA: 1, Rounding.RTZ: 2, Rounding.RUP: 3, Rounding.RDN: 4}


def flat(root: Path, out: Path) -> int:
    """One line per vector: kind, four parameters, the two format names, then
    sr, a, b, c, result, flags in hex; fields separated by tabs.
    kind 0: vf_op (p1 = operation); 1: comparison (p1 = which, p2 = signaling);
    2: convert; 3: round_to_integral (p1 = rounding, p2 = exact);
    4: to_int (bits, signed, rounding, exact); 5: from_int (bits, signed)."""
    info = json.loads((root / "MANIFEST.json").read_text())
    n = 0
    with open(out, "w") as fh:
        for name in sorted(info["files"]):
            path = root / name
            spec = vectors.VectorSpec.from_meta(vectors.read_meta(path))
            p = [0, 0, 0, 0]
            if spec.op in C_OPS:
                kind, p[0] = 0, C_OPS[spec.op]
            elif spec.op in C_COMPARE:
                kind, (p[0], p[1]) = 1, C_COMPARE[spec.op]
            elif spec.op == "convert":
                kind = 2
            elif spec.op == "round_to_integral":
                kind, p[0], p[1] = 3, C_ROUND[spec.rounding], int(spec.exact)
            elif spec.op == "to_int":
                kind, p = 4, [spec.bits, int(spec.signed), C_ROUND[spec.rounding], int(spec.exact)]
            else:
                kind, p[0], p[1] = 5, spec.bits, int(spec.signed)
            head = "\t".join([str(kind), *map(str, p), str(spec.fmt), str(spec.to if spec.to is not None else spec.fmt)])
            for v in vectors.read(path):
                ops = [*v.operands, 0, 0][:3]
                if kind == 5 and spec.signed and ops[0] >> (spec.bits - 1):
                    ops[0] = (ops[0] - (1 << spec.bits)) & ((1 << 64) - 1)        # sign-extended to 64 bits
                fh.write(f"{head}\t{v.sr or 0:x}\t{ops[0]:x}\t{ops[1]:x}\t{ops[2]:x}\t{v.result:x}\t{v.flags:x}\n")
                n += 1
    print(f"wrote {n} vectors to {out}")
    return 0


def main(argv) -> int:
    if len(argv) == 4 and argv[1] == "flat":
        return flat(Path(argv[2]), Path(argv[3]))
    if len(argv) != 3 or argv[1] not in ("generate", "verify"):
        print(__doc__.split("\n\n")[1], file=sys.stderr)
        return 2
    return (generate if argv[1] == "generate" else verify)(Path(argv[2]))


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