#!/usr/bin/env python3
"""AQUA: Activation-Quantile Utility Allocation for H3 mixed-precision GGUF.

AQUA consumes a stable-diffusion.cpp importance matrix plus the Q8 GGUF tensor
inventory. It starts each eligible matrix at an architecture-specific low-bit
floor, then greedily spends a target bit budget on the upgrades with the
largest activation-weighted utility per added bit. The output is a deterministic
first-match --tensor-type-rules string and a full JSON decision manifest.

This is a tensor-allocation policy, not a new block codec: it composes existing
GGML codecs (Q1_0, IQ1_M, IQ2_S, IQ3_S, Q4_K, Q5_K, Q8_0) according to measured
H3 activation importance.
"""
from __future__ import annotations
import argparse, json, math, re, struct
from collections import defaultdict
from pathlib import Path
import numpy as np
from gguf import GGUFReader

MAIN = re.compile(r"blocks\.(\d+)\.(attn\.(?:qkv_proj|out_proj)|mlp\.(?:fc1|fc2))\.weight$")
ROLE = {
    "attn.qkv_proj": "qkv",
    "attn.out_proj": "out",
    "mlp.fc1": "fc1",
    "mlp.fc2": "fc2",
}
# Effective bits/weight including block metadata, close enough for budget search.
BPW = {"q1_0": 1.125, "iq1_m": 1.75, "iq2_s": 2.5625, "iq3_s": 3.4375,
       "q4_K": 4.5, "q5_K": 5.5, "q8_0": 8.5}
# Monotonic proxy for reconstruction utility. Exact quality is measured later.
QUALITY = {"q1_0": 0.00, "iq1_m": 0.18, "iq2_s": 0.48, "iq3_s": 0.70,
           "q4_K": 0.84, "q5_K": 0.93, "q8_0": 1.00}
LADDERS = {
    "fc1": ["q1_0", "iq1_m", "iq2_s", "iq3_s", "q4_K", "q5_K", "q8_0"],
    "qkv": ["iq2_s", "iq3_s", "q4_K", "q5_K", "q8_0"],
    "out": ["iq3_s", "q4_K", "q5_K", "q8_0"],
    "fc2": ["iq2_s", "iq3_s", "q4_K", "q5_K", "q8_0"],
}
HARD_Q8_BLOCKS = {47, 48, 49}

def read_imatrix(path: Path) -> dict[str, float]:
    out = {}
    with path.open("rb") as f:
        (entries,) = struct.unpack("<i", f.read(4))
        for _ in range(entries):
            (ln,) = struct.unpack("<i", f.read(4))
            name = f.read(ln).decode("utf-8")
            calls, count = struct.unpack("<ii", f.read(8))
            values = np.frombuffer(f.read(count * 4), dtype="<f4")
            if calls > 0 and count > 0 and np.isfinite(values).all():
                out[name] = float(np.mean(np.abs(values)))
        trailer = f.read()
        if len(trailer) != 4:
            raise ValueError(f"unexpected imatrix trailer: {len(trailer)} bytes")
    return out

def percentile_scores(items: list[dict]) -> None:
    by_role = defaultdict(list)
    for item in items:
        by_role[item["role"]].append(item)
    for role, rows in by_role.items():
        rows.sort(key=lambda x: (math.log1p(x["importance"]), x["name"]))
        n = len(rows)
        for i, row in enumerate(rows):
            # Nonzero floor avoids treating the least-active matrix as worthless.
            row["score"] = 0.05 + 0.95 * ((i + 0.5) / n)

def make_rule(names: list[str], qtype: str) -> str | None:
    if not names:
        return None
    escaped = [re.escape(name) + "$" for name in sorted(names)]
    return "(?:" + "|".join(escaped) + ")=" + qtype

