#!/usr/bin/env python3
"""Rank H3 transformer matrices by mean activation importance."""

from __future__ import annotations

import argparse
import json
import re
import struct
from pathlib import Path

import numpy as np


FC1 = re.compile(r"model\.diffusion_model\.blocks\.(\d+)\.mlp\.fc1\.weight$")


def read_entries(path: Path) -> list[dict[str, object]]:
    rows: list[dict[str, object]] = []
    with path.open("rb") as stream:
        (entry_count,) = struct.unpack("<i", stream.read(4))
        for _ in range(entry_count):
            (name_length,) = struct.unpack("<i", stream.read(4))
            name = stream.read(name_length).decode("utf-8")
            calls, value_count = struct.unpack("<ii", stream.read(8))
            values = np.frombuffer(stream.read(value_count * 4), dtype="<f4")
            match = FC1.fullmatch(name)
            if match:
                rows.append(
                    {
                        "block": int(match.group(1)),
                        "calls": calls,
                        "values": value_count,
                        "mean": float(values.mean()),
                        "median": float(np.median(values)),
                        "p95": float(np.quantile(values, 0.95)),
                        "max": float(values.max()),
                    }
                )
        # stable-diffusion.cpp appends the collector's final call index.
        trailer = stream.read()
        if len(trailer) != 4:
            raise ValueError(f"unexpected imatrix trailer length: {len(trailer)}")
        struct.unpack("<i", trailer)
    return sorted(rows, key=lambda row: float(row["mean"]), reverse=True)


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("imatrix", type=Path)
    parser.add_argument("--output", type=Path)
    args = parser.parse_args()
    rows = read_entries(args.imatrix)
    rendered = json.dumps(rows, indent=2) + "\n"
    if args.output:
        args.output.write_text(rendered, encoding="utf-8")
    else:
        print(rendered, end="")


if __name__ == "__main__":
    main()