#!/usr/bin/env python
"""Phase 2b pass B: choose lambda per layer, then fit the rank-32 low-rank branch.

Both stages follow `deepcompressor`'s recipe rather than the paper's prose, because the two
differ in ways that change the result.

**The lambda grid is 39 candidates, not 20.** `calib/config/smooth.py:152-161` with the recipe's
`alpha: 0.5, beta: -2, num_grids: 20`:

    choices = [i/20 for i in range(1, 20)]                       # 0.05 .. 0.95
    beta == -2  ->  [(0,0)] + [(a, 0) for a] + [(a, 1-a) for a]  # 1 + 19 + 19

`num_grids` sets the granularity of alpha; `beta: -2` unions two families. The scale itself is
`calib/smooth.py:117` -- `lambda = alpha_base^alpha / beta_base^beta`, with `alpha_base` the
activation AbsMax span and `beta_base` the weight AbsMax span, both per input channel. So the
family with `beta = 0` is purely activation-driven and the family with `beta = 1-alpha` is the
SmoothQuant form; the search picks between them per layer instead of assuming one.

**The low-rank branch is not a truncated SVD.** `configs/svdquant/__default__.yaml` sets
`low_rank: {rank: 32, num_iters: 100, objective: OutputsError, degree: 2, early_stop: true}`. The
SVD gives the initial factors; the iterations then alternate between refitting the low-rank
branch to the *quantization residual* and re-quantizing what is left. Truncating the SVD once and
stopping -- which is what the paper's Eq. 5 reads like on its own -- leaves the branch fitted to
the wrong target, because what it has to absorb is the part the 4-bit grid cannot represent, not
the part with the largest singular values.

Every candidate and every iteration is scored by output error on the sampled rows:

    err(lambda) = || Q(X diag(lambda)^-1) Q(diag(lambda) W^T) - X W^T ||_2

with the quantizer matching the NVFP4 recipe (`configs/svdquant/nvfp4.yaml`): E2M1 elements,
block-16 groups along the input channels, and an FP8 E4M3 scale per group under a per-tensor
scale. Activations are dynamic (`static: false`), so they are quantized per call, not calibrated
to a fixed range.

Layers are independent, so this shards by layer: each worker holds one layer's weight and sample
rather than the whole model.
"""

from __future__ import annotations

import argparse
import json
import sys
import time
from pathlib import Path

import torch

REPO = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(REPO / "scripts"))
sys.path.insert(0, str(REPO / "src"))

FP4_E2M1_LEVELS = (0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0)
FP4_MAX = 6.0
GROUP = 16  # configs/svdquant/nvfp4.yaml: group_shapes [[-1,-1], [1,16,1,1,1]]


def _e2m1(x: torch.Tensor) -> torch.Tensor:
    """Round to the nearest representable E2M1 magnitude, keeping sign."""
    levels = torch.tensor(FP4_E2M1_LEVELS, dtype=x.dtype, device=x.device)
    mag = x.abs().clamp(max=FP4_MAX)
    idx = torch.bucketize(mag, (levels[1:] + levels[:-1]) / 2)
    return torch.sign(x) * levels[idx]


def quantize_nvfp4(t: torch.Tensor, group: int = GROUP) -> torch.Tensor:
    """NVFP4 round-trip along the last dim: per-tensor scale, then FP8-E4M3 per-group scale."""
    *lead, n = t.shape
    pad = (-n) % group
    if pad:
        t = torch.nn.functional.pad(t, (0, pad))
    blocks = t.reshape(*lead, -1, group).float()

    global_scale = blocks.abs().amax().clamp_min(1e-12) / (FP4_MAX * 448.0)
    scales = blocks.abs().amax(dim=-1, keepdim=True) / FP4_MAX / global_scale
    # The second-level scale is stored as FP8 E4M3, so it must be rounded to that grid or the
    # dequantization here would be more accurate than the hardware's.
    scales = scales.clamp(1e-8, 448.0).to(torch.float8_e4m3fn).float().clamp_min(1e-8)
    eff = (scales * global_scale).clamp_min(1e-12)
    out = (_e2m1(blocks / eff) * eff).reshape(*lead, -1)
    return out[..., :n] if pad else out


