#!/usr/bin/env python3
"""Per-expert precision assignment from routing stats, budget-constrained.

Tiers per expert within each hybrid layer (whole-NVFP4 layers and MTP are
untouched):
  hot  = NVFP4                          21.23 MB/expert  (4.5 bpw)
  base = AQLM w13 1x16, w2 2x16         12.58 MB/expert  (2 / 4 bpw)
  cold = AQLM w13 1x16, w2 1x16          9.44 MB/expert  (2 / 2 bpw)

Greedy global assignment: experts ranked by routing mass; the hottest get
NVFP4, the coldest fund them by dropping w2 to one book, subject to the
total expert-byte budget. Per-layer floor/cap keeps every layer sane.

Output: /data/glm52-expert-assignment.json
  {layer: {"hot": [ids], "cold": [ids]}}  (rest = base)
"""
import json
import os
import sys

import numpy as np

STATS = sys.argv[1] if len(sys.argv) > 1 else "/data/glm52-expert-stats.npz"
OUT = "/data/glm52-expert-assignment.json"
ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))

MB = 1e6
COST_HOT, COST_BASE, COST_COLD = 21.234, 12.583, 9.437  # MB / expert

# weight budget: total <= 292 GB
#   non-expert 38.0 GB, whole-NVFP4 layers (3,4,5,8,74-77) 43.5 GB, MTP 5.4 GB
HYBRID_EXPERT_BUDGET_GB = 310.0 - 38.0 - 43.5 - 5.4

# knobs
HOT_FRACTION_TARGET = 0.08    # ~20 experts/layer on average
MIN_HOT, MAX_HOT = 8, 96      # per-layer bounds
MAX_COLD = 176                # never cold more than this many per layer


def main():
    d = np.load(STATS)
    counts = d["counts"].astype(np.float64)          # [L, 256]
    layer_ids = d["layer_ids"].tolist()
    # exclude whole-NVFP4 layers (incl. MTP) from per-expert assignment
    plan = json.load(open(os.path.join(ROOT, "hybrid_plan.json")))
    whole_nvfp4 = set(plan["nvfp4_layers"])
    keep_rows = [i for i, li in enumerate(layer_ids) if li not in whole_nvfp4]
    counts = counts[keep_rows]
    layer_ids = [layer_ids[i] for i in keep_rows]
    n_layers, n_exp = counts.shape
    mass = counts / counts.sum(axis=1, keepdims=True).clip(min=1)

    total_experts = n_layers * n_exp
    n_hot_target = int(total_experts * HOT_FRACTION_TARGET)

    # global ranking for hot with per-layer floor/cap
    order = np.dstack(np.unravel_index(
        np.argsort(mass, axis=None)[::-1], mass.shape))[0]  # (layer, expert)
    hot = [set() for _ in range(n_layers)]
    # floors first
    for li in range(n_layers):
        top = np.argsort(mass[li])[::-1][:MIN_HOT]
        hot[li].update(top.tolist())
    n_hot = sum(len(h) for h in hot)
    for li, ei in order:
        if n_hot >= n_hot_target:
            break
        if len(hot[li]) >= MAX_HOT or ei in hot[li]:
            continue
        hot[li].add(int(ei))
        n_hot += 1

    # bytes so far with everyone else at base
    base_all = total_experts * COST_BASE
    extra_hot = n_hot * (COST_HOT - COST_BASE)
    budget_bytes = HYBRID_EXPERT_BUDGET_GB * 1e3  # MB
    need_savings = base_all + extra_hot - budget_bytes  # MB to claw back
    per_cold_saving = COST_BASE - COST_COLD

    cold = [set() for _ in range(n_layers)]
    if need_savings > 0:
        n_cold_needed = int(np.ceil(need_savings / per_cold_saving))
        # coldest experts globally (ascending mass), respecting caps
        order_asc = np.dstack(np.unravel_index(
            np.argsort(mass, axis=None), mass.shape))[0]
        n_cold = 0
        for li, ei in order_asc:
            if n_cold >= n_cold_needed:
                break
            ei = int(ei)
            if ei in hot[li] or len(cold[li]) >= MAX_COLD:
                continue
            cold[li].add(ei)
            n_cold += 1

    total = (n_hot * COST_HOT
             + sum(len(c) for c in cold) * COST_COLD
             + (total_experts - n_hot - sum(len(c) for c in cold)) * COST_BASE)
    out = {}
    hot_mass, cold_mass = [], []
    for li in range(n_layers):
        gid = layer_ids[li]
        out[str(gid)] = {
            "hot": sorted(hot[li]),
            "cold": sorted(cold[li]),
        }
        hot_mass.append(mass[li, sorted(hot[li])].sum())
        if cold[li]:
            cold_mass.append(mass[li, sorted(cold[li])].sum())

    json.dump(out, open(OUT, "w"))
    n_cold_t = sum(len(c) for c in cold)
    print(f"hot experts:  {n_hot} ({n_hot/total_experts:.1%}), "
          f"mean mass covered {np.mean(hot_mass):.1%}")
    print(f"cold experts: {n_cold_t} ({n_cold_t/total_experts:.1%}), "
          f"mean mass affected {np.mean(cold_mass) if cold_mass else 0:.2%}")
    print(f"hybrid expert bytes: {total/1e3:.1f} GB "
          f"(budget {HYBRID_EXPERT_BUDGET_GB:.1f} GB)")
    print(f"projected total weights: {total/1e3 + 38.0 + 43.5 + 5.4:.1f} GB (ceiling 310)")
    print("saved", OUT)


if __name__ == "__main__":
    main()