"""N1: two arms' drafts on the same captures (native_parity *_drafts.pt): block / per-position agreement per window,
and how often their routing (stage expert sets, window B) differs. Aggregates only.
Usage (x86host): python -m offload.dspark.native.tools.native_compare A_drafts.pt B_drafts.pt OUT.json
"""
from __future__ import annotations

import collections
import json
import sys

import torch


def main() -> None:
    a = torch.load(sys.argv[1], weights_only=False)
    b = torch.load(sys.argv[2], weights_only=False)
    if [r["round"] for r in a] != [r["round"] for r in b]:
        raise SystemExit("the two runs replayed different rounds")
    c = collections.Counter()
    for ra, rb in zip(a, b):
        for w in "AB":
            for da, db in zip(ra[w], rb[w]):
                c[f"{w}_blocks"] += 1
                c[f"{w}_block_equal"] += int(da == db)
                for i, (x, y) in enumerate(zip(da, db)):
                    c[f"{w}_pos{i}_equal"] += int(x == y)
                    if x != y:
                        c[f"{w}_first_diff_pos{i}"] += 1
                        break
        for i, (ea, eb) in enumerate(zip(ra["experts"], rb["experts"])):
            for ta, tb in zip(ea, eb):
                c[f"stage{i}_rows"] += 1
                c[f"stage{i}_expert_set_equal"] += int(sorted(ta) == sorted(tb))
    out = {"a": sys.argv[1].rsplit("/", 1)[-1], "b": sys.argv[2].rsplit("/", 1)[-1], "rounds": len(a)}
    for w in "AB":
        n = max(c[f"{w}_blocks"], 1)
        out[w] = {"blocks": c[f"{w}_blocks"], "block_equal": round(c[f"{w}_block_equal"] / n, 4),
                  "first_diff_pos": [c[f"{w}_first_diff_pos{i}"] for i in range(5)]}
    out["routing_rows_same_expert_set"] = {f"stage{i}": round(c[f"stage{i}_expert_set_equal"] /
                                                              max(c[f"stage{i}_rows"], 1), 4) for i in range(3)}
    with open(sys.argv[3], "w") as f:
        json.dump(out, f, indent=1)
    print(json.dumps(out, indent=1))


if __name__ == "__main__":
    main()