#!/usr/bin/env python3
"""Report FFN activation importance from a llama.cpp imatrix GGUF.

The script ranks FFN down and gate/up activation streams from imatrix tensors.
For MoE tensors, it also tries to infer the routed expert count and report the
highest-activation block/expert pairs.
"""

from __future__ import annotations

import argparse
import math
import re
import sys
from collections import Counter, defaultdict
from dataclasses import dataclass
from pathlib import Path
from typing import Iterable

import numpy as np

try:
    from gguf import GGUFReader
except ImportError as exc:  # pragma: no cover - exercised by users without gguf.
    raise SystemExit(
        "Could not import gguf. Run this with a Python environment that has "
        "llama.cpp's gguf package installed, for example:\n"
        "  ~/code/llama.cpp/.venv/bin/python scripts/report-imatrix-activations.py imatrix.gguf"
    ) from exc


BLOCK_PATTERNS = (
    re.compile(r"^blk\.(\d+)\.(.+)\.weight$"),
    re.compile(r"^layers\.(\d+)\.(.+)\.weight$"),
    re.compile(r"^model\.layers\.(\d+)\.(.+)\.weight$"),
)


@dataclass(frozen=True)
class Candidate:
    base: str
    tensor_name: str
    count_name: str
    block: int | None
    module: str
    kind: str
    projection: str
    family: str
    n: int


@dataclass
class Row:
    block: int | None
    family: str
    kind: str
    projection: str
    tensor: str
    count: float
    n: int
    mean_mse: float
    rms: float
    p95: float
    p99: float
    max_value: float
    gt1: float
    gt4: float


@dataclass
class ExpertRow:
    block: int
    expert: int
    family: str
    tensor: str
    count: float
    n: int
    mean_mse: float
    rms: float
    p95: float
    p99: float
    max_value: float


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Rank FFN down and gate/up activation streams from an imatrix GGUF.",
    )
    parser.add_argument("imatrix", type=Path, help="Path to an imatrix .gguf file")
    parser.add_argument(
        "--top",
        type=int,
        default=15,
        help="Rows to show in the top tensor-stream table. Default: 15",
    )
    parser.add_argument(
        "--family-top",
        type=int,
        default=5,
        help="Rows to show per tensor family. Default: 5",
    )
    parser.add_argument(
        "--top-experts",
        type=int,
        default=25,
        help="Rows to show in the combined routed expert table. Default: 25",
    )
    parser.add_argument(
        "--top-expert-streams",
        type=int,
        default=15,
        help="Rows to show in the per-stream routed expert table. Default: 15",
    )
    parser.add_argument(
        "--last-blocks",
        type=int,
        default=10,
        help="Number of final MoE blocks to show in the compact table. Default: 10",
    )
    parser.add_argument(
        "--experts",
        type=int,
        default=None,
        help="Override routed expert count. Useful when it cannot be inferred.",
    )
    parser.add_argument(
        "--no-experts",
        action="store_true",
        help="Skip per-expert routed rankings.",
    )
    parser.add_argument(
        "--include-router-input",
        action="store_true",
        help="Include ffn_gate_inp/router input tensors when present.",
    )
    return parser.parse_args()


def field_type_name(field: object, index: int) -> str:
    try:
        return field.types[index].name
    except Exception:
        return ""


def part_to_list(part: object) -> object:
    return part.tolist() if hasattr(part, "tolist") else part


def decode_bytes(values: object) -> str:
    if isinstance(values, list):
        return bytes(values).decode("utf-8", errors="replace")
    if isinstance(values, np.ndarray):
        return bytes(values.tolist()).decode("utf-8", errors="replace")
    if isinstance(values, bytes):
        return values.decode("utf-8", errors="replace")
    return str(values)


