# Prompt: 弄成input parameters
import argparse
import os
import sys
import types
from typing import Any, Dict, List, Optional, Tuple

# Add models directory to path for imports
_script_dir = os.path.dirname(os.path.abspath(__file__))
_models_dir = os.path.join(os.path.dirname(_script_dir), "models")
_demo_images_dir = os.path.join(os.path.dirname(_script_dir), "demo_images")
if _models_dir not in sys.path:
    sys.path.insert(0, _models_dir)

from transformers import AutoTokenizer, AutoProcessor, AutoModel

from generation_functions import Fast_dLLM_QwenForCausalLM
from qwen_vl_utils import process_vision_info


# Demo samples: (image_path, prompt)
DEMO_SAMPLES: List[Tuple[str, str]] = [
    (
        None,
        """Who are you?"""
    ),
    (
        None,
        """Question: Natalia sold clips to 48 of her friends in April, and then she sold half as many clips in May. How many clips did Natalia sell altogether in April and May? Please think step by step before answering.\nAnswer:"""
    ),
    (
        None,
        """Question: A deep-sea monster rises from the waters once every hundred years to feast on a ship and sate its hunger. Over three hundred years, it has consumed 847 people. Ships have been built larger over time, so each new ship has twice as many people as the last ship. How many people were on the ship the monster ate in the first hundred years? Please think step by step before answering.\nAnswer:"""
    ),
    (
        "chengyue_script/dvlm/demo_images/mmmu_pro.png",
        "Answer the following multiple-choice question in the image. The last line of your response should be of the following format: 'Answer: $LETTER' (without quotes) where LETTER is one of the options. Think step by step before answering."
    ),
    (
        "chengyue_script/dvlm/demo_images/chart_qa.png",
        "What's the no.1 country/region in the chart?\nAnswer the question using a single word or phrase."
    ),
    (
        os.path.join(_demo_images_dir, "beach_dog.jpeg"),
        "Describe this image."
    ),
    (
        os.path.join(_demo_images_dir, "mathvista_physics.png"),
        """When a spring does work on an object, we cannot find the work by simply multiplying the spring force by the object's displacement. The reason is that there is no one value for the force-it changes. However, we can split the displacement up into an infinite number of tiny parts and then approximate the force in each as being constant. Integration sums the work done in all those parts. Here we use the generic result of the integration.

In Figure, a cumin canister of mass $m=0.40 \\mathrm{~kg}$ slides across a horizontal frictionless counter with speed $v=0.50 \\mathrm{~m} / \\mathrm{s}$. It then runs into and compresses a spring of spring constant $k=750 \\mathrm{~N} / \\mathrm{m}$. When the canister is momentarily stopped by the spring, by what distance $d$ is the spring compressed?"""
    ),
    (
        os.path.join(_demo_images_dir, "mathverse.jpg"),
        """As shown in the figure, in triangle ABC, it is known that angle A = 80.0, angle B = 60.0, point D is on AB and point E is on AC, DE parallel BC, then the size of angle CED is ()
Choices:
A:40°
B:60°
C:120°
D:140°
"""
    ),
    (
        os.path.join(_demo_images_dir, "mathvision.jpg"),
        """Which number should be written in place of the question mark? Please think step by step before answering."""
    ),
]


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description="Run Qwen-VL chatbot generation with configurable inputs.")
    parser.add_argument(
        "--model-name",
        default="output_models/j63_math_8node_bsz256_anneal_bd32_block_causal_lr_1e-6_unfreeze_vision/checkpoint-7400",
        help="Model identifier or local path passed to AutoModel/AutoTokenizer.",
    )
    parser.add_argument(
        "--processor-name",
        default="Qwen/Qwen2.5-VL-3B-Instruct",
        help="Processor identifier passed to AutoProcessor.",
    )
    parser.add_argument(
        "--image",
        default=None,
        help="Image URL or path used in the user message. Leave empty to use demo samples.",
    )
    parser.add_argument(
        "--prompt",
        default=None,
        help="Text prompt included in the user message. Leave empty to use demo samples.",
    )
    parser.add_argument(
        "--demo",
        action="store_true",
        help="Run demo samples instead of custom input.",
    )
    parser.add_argument(
        "--block-size",
        type=int,
        default=32,
        help="Block size passed to model.mdm_sample.",
    )
    parser.add_argument(
        "--max-tokens",
        type=int,
        default=1024,
        help="Maximum number of tokens to generate.",
    )
    parser.add_argument(
        "--mask-token",
        default="|<MASK>|",
        help="Mask token string used to compute mask_id.",
    )
    parser.add_argument(
        "--device",
        default="cuda",
        help="Device to move processor outputs to (e.g., 'cuda' or 'cpu').",
    )
    return parser


