"""Load + run a TurboQuant-converted DiffusionGemma via mlx-vlm's diffusion sampler.

mlx_vlm.utils.load_model() can't be used directly: it sees config["quantization"]
and runs nn.quantize (affine), which doesn't understand TurboQuant checkpoints.
This loader replicates its model-construction steps, swaps in PolarQuantized
layers (turboquant_mlx.generate._prepare_polar_layers), then hands the model to
mlx-vlm's block-diffusion generate().

Usage:
    python scripts/generate_tq_diffusiongemma.py \
        --model ./diffusiongemma-26B-A4B-it-tq3-g32 \
        --prompt "Write a short paragraph about the ocean." \
        --max-tokens 256 --temperature 0.0
"""

import argparse
import glob
import time
from pathlib import Path

import mlx.core as mx

import turboquant_mlx.compat  # noqa: F401
from turboquant_mlx.config import TurboQuantConfig
from turboquant_mlx.generate import _prepare_polar_layers

from fast_switch_patch import patch_switch_fast

from mlx_vlm import generate
from mlx_vlm.prompt_utils import apply_chat_template
from mlx_vlm.utils import (
    apply_generation_config_defaults,
    get_model_and_args,
    load_config,
    load_processor,
    update_module_configs,
)


def load_tq_vlm(model_path):
    """Build the mlx-vlm model with PolarQuantized layers and load TQ weights."""
    model_path = Path(model_path)
    config = load_config(model_path)

    tq_dict = config.pop("quantization", None)
    if tq_dict is None or tq_dict.get("mode") != "turboquant":
        raise ValueError(f"{model_path} is not a TurboQuant checkpoint")
    tq_config = TurboQuantConfig.from_dict(tq_dict)
    config.pop("quantization_config", None)

    model_class, _ = get_model_and_args(config=config)
    config.setdefault("text_config", config.pop("llm_config", {}))
    config.setdefault("vision_config", {})
    config.setdefault("audio_config", {})
    model_config = model_class.ModelConfig.from_dict(config)
    model_config = update_module_configs(
        model_config, model_class, config,
        ["text", "vision", "perceiver", "projector", "audio"],
    )
    model_config = apply_generation_config_defaults(model_config, config)
    model = model_class.Model(model_config)

    weights = {}
    for wf in sorted(glob.glob(str(model_path / "model*.safetensors"))):
        weights.update(mx.load(wf))

    # Checkpoint was saved from the model tree (mlx format) — no sanitize.
    _prepare_polar_layers(model, weights, tq_config)
    model.load_weights(list(weights.items()), strict=False)
    mx.eval(model.parameters())
    model.model_path = model_path
    model.eval()

    processor = load_processor(model_path)
    return model, processor, config


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--model", required=True)
    ap.add_argument("--prompt", default="Write a short paragraph about the ocean.")
    ap.add_argument("--max-tokens", type=int, default=256)
    ap.add_argument("--temperature", type=float, default=0.0)
    ap.add_argument("--no-fast-switch", action="store_true",
                    help="disable the batched dequant+gather_mm switch path")
    args = ap.parse_args()

    if not args.no_fast_switch:
        patch_switch_fast()

    t0 = time.time()
    model, processor, config = load_tq_vlm(args.model)
    print(f"[load] {time.time() - t0:.1f}s", flush=True)

    formatted = apply_chat_template(processor, config, args.prompt, num_images=0)
    result = generate(
        model, processor, formatted,
        max_tokens=args.max_tokens,
        temperature=args.temperature,
        verbose=True,
    )
    print(f"\npeak memory: {mx.get_peak_memory() / 1024**3:.2f} GB")


if __name__ == "__main__":
    main()