#!/usr/bin/env python3
"""Write an LLAMA_EXPERT_LOCK_LIST file: the most used kept experts of the streamed layers.

Slots are positions in the pruned model's expert dimension = index in the ascending list of kept expert ids
(the order prune_experts.py writes). Usage counts come from mimo_expert_importance.json (count_<source>).

usage: python hot_lock_list.py --importance mimo_expert_importance.json --prune prune_..._keep144.json \
          --layers 14-30 --frac 0.6 [--source onpol] --out hot60.txt
"""
import argparse
import json


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--importance", required=True)
    ap.add_argument("--prune", required=True, help="prune list of the model ({'pruned': {layer: [ids]}})")
    ap.add_argument("--layers", required=True, help="streamed layers, e.g. 14-30")
    ap.add_argument("--frac", type=float, default=0.6, help="share of kept experts to lock per layer")
    ap.add_argument("--source", default="onpol", help="usage source: onpol, de, baekpica, ...")
    ap.add_argument("--n-expert", type=int, default=256)
    ap.add_argument("--out", required=True)
    a = ap.parse_args()
    imp = json.load(open(a.importance, encoding="utf-8"))["layers"]
    pruned = json.load(open(a.prune, encoding="utf-8"))["pruned"]
    lo, hi = (int(x) for x in a.layers.split("-"))
    cover = []
    with open(a.out, "w", encoding="ascii") as f:
        for L in range(lo, hi + 1):
            kept = sorted(set(range(a.n_expert)) - set(pruned[str(L)]))
            use = {r["expert"]: r.get(f"count_{a.source}", 0.0) for r in imp[str(L)]}
            order = sorted(kept, key=lambda e: -use.get(e, 0.0))
            hot = order[: round(a.frac * len(kept))]
            slots = sorted(kept.index(e) for e in hot)
            f.write(f"{L} " + " ".join(map(str, slots)) + "\n")
            tot = sum(use.get(e, 0.0) for e in kept) or 1.0
            cover.append(sum(use.get(e, 0.0) for e in hot) / tot)
    print(f"{a.out}: layers {lo}-{hi}, {a.frac:.0%} of kept experts per layer, "
          f"covers {sum(cover) / len(cover):.1%} of {a.source} routings (min {min(cover):.1%})")


if __name__ == "__main__":
    main()