"""compress_baselines.py -- pruning / quantization baselines on LLaDA-8B-Base.

Methods: wanda, sparsegpt (pruning), gptq, awq (weight-only FAKE quant: quantize ->
dequantize -> fp16). Each nn.Linear -> a dense bf16 weight saved as {prefix}_weight.pt
(loaded by eval/llada_harness.py replace_with_decoded_weights -- auto-detected since no
*_A.pt is present). Uses the SAME calib .pt as the SVD arms (base/ours/asvd) so all
methods are calibration-matched. X^T X (second moment of inputs) is collected in one
forward pass, GPU-resident and memory-binned, offloaded to CPU per bin (like compress.py).

Pruning sparsity default 0.23 == keep 0.77 == the model_ratio 0.8 retention of the SVD arms.
"""
import os, sys, time, argparse
import torch
import torch.nn as nn

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
import common as C


# ── calib X^T X collection (GPU-resident, binned) — same as compress.py ──────────
# NB Traj-MC fix: with the noised (ours) calib ~47% of positions are MASK_ID, whose
# input embedding is a single degenerate vector. Accumulating those rows into X^T X
# injects a giant rank-1 spike (breaks the cholesky-inverse error comp in
# gptq/sparsegpt) and washes out the per-channel outlier signal (wanda/awq under-
# respond). The calib DATA is the full trajectory (t~U[0,1], t=1 full-mask kept) —
# we only restrict the Monte-Carlo *statistic* to the real-token rows, i.e. the
# activations the compressed layer actually has to preserve. mask_id=None (clean
# base calib, 0% mask) is a no-op, so the base arm is byte-for-byte unchanged.
def collect_xtx_cpu(model, input_ids, target_mods, device, log_prefix="", mask_id=None):
    """Legacy path: per-sample D2H copy + CPU accumulate (~33s/sample on A5000).
    Kept only for --xtx_cpu reproduction; collect_xtx_gpu computes the identical
    statistic on GPU."""
    cov, ntok = {}, {}
    _keep = {"m": None}  # per-sample real-token row mask, shape [L], set before each fwd

    def make_hook(name, in_dim):
        def hook(module, inp, out):
            x = inp[0].detach().reshape(-1, in_dim).float()
            k = _keep["m"]
            if k is not None and k.numel() == x.shape[0]:
                x = x[k]                      # drop MASK-position rows
            if x.shape[0] == 0:               # fully-masked sample -> contributes nothing
                return
            if name not in cov:
                cov[name] = torch.zeros(in_dim, in_dim, dtype=torch.float32)  # CPU 累加,GPU 只放模型
                ntok[name] = 0
            cov[name] += (x.t() @ x).cpu()
            ntok[name] += x.shape[0]
        return hook

    handles = [m.register_forward_hook(make_hook(n, m.in_features))
               for n, m in target_mods.items()]
    t0 = time.time()
    with torch.no_grad():
        for i in range(input_ids.shape[0]):
            ids = input_ids[i:i + 1].to(device)
            _keep["m"] = None if mask_id is None else (ids.reshape(-1) != mask_id)
            model(ids)
            torch.cuda.empty_cache()
            if i % 200 == 0:
                print(f"  {log_prefix}collect {i}/{input_ids.shape[0]} ({time.time()-t0:.0f}s)", flush=True)
    for h in handles:
        h.remove()
    out = {n: (cov[n] / ntok[n]) for n in list(cov)}
    for n in list(cov):
        cov[n] = None
    torch.cuda.empty_cache()
    return out


