#!/usr/bin/env python3
"""Two-tier per-expert assignment: hot = NVFP4 (4.5bpw), cold = full 2-bpw
AQLM (w13 1x16, w2 1x16). No base tier. Hot count is set byte-exactly from
the expert budget; hot experts chosen by global routing mass with per-layer
floor/cap.

Usage: solve_assignment_2tier.py <stats.npz> [expert_budget_gb]
Output: /data/glm52-expert-assignment.json  {layer: {hot: [...], cold: [...]}}
"""
import json
import os
import sys

import numpy as np

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

COST_HOT, COST_COLD = 21.234, 9.437  # MB/expert
MIN_HOT, MAX_HOT = 16, 176

d = np.load(STATS)
counts = d["counts"].astype(np.float64)
layer_ids = d["layer_ids"].tolist()
plan = json.load(open(os.path.join(ROOT, "hybrid_plan.json")))
whole = set(plan["nvfp4_layers"]) - {78}  # none expected now
keep = [i for i, li in enumerate(layer_ids) if li not in whole]
counts = counts[keep]
layer_ids = [layer_ids[i] for i in keep]
n_layers, n_exp = counts.shape
mass = counts / counts.sum(axis=1, keepdims=True).clip(min=1)

total = n_layers * n_exp
budget_mb = EXPERT_BUDGET_GB * 1e3
n_hot = int((budget_mb - total * COST_COLD) / (COST_HOT - COST_COLD))
print(f"budget {EXPERT_BUDGET_GB} GB -> {n_hot} hot experts "
      f"({n_hot/total:.1%})")

hot = [set() for _ in range(n_layers)]
for li in range(n_layers):
    hot[li].update(np.argsort(mass[li])[::-1][:MIN_HOT].tolist())
placed = sum(len(h) for h in hot)
order = np.dstack(np.unravel_index(
    np.argsort(mass, axis=None)[::-1], mass.shape))[0]
for li, ei in order:
    if placed >= n_hot:
        break
    if len(hot[li]) >= MAX_HOT or int(ei) in hot[li]:
        continue
    hot[li].add(int(ei))
    placed += 1

out = {}
hot_mass = []
for li in range(n_layers):
    h = sorted(int(x) for x in hot[li])
    out[str(layer_ids[li])] = {
        "hot": h,
        "cold": [e for e in range(n_exp) if e not in hot[li]],
    }
    hot_mass.append(mass[li, h].sum())
json.dump(out, open(OUT, "w"))
tot_mb = placed * COST_HOT + (total - placed) * COST_COLD
print(f"hot: {placed} ({placed/total:.1%}) covering mean "
      f"{np.mean(hot_mass):.1%} of routing mass "
      f"(min {np.min(hot_mass):.1%} max {np.max(hot_mass):.1%})")
print(f"expert bytes: {tot_mb/1e3:.1f} GB | hot/layer: "
      f"min {min(len(h) for h in hot)} max {max(len(h) for h in hot)} "
      f"mean {placed/n_layers:.0f}")
print("saved", OUT)