from __future__ import annotations

import argparse
import json
from pathlib import Path
import sys

ROOT = Path(__file__).resolve().parents[1]
if str(ROOT) not in sys.path:
    sys.path.insert(0, str(ROOT))

from orbitquant_wan_a2.source_quant import build_packed_from_official_source


def main():
    ap = argparse.ArgumentParser(
        description="Directly quantize the freshly-downloaded official Wan-Animate-2 distilled BF16 transformer into packed OrbitQuant W4."
    )
    ap.add_argument("--transformer", default="/content/Wan2_2_Animate_Repo/transformer")
    ap.add_argument("--output", default="/content/OrbitQuant_Animate2_W4A4_PACKED")
    ap.add_argument("--seed", type=int, default=20260702)
    ap.add_argument("--device", default="auto")
    ap.add_argument("--row-chunk", type=int, default=512)
    ap.add_argument("--max-shard-gib", type=float, default=0.75)
    ap.add_argument("--model-revision", default=None)
    ap.add_argument("--overwrite", action="store_true")
    args = ap.parse_args()

    manifest = build_packed_from_official_source(
        args.transformer,
        args.output,
        seed=args.seed,
        bits=4,
        device=args.device,
        row_chunk=args.row_chunk,
        max_shard_gib=args.max_shard_gib,
        overwrite=args.overwrite,
        model_revision=args.model_revision,
    )
    audits = [v["audit"] for v in manifest["targets"].values()]
    if len(audits) != 480 or not all(
        a["code_exact_fraction"] == 1.0
        and a["scale_exact_fraction"] == 1.0
        and a["fake_bf16_exact_fraction"] == 1.0
        for a in audits
    ):
        raise RuntimeError("direct-source packed W4 audit did not pass 480/480")
    if manifest["passthrough_count"] != 823 or manifest["full_transformer_tensor_count"] != 1303:
        raise RuntimeError("wrong final transformer inventory")

    print("\nDIRECT SOURCE -> PACKED ORBITQUANT W4 COMPLETE")
    print("  source tensors : 1303 / 1303")
    print("  packed W4      : 480 / 480")
    print("  passthrough    : 823 / 823")
    print("  row scales     : BF16")
    print("  packed GiB     : %.3f" % (manifest["packed_bytes"] / 2**30))
    print("  output         :", Path(args.output))


if __name__ == "__main__":
    main()