def collect_xtx_gpu(model, input_ids, target_mods, device, log_prefix="", mask_id=None):
    """Same statistic as collect_xtx_cpu (incl. the real-token row restriction), but
    accumulated ON GPU via addmm_ (no per-forward D2H copy), offloaded to CPU once at
    the end — ported from compress.py. Caller must ensure this group's total XtX
    (sum of in_features^2 * 4 bytes) fits on the GPU beside the model (use plan_bins)."""
    cov, ntok = {}, {}
    _keep = {"m": None}  # per-sample real-token row mask, shape [L], set before each fwd

    def make_hook(name, in_dim):
        def hook(module, inp, out):
            x = inp[0].detach().reshape(-1, in_dim).float()
            k = _keep["m"]
            if k is not None and k.numel() == x.shape[0]:
                x = x[k]                      # drop MASK-position rows
            if x.shape[0] == 0:               # fully-masked sample -> contributes nothing
                return
            if name not in cov:
                cov[name] = torch.zeros(in_dim, in_dim, dtype=torch.float32, device=device)
                ntok[name] = 0
            cov[name].addmm_(x.t(), x)  # cov += x^T x, in-place on GPU
            ntok[name] += x.shape[0]
        return hook

    handles = [m.register_forward_hook(make_hook(n, m.in_features))
               for n, m in target_mods.items()]
    t0 = time.time()
    with torch.no_grad():
        for i in range(input_ids.shape[0]):
            ids = input_ids[i:i + 1].to(device)
            _keep["m"] = None if mask_id is None else (ids.reshape(-1) != mask_id)
            model(ids)
            if i % 200 == 0:
                print(f"  {log_prefix}collect {i}/{input_ids.shape[0]} ({time.time()-t0:.0f}s)", flush=True)
    for h in handles:
        h.remove()
    out = {}
    for n in list(cov):
        out[n] = (cov[n] / ntok[n]).cpu()
        cov[n] = None
    del cov
    torch.cuda.empty_cache()
    return out


def plan_bins(targets, budget_bytes):
    bins, cur, cur_b = [], {}, 0
    for name, mod in targets:
        b = mod.in_features * mod.in_features * 4
        if cur and cur_b + b > budget_bytes:
            bins.append(cur); cur, cur_b = {}, 0
        cur[name] = mod; cur_b += b
    if cur:
        bins.append(cur)
    return bins


# ── methods (each returns a dense W' of the SAME shape) ──────────────────────────
def wanda_prune(W, XtX, sparsity, prune_n=2, prune_m=4):
    """Wanda: importance = |W| * ||X_col||.
    prune_n:prune_m = N:M STRUCTURED sparsity (default 2:4 = zero the 2 lowest of every
    4 consecutive input cols -> 50% sparse, hardware-accelerated). prune_n=0 falls back
    to UNSTRUCTURED at `sparsity` (no real speedup)."""
    scaler = torch.sqrt(torch.diag(XtX).clamp(min=0)).reshape(1, -1)  # [1, in]
    metric = W.abs() * scaler
    W = W.clone()
    if prune_n > 0:
        out, inc = W.shape
        g = inc // prune_m
        mt = metric[:, :g * prune_m].reshape(out, g, prune_m)
        idx = mt.topk(prune_n, dim=2, largest=False).indices          # 2 lowest per group of 4
        mask = torch.zeros_like(mt, dtype=torch.bool).scatter_(2, idx, True).reshape(out, g * prune_m)
        Wv = W[:, :g * prune_m]
        Wv[mask] = 0.0
        W[:, :g * prune_m] = Wv
    else:
        k = int(W.shape[1] * sparsity)
        if k > 0:
            idx = torch.topk(metric, k, dim=1, largest=False).indices
            W.scatter_(1, idx, 0.0)
    return W


def sparsegpt_prune(W, H, sparsity, blocksize=128, percdamp=0.01, prune_n=2, prune_m=4):
    """Faithful port of SparseGPT fasterprune (Sink-Aware-Pruning/lib/sparsegpt.py).
    H = X X^T (== our XtX). Column-wise pruning with Hessian error compensation.
    prune_n:prune_m = N:M STRUCTURED (default 2:4, hardware-accelerated). prune_n=0 ->
    UNSTRUCTURED at `sparsity`."""
    W = W.clone().float()
    H = H.clone().float()
    cols = W.shape[1]
    dead = torch.diag(H) == 0
    H[dead, dead] = 1.0
    W[:, dead] = 0.0
    damp = percdamp * torch.mean(torch.diag(H))
    diag = torch.arange(cols, device=W.device)
    H[diag, diag] += damp
    H = torch.linalg.cholesky(H)
    H = torch.cholesky_inverse(H)
    H = torch.linalg.cholesky(H, upper=True)
    Hinv = H
    for i1 in range(0, cols, blocksize):
        i2 = min(i1 + blocksize, cols)
        count = i2 - i1
        W1 = W[:, i1:i2].clone()
        Q1 = torch.zeros_like(W1)
        Err1 = torch.zeros_like(W1)
        Hinv1 = Hinv[i1:i2, i1:i2]
        mask1 = torch.zeros_like(W1, dtype=torch.bool)
        if prune_n == 0:   # unstructured: one threshold over the block
            tmp = W1 ** 2 / (torch.diag(Hinv1).reshape((1, -1))) ** 2
            thresh = torch.sort(tmp.flatten())[0][int(tmp.numel() * sparsity)]
            mask1 = tmp <= thresh
        for i in range(count):
            w = W1[:, i]
            d = Hinv1[i, i]
            if prune_n != 0 and i % prune_m == 0:   # N:M: pick prune_n lowest of the next prune_m cols
                tmp = W1[:, i:i + prune_m] ** 2 \
                    / (torch.diag(Hinv1)[i:i + prune_m].reshape((1, -1))) ** 2
                mask1.scatter_(1, i + torch.topk(tmp, prune_n, dim=1, largest=False).indices, True)
            q = w.clone()
            q[mask1[:, i]] = 0.0
            Q1[:, i] = q
            err1 = (w - q) / d
            W1[:, i:] -= err1.unsqueeze(1).matmul(Hinv1[i, i:].unsqueeze(0))
            Err1[:, i] = err1
        W[:, i1:i2] = Q1
        W[:, i2:] -= Err1.matmul(Hinv[i1:i2, i2:])
    return W