def decode_field(field: object) -> object:
    """Best-effort conversion of a GGUF ReaderField into a printable value."""
    if not getattr(field, "parts", None):
        return ""

    if field_type_name(field, 0) == "STRING":
        return decode_bytes(part_to_list(field.parts[-1]))

    if field_type_name(field, 0) == "ARRAY" and field_type_name(field, 1) == "STRING":
        # llama.cpp imatrix files usually store a single dataset string.
        return decode_bytes(part_to_list(field.parts[-1]))

    value = part_to_list(field.parts[-1])
    if isinstance(value, list) and len(value) == 1:
        return value[0]
    return value


def parse_block_module(base: str) -> tuple[int | None, str]:
    for pattern in BLOCK_PATTERNS:
        match = pattern.match(base)
        if match:
            return int(match.group(1)), match.group(2)
    if base.endswith(".weight"):
        return None, base[: -len(".weight")]
    return None, base


def classify_module(module: str, include_router_input: bool = False) -> tuple[str, str, str] | None:
    """Return (kind, projection, family) for FFN modules that this report handles."""
    if module == "ffn_gate_inp" or module.endswith(".ffn_gate_inp"):
        if not include_router_input:
            return None
        return "router", "gate input", "router gate input"

    if "ffn_down" in module:
        projection = "down"
    elif "ffn_gate" in module or "ffn_up" in module:
        projection = "gate/up"
    else:
        return None

    if "_exps" in module or ".experts." in module or "_experts" in module:
        kind = "routed expert"
    elif "_shexp" in module or "shared_expert" in module or "shared_experts" in module:
        kind = "shared expert"
    else:
        kind = "dense"

    return kind, projection, f"{kind} {projection}"


def collect_lengths(reader: GGUFReader) -> dict[tuple[int | None, str], int]:
    lengths: dict[tuple[int | None, str], int] = {}
    for tensor in reader.tensors:
        if not tensor.name.endswith(".in_sum2"):
            continue
        base = tensor.name[: -len(".in_sum2")]
        block, module = parse_block_module(base)
        lengths[(block, module)] = int(np.asarray(tensor.data).size)
    return lengths


def collect_candidates(reader: GGUFReader, include_router_input: bool) -> list[Candidate]:
    tensor_names = {tensor.name for tensor in reader.tensors}
    candidates: list[Candidate] = []

    for tensor in reader.tensors:
        if not tensor.name.endswith(".in_sum2"):
            continue

        base = tensor.name[: -len(".in_sum2")]
        count_name = base + ".counts"
        if count_name not in tensor_names:
            continue

        block, module = parse_block_module(base)
        classified = classify_module(module, include_router_input=include_router_input)
        if classified is None:
            continue

        kind, projection, family = classified
        candidates.append(
            Candidate(
                base=base,
                tensor_name=tensor.name,
                count_name=count_name,
                block=block,
                module=module,
                kind=kind,
                projection=projection,
                family=family,
                n=int(np.asarray(tensor.data).size),
            )
        )

    return candidates


def infer_expert_count(candidates: Iterable[Candidate], lengths: dict[tuple[int | None, str], int]) -> tuple[int | None, Counter[int]]:
    candidates = list(candidates)
    gate_ref_modules = {
        "attn_gate",
        "attn_k",
        "attn_q",
        "attn_v",
        "ffn_gate",
        "ffn_up",
        "ffn_gate_shexp",
        "ffn_up_shexp",
    }
    down_ref_modules = {
        "ffn_down",
        "ffn_down_shexp",
    }

    gate_refs: dict[int | None, set[int]] = defaultdict(set)
    down_refs: dict[int | None, set[int]] = defaultdict(set)
    for (block, module), size in lengths.items():
        if module in gate_ref_modules:
            gate_refs[block].add(size)
        if module in down_ref_modules:
            down_refs[block].add(size)

    # Prefer dimensions discovered from the FFN candidates themselves. This
    # keeps inference useful for MoE GGUFs that have shared experts, dense
    # leading layers, or non-standard attention tensor availability.
    for candidate in candidates:
        if candidate.block is None or candidate.kind == "routed expert":
            continue
        if candidate.projection == "gate/up":
            gate_refs[candidate.block].add(candidate.n)
        elif candidate.projection == "down":
            down_refs[candidate.block].add(candidate.n)

    votes: Counter[int] = Counter()
    for candidate in candidates:
        if candidate.kind != "routed expert" or candidate.block is None:
            continue
        refs = gate_refs[candidate.block] if candidate.projection == "gate/up" else down_refs[candidate.block]
        for ref in refs:
            if ref <= 0 or candidate.n % ref != 0:
                continue
            expert_count = candidate.n // ref
            if 1 < expert_count <= 4096:
                votes[expert_count] += 1

    if not votes:
        return None, votes
    return votes.most_common(1)[0][0], votes