def output_error(x: torch.Tensor, w: torch.Tensor, lam: torch.Tensor | None) -> float:
    """Relative L2 error of the quantized layer against the bf16 layer, on sampled rows."""
    ref = x @ w.T
    if lam is None:
        xq, wq = quantize_nvfp4(x), quantize_nvfp4(w)
    else:
        xq = quantize_nvfp4(x / lam)
        wq = quantize_nvfp4(w * lam)
    err = (xq @ wq.T) - ref
    return (err.norm() / ref.norm().clamp_min(1e-12)).item()


def alpha_beta_pairs(num_grids: int = 20) -> list[tuple[float, float]]:
    choices = [i / num_grids for i in range(1, num_grids)]
    return [(0.0, 0.0)] + [(a, 0.0) for a in choices] + [(a, 1.0 - a) for a in choices]


def smooth_scale(x_span: torch.Tensor, w_span: torch.Tensor, alpha: float, beta: float) -> torch.Tensor:
    if alpha <= 0 and beta <= 0:
        return torch.ones_like(x_span)
    scale = x_span.clamp_min(1e-8).pow(alpha) if alpha > 0 else torch.ones_like(x_span)
    if beta > 0:
        scale = scale / w_span.clamp_min(1e-8).pow(beta)
    return scale.clamp_min(1e-8)


def fit_low_rank(x: torch.Tensor, w: torch.Tensor, rank: int, num_iters: int,
                 early_stop: bool = True) -> tuple[torch.Tensor, torch.Tensor, float, int]:
    """Alternate: fit rank-r to the current residual, re-quantize what is left, repeat.

    Returns `(l1, l2, best_error, best_iter)` with `w ~= l1 @ l2 + Q(w - l1 @ l2)`.
    """
    ref = x @ w.T
    ref_norm = ref.norm().clamp_min(1e-12)
    best = (None, None, float("inf"), -1)
    target = w
    for it in range(num_iters):
        u, s, vh = torch.linalg.svd(target.float(), full_matrices=False)
        l1 = u[:, :rank] * s[:rank]
        l2 = vh[:rank]
        residual = w - (l1 @ l2)
        err = ((quantize_nvfp4(x) @ quantize_nvfp4(residual).T + x @ (l1 @ l2).T) - ref)
        rel = (err.norm() / ref_norm).item()
        if rel < best[2] - 1e-6:
            best = (l1.clone(), l2.clone(), rel, it)
        elif early_stop and it - best[3] >= 5:
            break
        # Next iteration fits the low-rank branch to what quantization could not represent.
        target = w - quantize_nvfp4(w - (l1 @ l2))
    return best


