import argparse
import os
from typing import Any, Dict, List, Optional, Tuple

import torch
from transformers import AutoModel, AutoTokenizer, AutoConfig

_script_dir = os.path.dirname(os.path.abspath(__file__))
_demo_images_dir = os.path.join(
    os.path.dirname(_script_dir), "dvlm", "demo_images"
)

DEMO_SAMPLES: List[Tuple[Optional[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? Please think step by step before answering.""",
    ),
    (
        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 Ministral MDM block-causal demo with configurable inputs."
    )
    parser.add_argument(
        "--model-name",
        default="output_models/j63_mdm_ministral_8node_bd32_block_causal_ds_final_template/checkpoint-100",
        help="Model identifier or local path.",
    )
    parser.add_argument(
        "--image",
        default=None,
        help="Image URL or local path. Leave empty to use demo samples.",
    )
    parser.add_argument(
        "--prompt",
        default=None,
        help="Text prompt. Leave empty to use demo samples.",
    )
    parser.add_argument(
        "--demo",
        action="store_true",
        help="Run all demo samples instead of custom input.",
    )
    parser.add_argument(
        "--max-new-tokens",
        type=int,
        default=512,
        help="Maximum number of new tokens to generate.",
    )
    parser.add_argument(
        "--steps",
        type=int,
        default=512,
        help="Number of diffusion steps.",
    )
    parser.add_argument(
        "--block-length",
        type=int,
        default=32,
        help="Block length for block-causal generation.",
    )
    parser.add_argument(
        "--threshold",
        type=float,
        default=0.9,
        help="Confidence threshold for early stopping.",
    )
    parser.add_argument(
        "--device",
        default="cuda",
        help="Device (e.g. 'cuda' or 'cpu').",
    )
    parser.add_argument(
        "--dlm-paradigm",
        default="bidirectional",
        help="DLM paradigm to set in model config.",
    )
    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,
    tokenizer,
    image: Optional[str],
    prompt: str,
    args,
) -> Tuple[str, int]:
    """Run a single inference and return (output_text, nfe)."""
    messages = build_messages(image, prompt)

    processed = tokenizer.process_messages(messages, add_generation_prompt=True, enable_thinking=False)
    prompt_ids = processed["input_ids"].to(device=args.device)

    decoded_prompt = tokenizer.decode(prompt_ids[0])
    print(f"=== Expanded prompt ===")
    print(decoded_prompt)
    print(f"=== Prompt token count: {prompt_ids.shape[1]} ===")

    pixel_values = None
    image_sizes = None
    if "pixel_values" in processed:
        pixel_values = processed["pixel_values"].to(
            dtype=torch.bfloat16, device=args.device
        )
        image_sizes = torch.tensor(
            processed["image_sizes"], dtype=torch.long, device=args.device
        )
        print(f"Pixel values shape: {pixel_values.shape}")
        print(f"Image sizes (H, W): {processed['image_sizes']}")

    out_ids, nfe = model.generate(
        prompt_ids,
        max_new_tokens=args.max_new_tokens,
        steps=args.steps,
        block_length=args.block_length,
        shift_logits=False,
        causal_context=True,
        threshold=args.threshold,
        pixel_values=pixel_values,
        image_sizes=image_sizes,
        eos_token_id=tokenizer.eos_token_id,
    )

    response = out_ids[:, prompt_ids.shape[1]:]
    if 11 in response:
        eos_idx = torch.where(response == 11)[1][0]
    else:
        eos_idx = response.shape[1]

    output_text = tokenizer.batch_decode(
        response[:, :eos_idx], skip_special_tokens=True
    )[0]
    return output_text, nfe


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

    print("=" * 80)
    print("Loading model...")
    print("=" * 80)

    tokenizer = AutoTokenizer.from_pretrained(args.model_name, trust_remote_code=True)
    config = AutoConfig.from_pretrained(args.model_name, trust_remote_code=True)
    config.dlm_paradigm = args.dlm_paradigm
    model = AutoModel.from_pretrained(
        args.model_name, config=config, trust_remote_code=True
    )
    model = model.to(args.device).to(torch.bfloat16)

    if args.demo or (args.image is None and args.prompt is None):
        samples = DEMO_SAMPLES
        print(f"\nRunning {len(samples)} demo samples...\n")
    else:
        samples = [(args.image, args.prompt or "Describe this image.")]

    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)

        output, nfe = run_inference(model, tokenizer, image, prompt, args)

        print(f"Model: {output}")
        print(f"[Num Function Eval (NFE)={nfe}]")
        print()

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


if __name__ == "__main__":
    main()