def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--gguf", type=Path, required=True)
    ap.add_argument("--imatrix", type=Path, required=True)
    ap.add_argument("--target-bpw", type=float, default=3.15)
    ap.add_argument("--manifest", type=Path, required=True)
    ap.add_argument("--rules-file", type=Path, required=True)
    args = ap.parse_args()

    importance = read_imatrix(args.imatrix)
    reader = GGUFReader(args.gguf, "r")
    items = []
    for tensor in reader.tensors:
        m = MAIN.fullmatch(tensor.name)
        if not m:
            continue
        block = int(m.group(1))
        role = ROLE[m.group(2)]
        items.append({
            "name": tensor.name,
            "block": block,
            "role": role,
            "elements": int(tensor.n_elements),
            "importance": float(importance.get("model.diffusion_model." + tensor.name, importance.get(tensor.name, 0.0))),
        })
    if not items:
        raise SystemExit("no eligible H3 tensors found")
    percentile_scores(items)

    total_elements = sum(x["elements"] for x in items)
    budget = args.target_bpw * total_elements
    used = 0.0
    for x in items:
        ladder = LADDERS[x["role"]]
        if x["block"] in HARD_Q8_BLOCKS:
            idx = len(ladder) - 1
        else:
            idx = 0
        x["ladder"] = ladder
        x["index"] = idx
        x["type"] = ladder[idx]
        used += BPW[x["type"]] * x["elements"]

    # If hard floors exceed the requested budget, retain the safe floors.
    target = max(budget, used)
    while True:
        best = None
        for x in items:
            if x["block"] in HARD_Q8_BLOCKS:
                continue
            idx = x["index"]
            ladder = x["ladder"]
            if idx + 1 >= len(ladder):
                continue
            cur, nxt = ladder[idx], ladder[idx + 1]
            added = (BPW[nxt] - BPW[cur]) * x["elements"]
            if used + added > target + 1e-6:
                continue
            quality_gain = QUALITY[nxt] - QUALITY[cur]
            # High-importance tensors get first refusal; sqrt(elements) prevents
            # a single huge fc1 matrix from swallowing the entire budget.
            utility = x["score"] * math.sqrt(x["elements"]) * quality_gain
            ratio = utility / max(added, 1.0)
            key = (ratio, x["score"], -x["block"], x["name"])
            if best is None or key > best[0]:
                best = (key, x, nxt, added)
        if best is None:
            break
        _, x, nxt, added = best
        x["index"] += 1
        x["type"] = nxt
        used += added

    # First-match protection, followed by exact per-tensor decisions.
    rules = [
        r"condition_proj\.weight$=bf16",
        r"token_refiner\.=q8_0",
        r"(?:audio|video)_patch_proj\.weight$=q8_0",
    ]
    grouped = defaultdict(list)
    for x in items:
        grouped[x["type"]].append(x["name"])
    order = ["q8_0", "q5_K", "q4_K", "iq3_s", "iq2_s", "iq1_m", "q1_0"]
    for qtype in order:
        rule = make_rule(grouped[qtype], qtype)
        if rule:
            rules.append(rule)
    rules_text = ",".join(rules)
    args.rules_file.write_text(rules_text + "\n", encoding="utf-8")

    type_counts = defaultdict(lambda: {"tensors": 0, "parameters": 0})
    for x in items:
        type_counts[x["type"]]["tensors"] += 1
        type_counts[x["type"]]["parameters"] += x["elements"]
    manifest = {
        "schema_version": "1.0",
        "algorithm": "AQUA-v1",
        "expanded_name": "Activation-Quantile Utility Allocation",
        "target_bpw_eligible": args.target_bpw,
        "achieved_estimated_bpw_eligible": used / total_elements,
        "eligible_parameters": total_elements,
        "hard_q8_blocks": sorted(HARD_Q8_BLOCKS),
        "importance_normalization": "within-role percentile of mean absolute activation importance",
        "optimizer": "greedy marginal activation-weighted utility per added bit",
        "type_inventory": dict(sorted(type_counts.items())),
        "rules": rules,
        "decisions": sorted(({k: v for k, v in x.items() if k not in {"ladder", "index"}} for x in items), key=lambda x: x["name"]),
        "note": "AQUA allocates existing GGML codecs; it is a new mixed-precision allocation policy, not a new binary block format.",
    }
    args.manifest.write_text(json.dumps(manifest, indent=2, sort_keys=True) + "\n", encoding="utf-8")
    print(json.dumps({k: manifest[k] for k in ("algorithm", "target_bpw_eligible", "achieved_estimated_bpw_eligible", "eligible_parameters", "type_inventory")}, indent=2))

if __name__ == "__main__":
    main()