def build_messages(image: Optional[str], prompt: str) -> List[Dict[str, Any]]:
    content: List[Dict[str, Any]] = []
    if image:
        content.append({"type": "image", "image": image})
    content.append({"type": "text", "text": prompt})
    return [{"role": "user", "content": content}]


def run_inference(model, processor, image: Optional[str], prompt: str, args) -> str:
    """Run inference on a single sample and return the output text."""
    messages = build_messages(image, prompt)

    text = processor.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True,
    )
    image_inputs, video_inputs = process_vision_info(messages)
    inputs = processor(
        text=[text],
        images=image_inputs,
        videos=video_inputs,
        padding=True,
        return_tensors="pt",
    )
    inputs = inputs.to(args.device)
    block_size = args.block_size
    mask_id = processor.tokenizer.encode(args.mask_token)[0]
    mdm_kwargs = {
        "input_ids": inputs.input_ids,
        "tokenizer": processor.tokenizer,
        "block_size": block_size,
        "max_tokens": args.max_tokens,
        "mask_id": mask_id,
        "debug": False,
        "stream": True,  # Enable real-time token streaming
        # "output_file": "decoded_output_speculative.txt",
    }
    if hasattr(inputs, "pixel_values"):
        mdm_kwargs["pixel_values"] = inputs.pixel_values
    if hasattr(inputs, "image_grid_thw"):
        mdm_kwargs["image_grid_thw"] = inputs.image_grid_thw

    generated_ids = model.speculative_block_causal_sample_cache(**mdm_kwargs)
    generated_ids_trimmed = [
        out_ids[len(in_ids) :] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
    ]
    output_text = processor.batch_decode(
        generated_ids_trimmed,
        skip_special_tokens=True,
        clean_up_tokenization_spaces=False,
    )
    return output_text[0]


def main() -> None:
    args = build_parser().parse_args()

    print("=" * 80)
    print("Loading model...")
    print("=" * 80)
    
    from lmflow.models.qwen2_5_vl_mdm import Fast_dLLM_Qwen2_5_VLForConditionalGeneration
    model = Fast_dLLM_Qwen2_5_VLForConditionalGeneration.from_pretrained(
        args.model_name,
        torch_dtype="auto",
        trust_remote_code=True,
    ).to(args.device)

    model.speculative_block_causal_sample_cache = types.MethodType(Fast_dLLM_QwenForCausalLM.speculative_block_causal_sample_cache, model)

    processor = AutoProcessor.from_pretrained(args.processor_name, use_fast=False)
    tokenizer = AutoTokenizer.from_pretrained(args.model_name)
    processor.tokenizer = tokenizer

    # Determine which samples to run
    if args.demo or (args.image is None and args.prompt is None):
        # Run demo samples
        samples = DEMO_SAMPLES
        print(f"\nRunning {len(samples)} demo samples...\n")
    else:
        # Run custom input
        image = args.image if args.image else "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen-VL/assets/demo.jpeg"
        prompt = args.prompt if args.prompt else "Describe this image."
        samples = [(image, prompt)]

    # Run inference on each sample
    for i, (image, prompt) in enumerate(samples):
        print("=" * 80)
        print(f"Sample {i + 1}/{len(samples)}")
        print("=" * 80)
        print(f"Image: {image}")
        print(f"Prompt: {prompt[:100]}..." if len(prompt) > 100 else f"Prompt: {prompt}")
        print("-" * 80, flush=True)
        
        output = run_inference(model, processor, image, prompt, args)
        
        print(f"Output:\n{output}")
        print()

    print("=" * 80)
    print("Done!")
    print("=" * 80)


if __name__ == "__main__":
    main()