# ── GPTQ (IST-DASLab/gptq) weight-only fake-quant, asym per-row per-group ────────
def _gptq_find_params(W, bits, sym=False):
    maxq = 2 ** bits - 1
    x = W.flatten(1)
    zeros = torch.zeros(x.shape[0], device=x.device)
    xmin = torch.minimum(x.min(1)[0], zeros)
    xmax = torch.maximum(x.max(1)[0], zeros)
    if sym:
        xmax = torch.maximum(xmin.abs(), xmax)
        neg = xmin < 0
        xmin[neg] = -xmax[neg]
    deg = (xmin == 0) & (xmax == 0)
    xmin[deg], xmax[deg] = -1, 1
    scale = (xmax - xmin) / maxq
    zero = torch.full_like(scale, (maxq + 1) / 2) if sym else torch.round(-xmin / scale)
    return scale.unsqueeze(1), zero.unsqueeze(1), maxq


def _gptq_quant(x, scale, zero, maxq):
    q = torch.clamp(torch.round(x / scale) + zero, 0, maxq)
    return scale * (q - zero)


@torch.no_grad()
def gptq_quantize(W, H, bits=3, groupsize=128, percdamp=0.01, blocksize=128, sym=False, actorder=False):
    orig_dtype = W.dtype
    W = W.clone().float()
    H = H.clone().float()
    rows, columns = W.shape
    dev = W.device
    dead = torch.diag(H) == 0
    H[dead, dead] = 1
    W[:, dead] = 0
    invperm = None
    if actorder:
        perm = torch.argsort(torch.diag(H), descending=True)
        W = W[:, perm]; H = H[perm][:, perm]
        invperm = torch.argsort(perm)
    if groupsize == -1:
        scale, zero, maxq = _gptq_find_params(W, bits, sym)
    else:
        maxq = 2 ** bits - 1
        scale = zero = None
    Q = torch.zeros_like(W)
    damp = percdamp * torch.mean(torch.diag(H))
    diag = torch.arange(columns, device=dev)
    H[diag, diag] += damp
    H = torch.linalg.cholesky(H)
    H = torch.cholesky_inverse(H)
    Hinv = torch.linalg.cholesky(H, upper=True)
    for i1 in range(0, columns, blocksize):
        i2 = min(i1 + blocksize, columns)
        count = i2 - i1
        W1 = W[:, i1:i2].clone()
        Q1 = torch.zeros_like(W1)
        Err1 = torch.zeros_like(W1)
        Hinv1 = Hinv[i1:i2, i1:i2]
        for i in range(count):
            w = W1[:, i]
            d = Hinv1[i, i]
            if groupsize != -1 and (i1 + i) % groupsize == 0:
                scale, zero, _ = _gptq_find_params(W[:, (i1 + i):(i1 + i + groupsize)], bits, sym)
            q = _gptq_quant(w.unsqueeze(1), scale, zero, maxq).flatten()
            Q1[:, i] = q
            err1 = (w - q) / d
            W1[:, i:] -= err1.unsqueeze(1).matmul(Hinv1[i, i:].unsqueeze(0))
            Err1[:, i] = err1
        Q[:, i1:i2] = Q1
        W[:, i2:] -= Err1.matmul(Hinv[i1:i2, i2:])
    if actorder:
        Q = Q[:, invperm]
    return Q.to(orig_dtype)