def deduplicate_gate_up(candidates: Iterable[Candidate]) -> list[Candidate]:
    """Keep one tensor for each shared gate/up input stream."""
    by_key: dict[tuple[int | None, str], Candidate] = {}
    for candidate in candidates:
        key = (candidate.block, candidate.family)
        current = by_key.get(key)
        if current is None:
            by_key[key] = candidate
            continue

        # Prefer gate over up because it is the conventional representative.
        if "ffn_gate" in candidate.module and "ffn_up" in current.module:
            by_key[key] = candidate

    return list(by_key.values())


def stats(values: np.ndarray) -> tuple[float, float, float, float, float, float, float]:
    mean_mse = float(np.mean(values))
    rms = math.sqrt(mean_mse) if mean_mse >= 0 and math.isfinite(mean_mse) else float("nan")
    return (
        mean_mse,
        rms,
        float(np.percentile(values, 95)),
        float(np.percentile(values, 99)),
        float(np.max(values)),
        float(np.mean(values > 1.0)),
        float(np.mean(values > 4.0)),
    )


def analyze(reader: GGUFReader, args: argparse.Namespace) -> tuple[list[Row], list[ExpertRow], int | None, Counter[int]]:
    tensor_by_name = {tensor.name: tensor for tensor in reader.tensors}
    lengths = collect_lengths(reader)
    candidates = collect_candidates(reader, include_router_input=args.include_router_input)
    inferred_experts, expert_votes = infer_expert_count(candidates, lengths)
    expert_count = args.experts if args.experts is not None else inferred_experts
    candidates = deduplicate_gate_up(candidates)

    rows: list[Row] = []
    expert_rows: list[ExpertRow] = []

    for candidate in candidates:
        count = float(np.asarray(tensor_by_name[candidate.count_name].data).reshape(-1)[0])
        raw = np.asarray(tensor_by_name[candidate.tensor_name].data, dtype=np.float64)
        values = raw / count if count else np.full_like(raw, np.nan, dtype=np.float64)
        mean_mse, rms, p95, p99, max_value, gt1, gt4 = stats(values)
        rows.append(
            Row(
                block=candidate.block,
                family=candidate.family,
                kind=candidate.kind,
                projection=candidate.projection,
                tensor=candidate.base,
                count=count,
                n=candidate.n,
                mean_mse=mean_mse,
                rms=rms,
                p95=p95,
                p99=p99,
                max_value=max_value,
                gt1=gt1,
                gt4=gt4,
            )
        )

        if (
            not args.no_experts
            and expert_count
            and candidate.kind == "routed expert"
            and candidate.block is not None
            and candidate.n % expert_count == 0
        ):
            per_expert = values.reshape(expert_count, candidate.n // expert_count)
            for expert_id, expert_values in enumerate(per_expert):
                exp_mean, exp_rms, exp_p95, exp_p99, exp_max, _, _ = stats(expert_values)
                expert_rows.append(
                    ExpertRow(
                        block=candidate.block,
                        expert=expert_id,
                        family=candidate.family,
                        tensor=candidate.base,
                        count=count,
                        n=int(expert_values.size),
                        mean_mse=exp_mean,
                        rms=exp_rms,
                        p95=exp_p95,
                        p99=exp_p99,
                        max_value=exp_max,
                    )
                )

    return rows, expert_rows, expert_count, expert_votes


def fmt(value: object) -> str:
    if isinstance(value, float):
        if math.isnan(value):
            return ""
        return f"{value:.4g}"
    if isinstance(value, int):
        return str(value)
    return str(value)


def pct(value: float) -> str:
    return f"{100.0 * value:.1f}%"


def block_label(block: int | None) -> str:
    return "?" if block is None else str(block)


def print_table(headers: list[str], rows: Iterable[list[object]]) -> None:
    print("| " + " | ".join(headers) + " |")
    print("| " + " | ".join("---" for _ in headers) + " |")
    for row in rows:
        print("| " + " | ".join(fmt(value) for value in row) + " |")


def metadata(reader: GGUFReader) -> dict[str, object]:
    wanted = ("general.type", "imatrix.datasets", "imatrix.chunk_count", "imatrix.chunk_size")
    return {key: decode_field(reader.fields[key]) for key in wanted if key in reader.fields}


def aggregate_experts(expert_rows: Iterable[ExpertRow]) -> list[dict[str, object]]:
    acc: dict[tuple[int, int], dict[str, object]] = {}
    for row in expert_rows:
        key = (row.block, row.expert)
        item = acc.setdefault(
            key,
            {
                "block": row.block,
                "expert": row.expert,
                "weighted_mse_sum": 0.0,
                "n": 0,
                "max_stream_rms": 0.0,
                "streams": [],
            },
        )
        item["weighted_mse_sum"] = float(item["weighted_mse_sum"]) + row.mean_mse * row.n
        item["n"] = int(item["n"]) + row.n
        item["max_stream_rms"] = max(float(item["max_stream_rms"]), row.rms)
        item["streams"].append(row.family)

    result: list[dict[str, object]] = []
    for item in acc.values():
        mean_mse = float(item["weighted_mse_sum"]) / int(item["n"])
        result.append(
            {
                "block": item["block"],
                "expert": item["expert"],
                "rms": math.sqrt(mean_mse),
                "mean_mse": mean_mse,
                "max_stream_rms": item["max_stream_rms"],
                "streams": ", ".join(sorted(set(item["streams"]))),
            }
        )
    return result


def render(path: Path, reader: GGUFReader, rows: list[Row], expert_rows: list[ExpertRow], expert_count: int | None, expert_votes: Counter[int], args: argparse.Namespace) -> None:
    print(f"# {path.name}")
    print()
    print("Score: `RMS = sqrt(mean(in_sum2 / counts))`.")
    print()

    meta = metadata(reader)
    if meta:
        print("Metadata:")
        print()
        print_table(["Key", "Value"], ([key, value] for key, value in meta.items()))
        print()

    blocks = sorted(row.block for row in rows if row.block is not None)
    if blocks:
        print(f"Blocks covered: `{blocks[0]}-{blocks[-1]}`")
        print()

    if expert_count:
        source = "command line" if args.experts is not None else "inferred from tensor lengths"
        print(f"Routed expert count: `{expert_count}` ({source}).")
        if expert_votes and args.experts is None:
            vote_text = ", ".join(f"{count}: {votes}" for count, votes in expert_votes.most_common(5))
            print(f"Expert-count inference votes: `{vote_text}`.")
        print()
    elif not args.no_experts:
        print("Routed expert count could not be inferred. Use `--experts N` to enable per-expert rankings.")
        print()

    rows_sorted = sorted(rows, key=lambda row: row.rms, reverse=True)
    if rows_sorted:
        print(f"Top {min(args.top, len(rows_sorted))} `down` / `gate/up` tensor streams:")
        print()
        print_table(
            ["Block", "Family", "Tensor", "RMS", "Mean MSE", "P95", "Max", "Count"],
            (
                [
                    block_label(row.block),
                    row.family,
                    f"`{row.tensor}`",
                    row.rms,
                    row.mean_mse,
                    row.p95,
                    row.max_value,
                    int(row.count),
                ]
                for row in rows_sorted[: args.top]
            ),
        )
        print()

    if expert_rows:
        print(
            "Per-expert rankings assume llama.cpp-style routed expert tensors are flattened "
            "with each expert stored contiguously. Counts are tensor-level counts, so these "
            "rank aggregate activation contribution per expert stream."
        )
        print()

        aggregated = sorted(aggregate_experts(expert_rows), key=lambda row: float(row["rms"]), reverse=True)
        print(f"Top {min(args.top_experts, len(aggregated))} routed experts, combined across routed down and gate/up:")
        print()
        print_table(
            ["Block", "Expert", "Combined RMS", "Mean MSE", "Max stream RMS", "Streams"],
            (
                [
                    row["block"],
                    row["expert"],
                    row["rms"],
                    row["mean_mse"],
                    row["max_stream_rms"],
                    row["streams"],
                ]
                for row in aggregated[: args.top_experts]
            ),
        )
        print()

        expert_streams = sorted(expert_rows, key=lambda row: row.rms, reverse=True)
        print(f"Top {min(args.top_expert_streams, len(expert_streams))} routed expert streams:")
        print()
        print_table(
            ["Block", "Expert", "Family", "Tensor", "RMS", "Mean MSE", "P95", "Max", "Count"],
            (
                [
                    row.block,
                    row.expert,
                    row.family,
                    f"`{row.tensor}`",
                    row.rms,
                    row.mean_mse,
                    row.p95,
                    row.max_value,
                    int(row.count),
                ]
                for row in expert_streams[: args.top_expert_streams]
            ),
        )
        print()

    families = [
        "dense down",
        "dense gate/up",
        "routed expert down",
        "routed expert gate/up",
        "shared expert down",
        "shared expert gate/up",
        "router gate input",
    ]
    print(f"Top {args.family_top} by family:")
    print()
    family_rows: list[list[object]] = []
    for family in families:
        subset = [row for row in rows if row.family == family]
        for row in sorted(subset, key=lambda item: item.rms, reverse=True)[: args.family_top]:
            family_rows.append(
                [
                    family,
                    block_label(row.block),
                    f"`{row.tensor}`",
                    row.rms,
                    row.mean_mse,
                    row.p95,
                    row.p99,
                    row.max_value,
                    pct(row.gt1),
                    pct(row.gt4),
                    int(row.count),
                ]
            )
    print_table(
        ["Family", "Block", "Tensor", "RMS", "Mean MSE", "P95", "P99", "Max", ">1", ">4", "Count"],
        family_rows,
    )
    print()

    moe_blocks = sorted(
        {
            row.block
            for row in rows
            if row.block is not None and row.kind in {"routed expert", "shared expert"}
        }
    )
    if args.last_blocks > 0 and moe_blocks:
        selected_blocks = moe_blocks[-args.last_blocks :]
        by_block_family = {(row.block, row.family): row for row in rows}
        print(f"Last {len(selected_blocks)} MoE block RMS:")
        print()
        compact_rows: list[list[object]] = []
        for block in selected_blocks:
            compact_rows.append(
                [
                    block,
                    by_block_family.get((block, "shared expert down"), "").rms
                    if (block, "shared expert down") in by_block_family
                    else "",
                    by_block_family.get((block, "shared expert gate/up"), "").rms
                    if (block, "shared expert gate/up") in by_block_family
                    else "",
                    by_block_family.get((block, "routed expert down"), "").rms
                    if (block, "routed expert down") in by_block_family
                    else "",
                    by_block_family.get((block, "routed expert gate/up"), "").rms
                    if (block, "routed expert gate/up") in by_block_family
                    else "",
                ]
            )
        print_table(["Block", "Shared down", "Shared gate/up", "Routed down", "Routed gate/up"], compact_rows)
        print()


def main() -> int:
    args = parse_args()
    if not args.imatrix.exists():
        print(f"error: file not found: {args.imatrix}", file=sys.stderr)
        return 2

    reader = GGUFReader(str(args.imatrix))
    rows, expert_rows, expert_count, expert_votes = analyze(reader, args)
    if not rows:
        print("error: no FFN down/gate/up imatrix tensors were found", file=sys.stderr)
        return 1

    render(args.imatrix, reader, rows, expert_rows, expert_count, expert_votes, args)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())