"""Fast-dDrive single-shot inference demo.

Loads the Fast-dDrive HF release (via ``trust_remote_code=True``) and decodes
one image+prompt using one of three paths:

* ``section_diffusion``  — iterative MDM denoising over the scaffold (SD).
* ``scaffold_spec``      — scaffold-aware self-speculative decoding (SS, paper canonical).
* ``inference_scaling``  — SS with multi-trajectory rollouts (test-time scaling).
"""

import argparse

import torch
from PIL import Image
from transformers import AutoModelForCausalLM, AutoProcessor, AutoTokenizer


_MODES = {
    "section_diffusion":  ("mdm_sample_deep_scaffold",        0.9),
    "scaffold_spec":      ("scaffold_speculative_sample",     0.0),
    "inference_scaling":  ("scaffold_spec_with_ss_multi_traj", 0.0),
}


def main():
    p = argparse.ArgumentParser()
    p.add_argument("--model_path", default="Efficient-Large-Model/Fast-dDrive",
                   help="Fast-dDrive checkpoint directory or HuggingFace id "
                        "(default: Efficient-Large-Model/Fast-dDrive on the Hugging Face Hub).")
    p.add_argument("--image", required=True, help="Path to a single image.")
    p.add_argument("--prompt", required=True, help="Text prompt.")
    p.add_argument("--mode", default="scaffold_spec", choices=sorted(_MODES),
                   help="Decoding path; scaffold_spec is the paper canonical (SS).")
    p.add_argument("--confidence_threshold", type=float, default=None,
                   help="Override the per-mode default (0.0 for scaffold_spec / "
                        "inference_scaling, 0.9 for section_diffusion).")
    args = p.parse_args()

    method_name, default_threshold = _MODES[args.mode]
    threshold = args.confidence_threshold if args.confidence_threshold is not None else default_threshold

    print(f"Loading model from {args.model_path} ...", flush=True)
    model = AutoModelForCausalLM.from_pretrained(
        args.model_path, torch_dtype=torch.bfloat16,
        device_map="cuda:0", trust_remote_code=True,
    ).eval()

    tokenizer = AutoTokenizer.from_pretrained(args.model_path, trust_remote_code=True)
    processor = AutoProcessor.from_pretrained(args.model_path, use_fast=False)
    processor.tokenizer = tokenizer

    image = Image.open(args.image).convert("RGB")
    messages = [{"role": "user", "content": [
        {"type": "image", "image": image},
        {"type": "text", "text": args.prompt},
    ]}]
    text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    inputs = processor(text=[text], images=[image], return_tensors="pt").to("cuda:0")

    mask_id = int(tokenizer.encode("|<MASK>|", add_special_tokens=False)[0])
    kwargs = dict(
        input_ids=inputs.input_ids,
        tokenizer=tokenizer,
        block_size=32,
        max_tokens=512,
        mask_id=mask_id,
        threshold=threshold,
    )
    if getattr(inputs, "pixel_values", None) is not None:
        kwargs["pixel_values"] = inputs.pixel_values
    if getattr(inputs, "image_grid_thw", None) is not None:
        kwargs["image_grid_thw"] = inputs.image_grid_thw

    print(f"Decoding with {method_name} (threshold={threshold}) ...", flush=True)
    with torch.inference_mode():
        out = getattr(model, method_name)(**kwargs)
    trimmed = out[0, inputs.input_ids.shape[1]:]
    print("\n--- Response ---\n" + tokenizer.decode(trimmed, skip_special_tokens=True))


if __name__ == "__main__":
    main()