# ── AWQ (mit-han-lab/llm-awq) activation-aware fake-quant, asym group-wise ────────
def _awq_pseudo_quant(w, n_bit=4, group=128, zero_point=True):
    shape = w.shape
    if group > 0:
        assert shape[-1] % group == 0
        w = w.reshape(-1, group)
    if zero_point:
        max_val = w.amax(dim=1, keepdim=True)
        min_val = w.amin(dim=1, keepdim=True)
        max_int, min_int = 2 ** n_bit - 1, 0
        scales = (max_val - min_val).clamp(min=1e-5) / max_int
        zeros = (-torch.round(min_val / scales)).clamp_(min_int, max_int)
        w = (torch.clamp(torch.round(w / scales) + zeros, min_int, max_int) - zeros) * scales
    else:
        max_val = w.abs().amax(dim=1, keepdim=True).clamp(min=1e-5)
        max_int, min_int = 2 ** (n_bit - 1) - 1, -(2 ** (n_bit - 1))
        scales = max_val / max_int
        w = torch.clamp(torch.round(w / scales), min_int, max_int) * scales
    return w.reshape(shape)


@torch.no_grad()
def awq_quantize(W, act_scale, bits=4, groupsize=128, n_grid=20, zero_point=True, H=None, do_clip=True):
    orig_dtype = W.dtype
    W = W.clone().float()
    act_scale = act_scale.float().to(W.device)

    def loss(Wq):
        d = W - Wq
        if H is not None:
            return torch.einsum('oi,ij,oj->', d, H.float(), d).item()
        return (d.pow(2) * act_scale.pow(2).unsqueeze(0)).sum().item()

    best_loss, best_s = float('inf'), torch.ones_like(act_scale)
    for i in range(n_grid):
        ratio = i / n_grid
        s = act_scale.pow(ratio).clamp(min=1e-4)
        s = s / (s.max() * s.min()).sqrt()
        Wq = _awq_pseudo_quant(W * s.unsqueeze(0), bits, groupsize, zero_point) / s.unsqueeze(0)
        L = loss(Wq)
        if L < best_loss:
            best_loss, best_s = L, s
    Ws = W * best_s.unsqueeze(0)
    if do_clip:
        g = groupsize if groupsize > 0 else Ws.shape[1]
        wg = Ws.reshape(Ws.shape[0], -1, g)
        org_max = wg.abs().amax(dim=-1, keepdim=True)
        best_max = org_max.clone()
        best_err = torch.full_like(org_max, float('inf'))
        for i in range(int(0.5 * n_grid)):
            mx = org_max * (1 - i / n_grid)
            q = _awq_pseudo_quant(torch.clamp(wg, -mx, mx).reshape(-1, g), bits, g, zero_point).reshape(wg.shape)
            err = (q - wg).pow(2).mean(dim=-1, keepdim=True)
            better = err < best_err
            best_err = torch.where(better, err, best_err)
            best_max = torch.where(better, mx, best_max)
        Ws = torch.clamp(wg, -best_max, best_max).reshape(Ws.shape)
    Wq = _awq_pseudo_quant(Ws, bits, groupsize, zero_point) / best_s.unsqueeze(0)
    return Wq.to(orig_dtype)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--calib", required=True)
    ap.add_argument("--method", required=True, choices=["wanda", "sparsegpt", "gptq", "awq"])
    ap.add_argument("--sparsity", type=float, default=0.23,
                    help="prune fraction (0.23 == keep 0.77 == model_ratio 0.8 retention)")
    ap.add_argument("--bits", type=int, default=4)
    ap.add_argument("--prune_n", type=int, default=2,
                    help="N:M structured sparsity for wanda/sparsegpt (default 2:4, "
                         "hardware-accelerated). Set 0 for unstructured at --sparsity.")
    ap.add_argument("--prune_m", type=int, default=4)
    ap.add_argument("--groupsize", type=int, default=128)
    ap.add_argument("--layer_type", choices=["all", "attn", "mlp"], default="all")
    ap.add_argument("--model_path", default=C.DEFAULT_MODEL_PATH)
    ap.add_argument("--save_path", required=True)
    ap.add_argument("--nsamples", type=int, default=0)
    ap.add_argument("--xtx_budget_gb", type=float, default=0.0,
                    help="max GB of XtX kept on GPU per collection pass. 0 = AUTO: "
                         "free GPU memory after the model is resident minus headroom, "
                         "so the same command runs on a 24GB or a 48GB card (only the "
                         "number of passes changes -- binning never touches math).")
    ap.add_argument("--xtx_headroom_gb", type=float, default=6.0)
    ap.add_argument("--xtx_cpu", action="store_true",
                    help="legacy per-sample CPU accumulation (slow; repro only)")
    args = ap.parse_args()

    device = "cuda" if torch.cuda.is_available() else "cpu"
    os.makedirs(args.save_path, exist_ok=True)
    t0 = time.time()

    blob = torch.load(args.calib, map_location="cpu")
    input_ids = blob["input_ids"]
    if args.nsamples and args.nsamples < input_ids.shape[0]:
        input_ids = input_ids[: args.nsamples]
    arm = blob.get("arm", "unknown")
    print(f"[baselines] method={args.method} calib={os.path.basename(args.calib)} "
          f"arm={arm} N={input_ids.shape[0]} sparsity={args.sparsity} bits={args.bits}", flush=True)

    model, tok = C.load_model(args.model_path, device=device)
    targets = list(C.iter_target_linears(model, args.layer_type))
    print(f"[baselines] target linears ({args.layer_type}) = {len(targets)}", flush=True)

    if args.xtx_cpu:
        cov = collect_xtx_cpu(model, input_ids, dict(targets), device, mask_id=C.MASK_ID)
    else:
        budget_gb = args.xtx_budget_gb
        if budget_gb <= 0:
            if device == "cuda":
                free_b, total_b = torch.cuda.mem_get_info()
                budget_gb = max(2.0, free_b / 1024 ** 3 - args.xtx_headroom_gb)
                print(f"[baselines] AUTO xtx budget: free={free_b/1024**3:.1f}GB "
                      f"total={total_b/1024**3:.1f}GB headroom={args.xtx_headroom_gb}GB "
                      f"-> budget={budget_gb:.1f}GB", flush=True)
            else:
                budget_gb = 8.0
        bins = plan_bins(targets, budget_gb * 1024 ** 3)
        print(f"[baselines] {len(targets)} target linears -> {len(bins)} XtX pass(es) "
              f"(budget {budget_gb:.1f}GB/pass)", flush=True)
        cov = {}
        for bi, b in enumerate(bins):
            gb = sum(m.in_features ** 2 * 4 for m in b.values()) / 1024 ** 3
            print(f"[baselines] XtX pass {bi+1}/{len(bins)}: {len(b)} layers, {gb:.1f}GB", flush=True)
            cov.update(collect_xtx_gpu(model, input_ids, b, device,
                                       log_prefix=f"p{bi+1} ", mask_id=C.MASK_ID))
    print(f"[baselines] XtX collected for {len(cov)} layers ({time.time()-t0:.0f}s)", flush=True)

    n_done = 0
    for name, mod in targets:
        if name not in cov:
            continue
        W = mod.weight.data.cpu().float().to(device)
        H = cov[name].to(device)
        if args.method == "wanda":
            Wq = wanda_prune(W, H, args.sparsity, args.prune_n, args.prune_m)
        elif args.method == "sparsegpt":
            Wq = sparsegpt_prune(W, H, args.sparsity, prune_n=args.prune_n, prune_m=args.prune_m)
        elif args.method == "gptq":
            Wq = gptq_quantize(W, H, args.bits, args.groupsize)
        elif args.method == "awq":
            act = torch.sqrt(torch.diag(H).clamp(min=0))  # per-channel RMS ~ mean|x| proxy
            Wq = awq_quantize(W, act, args.bits, args.groupsize, H=H)
        torch.save(Wq.to(torch.bfloat16).cpu(),
                   os.path.join(args.save_path, f"{name.replace('.', '_')}_weight.pt"))
        del W, H, Wq
        torch.cuda.empty_cache()
        n_done += 1
        if n_done % 30 == 0:
            print(f"  [{n_done}/{len(targets)}] {name}", flush=True)

    _structured = (args.method in ("wanda", "sparsegpt") and args.prune_n > 0)
    summary = {"method": args.method, "arm": arm, "sparsity": args.sparsity,
               "bits": args.bits, "groupsize": args.groupsize,
               "structured": f"{args.prune_n}:{args.prune_m}" if _structured else "unstructured",
               "real_sparsity": (args.prune_n / args.prune_m) if _structured else args.sparsity,
               "layer_type": args.layer_type, "n_layers": n_done,
               "calib_file": os.path.basename(args.calib), "git_hash": C.git_hash(),
               "wall_clock_sec": time.time() - t0, "save_path": args.save_path}
    C.dump_json(summary, os.path.join(args.save_path, "compression_summary.json"))
    print(f"[baselines] done: {n_done} layers -> {args.save_path}  ({(time.time()-t0)/60:.1f}min)", flush=True)


if __name__ == "__main__":
    main()