# -*- coding: utf-8 -*-
"""Streaming Layer-by-Layer SVDQuant NVFP4 Quantization for Qwen-Image-2.1.

Enables full-model SVDQuant NVFP4 (r32) quantization on 16GB VRAM GPUs
by streaming one transformer block at a time from disk to GPU.

Key architectural features of Qwen-Image-2.1:
- 32 single-stream blocks (dim=4096, heads=32, head_dim=128, mlp_ratio=3).
- 7 linear layers per block:
  - attn.to_q (4096, 4096)
  - attn.to_k (4096, 4096)
  - attn.to_v (4096, 4096)
  - attn.to_out.0 (4096, 4096)
  - img_mlp.proj (12288, 4096)
  - img_mlp.gate_layer (12288, 4096)
  - img_mlp.out (4096, 12288)
- Zero block-level AdaNorm modulations (eliminating cobblestone flutter!).
- Non-block parameters (~270 MB) preserved in native BF16 for lossless embeddings.
"""

import argparse
import gc
import glob
import json
import os
import sys
import time

import safetensors.torch as st
import torch
from tqdm import tqdm

# Setup paths
ROOT_DIR = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", ".."))
DEEPCOMPRESSOR_DIR = os.path.join(ROOT_DIR, "packages", "deepcompressor")
NUNCHAKU_DIR = os.path.join(ROOT_DIR, "packages", "nunchaku")

for p in [DEEPCOMPRESSOR_DIR, NUNCHAKU_DIR, ROOT_DIR]:
    if p not in sys.path:
        sys.path.insert(0, p)

from deepcompressor.data.dtype import QDType
from deepcompressor.quantizer.config.base import QuantizerConfig
from deepcompressor.quantizer.processor import Quantizer

from src.nunchaku.packer import quantize_and_pack_nvfp4_linear


def load_shard_mapping(index_path: str) -> dict[str, str]:
    with open(index_path, "r") as f:
        data = json.load(f)
    return data["weight_map"]


def load_safetensor_weights(keys: list[str], weight_map: dict[str, str], model_dir: str) -> dict[str, torch.Tensor]:
    """Selectively load only requested tensor keys from safetensors shards."""
    tensors = {}
    shards_needed: dict[str, list[str]] = {}
    for k in keys:
        shard_file = weight_map[k]
        shards_needed.setdefault(shard_file, []).append(k)

    for shard_file, shard_keys in shards_needed.items():
        shard_path = os.path.join(model_dir, shard_file)
        with st.safe_open(shard_path, framework="pt", device="cpu") as f:
            for k in shard_keys:
                tensors[k] = f.get_tensor(k)
    return tensors


def compute_compensated_svd_linear(
    weight: torch.Tensor,
    rank: int = 32,
    num_iters: int = 2,
    device: str = "cuda:0",
) -> tuple[torch.Tensor, torch.Tensor]:
    """Compute SVD low-rank branch with alternating FP4 error compensation."""
    w = weight.detach().to(device=device, dtype=torch.bfloat16)
    w_fp = w.float()

    cfg = QuantizerConfig(
        dtype=QDType.sfp4_e2m1_all,
        group_shapes=[[-1, -1], [1, 16, 1, 1, 1]],
        scale_dtypes=[None, QDType.sfp8_e4m3_nan],
    )
    q = Quantizer(config=cfg, develop_dtype=torch.float32)

    # Initial SVD on weights
    u, s, vh = torch.linalg.svd(w_fp, full_matrices=False)
    lu = (u[:, :rank] * s[:rank]).to(torch.bfloat16)
    ld = vh[:rank, :].to(torch.bfloat16)

    # Alternating optimization: SVD fits (W - Q), absorbing quantization error
    for it in range(1, num_iters):
        lw = lu.float() @ ld.float()
        rw = (w_fp - lw).to(torch.bfloat16)
        qw = q.quantize(rw, return_with_dequant=True).data.float()

        target = w_fp - qw
        u, s, vh = torch.linalg.svd(target, full_matrices=False)
        lu = (u[:, :rank] * s[:rank]).to(torch.bfloat16)
        ld = vh[:rank, :].to(torch.bfloat16)

    return ld, lu