def main() -> int:
    ap = argparse.ArgumentParser(description="Lambda grid search + rank-32 low-rank fit")
    ap.add_argument("--stats", required=True)
    ap.add_argument("--samples", required=True, help="directory of sample_*.pt shards")
    ap.add_argument("--out", required=True)
    ap.add_argument("--model-path", default=None)
    ap.add_argument("--shard", default="0/1", help="i/n over LAYERS")
    ap.add_argument("--rank", type=int, default=32)
    ap.add_argument("--num-iters", type=int, default=100)
    ap.add_argument("--num-grids", type=int, default=20)
    ap.add_argument("--max-rows", type=int, default=8192)
    ap.add_argument("--span", default="all", choices=["all", "video", "text", "audio"],
                    help="which modality's absmax defines lambda. `all` reproduces the first\n                         calibration; `video` excludes the 1.36%% of rows that are text/audio\n                         but set 79.8%% of channel spans in block 0. Needs a stats file from\n                         calib_stats_modal.py for anything but `all`.")
    ap.add_argument("--no-smooth", action="store_true",
                    help="skip the lambda search entirely and fit the low-rank branch to the raw\n                         weight. Not the same as passing an identity lambda to the existing\n                         factors: L1/L2 were fitted to W*lambda, so they have to be refitted to\n                         W or the two halves describe different matrices.")
    ap.add_argument("--eval-span", default="video", choices=["all", "video"],
                    help="which rows the OutputsError sample is drawn from")
    ap.add_argument("--device", default="cuda:0")
    args = ap.parse_args()

    import bench
    from h3opt.svdquant_rules import target_linears

    device = torch.device(args.device)
    stats = torch.load(args.stats, map_location="cpu", weights_only=False)["stats"]
    probe = stats[next(iter(stats))]
    per_modality = "video" in probe and isinstance(probe.get("video"), dict)
    if args.span != "all" and not args.no_smooth and not per_modality:
        raise SystemExit(f"--span {args.span} needs per-modality stats (calib_stats_modal.py)")

    def span_of(name: str) -> torch.Tensor:
        """The absmax that defines lambda for this layer."""
        e = stats[name]
        if not per_modality:
            return e["absmax"]
        if args.span == "all":
            return e["all"]["absmax"]
        v = e[args.span]["absmax"]
        # The token refiner never sees video or audio rows, so a video span there is all zeros;
        # fall back to its text span rather than emitting a degenerate lambda.
        return v if float(v.max()) > 0 else e["text"]["absmax"]
    shard_i, shard_n = (int(x) for x in args.shard.split("/"))

    sample_files = sorted(Path(args.samples).glob("sample_*.pt"))
    names = sorted(stats)[shard_i::shard_n]
    print(f"shard {shard_i}/{shard_n}: {len(names)} layers, {len(sample_files)} sample shards",
          flush=True)

    # Only the weights are needed here; load on CPU and move one layer at a time.
    from diffusers import ModularPipeline

    pipe = ModularPipeline.from_pretrained(args.model_path or bench.DEFAULT_MODEL)
    pipe.load_components(names=["transformer"], dtype=torch.bfloat16)
    # `load_components` may place the denoiser on the accelerator. Only one layer's weight is
    # needed on the GPU at a time here, and leaving the other 62 GB resident is what turns a
    # comfortable per-layer working set into an OOM two layers in.
    pipe.transformer.to("cpu")
    weights = {n: m.weight.detach().to("cpu") for n, m in target_linears(pipe.transformer).items()}
    del pipe
    torch.cuda.empty_cache()
    print(f"weights on cpu: {sum(w.numel() for w in weights.values())/1e9:.2f} B params", flush=True)

    samples: dict[str, list[torch.Tensor]] = {n: [] for n in names}
    for path in sample_files:
        blob = torch.load(path, map_location="cpu", weights_only=False)
        for n in names:
            if n in blob and blob[n].numel():
                samples[n].append(blob[n])

    pairs = alpha_beta_pairs(args.num_grids)
    results = {}
    t0 = time.perf_counter()
    for k, name in enumerate(names):
        x = torch.cat(samples[name])[: args.max_rows].to(device, torch.float32)
        w = weights[name].to(device, torch.float32)
        w_span = w.abs().amax(dim=0).clamp_min(1e-8)
        x_span = span_of(name).to(device).clamp_min(1e-8)

        baseline = output_error(x, w, None)
        if args.no_smooth:
            alpha, beta = 0.0, 0.0
            lam = torch.ones_like(x_span)
            err_s = baseline
        else:
            best = (float("inf"), None, None)
            for alpha, beta in pairs:
                lam = smooth_scale(x_span, w_span, alpha, beta)
                err = output_error(x, w, lam)
                if err < best[0]:
                    best = (err, alpha, beta)
            err_s, alpha, beta = best
            lam = smooth_scale(x_span, w_span, alpha, beta)

        l1, l2, err_lr, it = fit_low_rank(x / lam, w * lam, args.rank, args.num_iters)
        lam_cpu, l1_cpu, l2_cpu = lam.half().cpu(), l1.half().cpu(), l2.half().cpu()
        del l1, l2, lam, x, w, w_span, x_span
        torch.cuda.empty_cache()
        results[name] = {
            "span": args.span,
            "alpha": alpha, "beta": beta, "lambda": lam_cpu,
            "l1": l1_cpu, "l2": l2_cpu,
            "err_no_smooth": baseline, "err_smooth": err_s, "err_smooth_lowrank": err_lr,
            "lowrank_iter": it, "rows": args.max_rows,
        }
        print(f"[{k+1}/{len(names)}] {name}  a={alpha:.2f} b={beta:.2f}  "
              f"err {baseline:.4f} -> {err_s:.4f} -> {err_lr:.4f} (iter {it})", flush=True)

    out = Path(args.out)
    out.parent.mkdir(parents=True, exist_ok=True)
    torch.save({"rank": args.rank, "num_grids": args.num_grids, "shard": args.shard,
                "wall_s": round(time.perf_counter() - t0, 1), "layers": results}, out)
    print(f"wrote {out} in {time.perf_counter()-t0:.0f}s", flush=True)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())