#!/usr/bin/env python3
"""Compare opt-in layer dumps. All-row maxima plus last-token cosine."""
import argparse
import json
from pathlib import Path
import numpy as np

parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("baseline", type=Path)
parser.add_argument("candidate", type=Path)
args = parser.parse_args()
records = []
for position in (0, 128, 129):
    for layer in range(24):
        row = {"position": position, "layer": layer}
        for stage in ("attention", "post_attention", "ffn", "output", "routes"):
            name = f"{position}_layer{layer}_{stage}.f32"
            a_path, b_path = args.baseline / name, args.candidate / name
            if not a_path.exists() or not b_path.exists():
                continue
            a = np.fromfile(a_path, dtype=np.float32)
            b = np.fromfile(b_path, dtype=np.float32)
            if a.shape != b.shape:
                raise ValueError(name)
            if stage == "routes":
                a, b = a.reshape(-1, 16)[:, :8], b.reshape(-1, 16)[:, :8]
                row[stage] = {"changed_sets": int(np.any(np.sort(a) != np.sort(b), axis=1).sum()),
                              "total": len(a), "last_old": a[-1].tolist(), "last_new": b[-1].tolist()}
                continue
            last_a, last_b = a[-1536:].astype(np.float64), b[-1536:].astype(np.float64)
            cosine = np.dot(last_a, last_b) / (np.linalg.norm(last_a) * np.linalg.norm(last_b))
            row[stage] = {"last_cosine": float(cosine), "last_max": float(np.abs(last_a-last_b).max()),
                          "all_max": float(np.abs(a-b).max()), "different": int((a != b).sum())}
        if len(row) > 2:
            records.append(row)
print(json.dumps(records, indent=2))