def quantize_block(
    block_idx: int,
    weight_map: dict[str, str],
    model_dir: str,
    rank: int = 32,
    num_iters: int = 2,
    device: str = "cuda:0",
    per_channel: bool = True,
) -> dict[str, torch.Tensor]:
    """Quantize all 7 linear layers and retain norms for 1 block."""
    prefix = f"transformer_blocks.{block_idx}"
    block_keys = [k for k in weight_map if k.startswith(f"{prefix}.")]
    block_weights = load_safetensor_weights(block_keys, weight_map, model_dir)

    out_sd: dict[str, torch.Tensor] = {}

    # Linear layers to quantize
    linear_suffixes = [
        "attn.to_q",
        "attn.to_k",
        "attn.to_v",
        "attn.to_out.0",
        "img_mlp.proj",
        "img_mlp.gate_layer",
        "img_mlp.out",
    ]

    for suffix in linear_suffixes:
        w_key = f"{prefix}.{suffix}.weight"
        w = block_weights[w_key].to(device=device, dtype=torch.bfloat16)

        # SVD low-rank branch
        ld, lu = compute_compensated_svd_linear(w, rank=rank, num_iters=num_iters, device=device)

        # Pack into Nunchaku NVFP4 format
        packed = quantize_and_pack_nvfp4_linear(w, lora=(ld, lu), per_channel=per_channel, device=device)

        # Map to final safetensors parameter names
        for p_name, tensor in packed.items():
            out_sd[f"{prefix}.{suffix}.{p_name}"] = tensor.cpu()

        del w, ld, lu, packed
        torch.cuda.empty_cache()

    # Norm layers: preserve in BF16
    for norm_suffix in ["attn.norm_q.weight", "attn.norm_k.weight"]:
        k = f"{prefix}.{norm_suffix}"
        if k in block_weights:
            out_sd[k] = block_weights[k].to(torch.bfloat16).cpu()

    return out_sd


def run_streaming_quantization(
    model_dir: str,
    output_dir: str,
    rank: int = 32,
    num_iters: int = 2,
    start_block: int = 0,
    end_block: int | None = None,
    device: str = "cuda:0",
    per_channel: bool = True,
):
    os.makedirs(output_dir, exist_ok=True)
    index_path = os.path.join(model_dir, "diffusion_pytorch_model.safetensors.index.json")
    config_path = os.path.join(model_dir, "config.json")

    with open(config_path, "r") as f:
        config = json.load(f)

    weight_map = load_shard_mapping(index_path)
    total_blocks = config.get("num_layers", 32)

    b_start = max(0, start_block)
    b_end = min(total_blocks, end_block if end_block is not None else total_blocks)

    print("=" * 70)
    print(f"🧙‍♀️ SVDQuant NVFP4 Quantization for Qwen-Image-2.1")
    print(f"Blocks: {b_start} to {b_end - 1} (Total: {total_blocks})")
    print(f"Rank: {rank} | SVD Iterations: {num_iters} | Target: {device} | Per-Channel: {per_channel}")
    print(f"Output Directory: {output_dir}")
    print("=" * 70)

    shard_sd: dict[str, torch.Tensor] = {}

    # If starting from block 0, also extract non-block weights in native BF16
    if b_start == 0:
        print("Extracting non-block layers (embeddings, modulation, projections) in BF16...")
        non_block_keys = [k for k in weight_map if not k.startswith("transformer_blocks.")]
        non_block_weights = load_safetensor_weights(non_block_keys, weight_map, model_dir)
        for k, v in non_block_weights.items():
            shard_sd[k] = v.to(torch.bfloat16).cpu()
        print(f"  • Extracted {len(non_block_keys)} non-block tensors.")

    t_start = time.time()
    for b_idx in range(b_start, b_end):
        t0 = time.time()
        print(f"⚡ Quantizing Block {b_idx:02d}/{total_blocks - 1}...", end="", flush=True)
        block_sd = quantize_block(
            block_idx=b_idx,
            weight_map=weight_map,
            model_dir=model_dir,
            rank=rank,
            num_iters=num_iters,
            device=device,
            per_channel=per_channel,
        )
        shard_sd.update(block_sd)
        dt = time.time() - t0
        print(f" Done in {dt:.1f}s ({len(block_sd)} tensors)")
        gc.collect()

    shard_path = os.path.join(output_dir, f"svdq-fp4_r{rank}_blocks_{b_start:02d}_{b_end:02d}.safetensors")
    print(f"\nSaving shard to {shard_path}...")
    st.save_file(shard_sd, shard_path)
    total_time = time.time() - t_start
    print(f"✨ Shard complete in {total_time/60:.2f} minutes!")


def merge_shards(output_dir: str, model_dir: str, rank: int = 32):
    print("=" * 70)
    print(f"🧩 Merging SVDQuant NVFP4 Shards into Consolidated Checkpoint")
    print("=" * 70)

    pattern = os.path.join(output_dir, f"svdq-fp4_r{rank}_blocks_*.safetensors")
    shard_files = sorted(glob.glob(pattern))

    if not shard_files:
        raise FileNotFoundError(f"No shard files found matching {pattern}")

    print(f"Found {len(shard_files)} shards to merge:")
    for sf in shard_files:
        print(f"  • {os.path.basename(sf)}")

    merged_sd: dict[str, torch.Tensor] = {}
    for sf in shard_files:
        print(f"Loading {os.path.basename(sf)}...")
        with st.safe_open(sf, framework="pt", device="cpu") as f:
            for k in f.keys():
                merged_sd[k] = f.get_tensor(k)

    # Copy & augment config.json
    src_config = os.path.join(model_dir, "config.json")
    with open(src_config, "r") as f:
        cfg = json.load(f)

    cfg["quantization_config"] = {
        "quant_method": "nunchaku",
        "rank": rank,
        "precision": "nvfp4",
        "weight": {
            "dtype": "nvfp4",
            "group_size": 16,
        },
    }

    dst_config = os.path.join(output_dir, "config.json")
    with open(dst_config, "w") as f:
        json.dump(cfg, f, indent=2)

    final_safetensors = os.path.join(output_dir, f"svdq-fp4_r{rank}-qwen-image-2.1.safetensors")
    print(f"\nWriting consolidated model ({len(merged_sd)} tensors) to {final_safetensors}...")
    st.save_file(merged_sd, final_safetensors)

    file_size_gb = os.path.getsize(final_safetensors) / (1024**3)
    print(f"🎉 Final Checkpoint Size: {file_size_gb:.2f} GB")
    print("Consolidated NVFP4 model forged successfully!")


def main():
    parser = argparse.ArgumentParser(description="Streaming SVDQuant NVFP4 Quantizer for Qwen-Image-2.1")
    parser.add_argument(
        "--model-dir",
        type=str,
        default="/home/olegk/Nikola/models/Qwen/Qwen-Image-2.1/transformer",
        help="Path to source unquantized transformer directory",
    )
    parser.add_argument(
        "--output-dir",
        type=str,
        default="/home/olegk/Nikola/models/nunchaku-qwen-image-2.1",
        help="Path to output directory for quantized models and shards",
    )
    parser.add_argument("--rank", type=int, default=32, help="SVD low-rank dimension (default: 32)")
    parser.add_argument("--num-iters", type=int, default=2, help="Alternating SVD error compensation iterations")
    parser.add_argument("--start-block", type=int, default=0, help="Starting block index")
    parser.add_argument("--end-block", type=int, default=None, help="Ending block index (exclusive)")
    parser.add_argument("--device", type=str, default="cuda:0", help="CUDA device to use (default: cuda:0)")
    parser.add_argument("--merge", action="store_true", help="Merge shards into consolidated model")
    parser.add_argument("--per-channel", action="store_true", default=True, help="Use per-channel macro scales (wcscales)")
    parser.add_argument("--no-per-channel", dest="per_channel", action="store_false", help="Use per-tensor macro scale (wtscale)")

    args = parser.parse_args()

    if args.merge:
        merge_shards(output_dir=args.output_dir, model_dir=args.model_dir, rank=args.rank)
    else:
        run_streaming_quantization(
            model_dir=args.model_dir,
            output_dir=args.output_dir,
            rank=args.rank,
            num_iters=args.num_iters,
            start_block=args.start_block,
            end_block=args.end_block,
            device=args.device,
            per_channel=args.per_channel,
        )


if __name__ == "__main__":
    main()