#!/usr/bin/env python3
"""
Side-by-side comparison demo: AR Baseline vs Diffusion VLM

Usage:
    python demo_compare.py [--image PATH] [--prompt TEXT] [--max-tokens N]
"""

import argparse
import os
import sys
import time
import threading
from typing import Optional, Set

import torch

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

# Demo image
_demo_images_dir = os.path.join(os.path.dirname(_script_dir), "demo_images")
DEFAULT_IMAGE = os.path.join(_demo_images_dir, "beach_dog.jpeg")
DEFAULT_PROMPT = "Describe this image."


def _pick_latest_checkpoint(run_dir: str) -> Optional[str]:
    """Pick latest valid checkpoint-* directory from a run directory."""
    candidates = []
    for name in os.listdir(run_dir):
        if not name.startswith("checkpoint-"):
            continue
        ckpt_dir = os.path.join(run_dir, name)
        if not os.path.isdir(ckpt_dir):
            continue
        if not os.path.isfile(os.path.join(ckpt_dir, "config.json")):
            continue
        step_str = name[len("checkpoint-"):]
        step = int(step_str) if step_str.isdigit() else -1
        mtime = os.path.getmtime(ckpt_dir)
        candidates.append((step, mtime, ckpt_dir))

    if not candidates:
        return None
    candidates.sort(key=lambda x: (x[0], x[1]))
    return candidates[-1][2]


def _resolve_local_mdm_model_path(model_path: str) -> str:
    """Resolve MDM model path to an existing local checkpoint directory."""
    expanded = os.path.expanduser(model_path)
    repo_root = os.path.abspath(os.path.join(_script_dir, "..", "..", ".."))

    if os.path.isabs(expanded):
        path_candidates = [expanded]
    else:
        path_candidates = [
            os.path.abspath(expanded),
            os.path.abspath(os.path.join(repo_root, expanded)),
            os.path.abspath(os.path.join(_script_dir, expanded)),
        ]

    # De-duplicate while preserving order.
    unique_candidates = []
    seen = set()
    for path in path_candidates:
        if path in seen:
            continue
        unique_candidates.append(path)
        seen.add(path)

    for candidate in unique_candidates:
        if not os.path.isdir(candidate):
            continue
        if os.path.isfile(os.path.join(candidate, "config.json")):
            return candidate

        latest_ckpt = _pick_latest_checkpoint(candidate)
        if latest_ckpt is not None:
            return latest_ckpt

        raise FileNotFoundError(
            "Found local MDM directory but no loadable model config.\n"
            f"Directory: {candidate}\n"
            "Expected either:\n"
            "  1) config.json directly under this directory, or\n"
            "  2) checkpoint-*/config.json under this directory."
        )

    tried = "\n".join(f"  - {p}" for p in unique_candidates)
    raise FileNotFoundError(
        "MDM model must be a local directory, but it was not found.\n"
        f"Input: {model_path}\n"
        f"Tried:\n{tried}\n"
        "Tip: pass a checkpoint path like output_models/<run>/checkpoint-XXXX."
    )


class DualStreamDisplay:
    """Display two streams side by side in terminal."""
    
    def __init__(self, width: int = None):
        import shutil
        self.term_width = width or shutil.get_terminal_size().columns
        self.col_width = (self.term_width - 3) // 2  # -3 for " | " separator
        
        self.left_text = ""
        self.right_text = ""
        self.left_done = False
        self.right_done = False
        self.left_stats = ""
        self.right_stats = ""
        self.lock = threading.Lock()

    @staticmethod
    def _visible_len(text: str) -> int:
        """Return string length excluding ANSI escape sequences."""
        import re
        ansi_escape = re.compile(r'\x1B(?:[@-Z\\-_]|\[[0-?]*[ -/]*[@-~])')
        return len(ansi_escape.sub('', text))

    def _center_ansi(self, text: str, width: int) -> str:
        """Center text by visible width while preserving ANSI styles."""
        visible = self._visible_len(text)
        if visible >= width:
            return text
        pad_total = width - visible
        left_pad = pad_total // 2
        right_pad = pad_total - left_pad
        return (" " * left_pad) + text + (" " * right_pad)
        
    def _wrap_text(self, text: str, width: int) -> list:
        """Wrap text to specified width."""
        # Remove ANSI codes for length calculation
        import re
        ansi_escape = re.compile(r'\x1B(?:[@-Z\\-_]|\[[0-?]*[ -/]*[@-~])')
        
        lines = []
        for paragraph in text.split('\n'):
            if not paragraph:
                lines.append("")
                continue
            # Calculate visible length without ANSI codes
            while len(ansi_escape.sub('', paragraph)) > width:
                split_pos = paragraph.rfind(' ', 0, width + 20)  # Allow some slack for ANSI
                if split_pos == -1:
                    split_pos = width
                lines.append(paragraph[:split_pos])
                paragraph = paragraph[split_pos:].lstrip()
            lines.append(paragraph)
        return lines
    
    def _render(self):
        """Render both columns."""
        import re
        ansi_escape = re.compile(r'\x1B(?:[@-Z\\-_]|\[[0-?]*[ -/]*[@-~])')
        
        left_lines = self._wrap_text(self.left_text, self.col_width)
        right_lines = self._wrap_text(self.right_text, self.col_width)
        
        # Pad to same height (minimum 5 lines)
        max_lines = max(len(left_lines), len(right_lines), 5)
        left_lines.extend([""] * (max_lines - len(left_lines)))
        right_lines.extend([""] * (max_lines - len(right_lines)))
        
        # Clear screen and move to top
        sys.stdout.write("\033[2J\033[H")
        
        # Header
        header_left = "\033[1;97;44m AR Baseline (Qwen2.5-VL) \033[0m"
        header_right = "\033[1;97;45m Diffusion VLM (MDM) \033[0m"
        print("┌" + "─" * self.col_width + "┬" + "─" * self.col_width + "┐")
        print("│" + self._center_ansi(header_left, self.col_width) + "│" + self._center_ansi(header_right, self.col_width) + "│")
        print("├" + "─" * self.col_width + "┼" + "─" * self.col_width + "┤")
        
        # Content
        for left, right in zip(left_lines, right_lines):
            # Pad considering ANSI codes
            left_visible_len = len(ansi_escape.sub('', left))
            right_visible_len = len(ansi_escape.sub('', right))
            left_padded = left + " " * (self.col_width - left_visible_len)
            right_padded = right + " " * (self.col_width - right_visible_len)
            print(f"│{left_padded[:self.col_width + len(left) - left_visible_len]}│{right_padded[:self.col_width + len(right) - right_visible_len]}│")
        
        # Stats row
        print("├" + "─" * self.col_width + "┼" + "─" * self.col_width + "┤")
        left_stat = self.left_stats if self.left_done else "generating..."
        right_stat = self.right_stats if self.right_done else "generating..."
        print("│" + f"\033[90m{left_stat}\033[0m".center(self.col_width + 9) + "│" + 
              f"\033[92m{right_stat}\033[0m".center(self.col_width + 9) + "│")
        print("└" + "─" * self.col_width + "┴" + "─" * self.col_width + "┘")
        
        sys.stdout.flush()
    
    def update_left(self, text: str, stats: str = None):
        with self.lock:
            self.left_text = text
            if stats:
                self.left_stats = stats
                self.left_done = True
            self._render()
    
    def update_right(self, text: str, tokens_in_block: int = None, stats: str = None):
        with self.lock:
            # Add block marker for multi-token blocks
            if tokens_in_block and tokens_in_block > 1:
                self.right_text = text + f" \033[92m[+{tokens_in_block}]\033[0m"
            else:
                self.right_text = text
            if stats:
                self.right_stats = stats
                self.right_done = True
            self._render()


def load_ar_model(model_path: str, device: str = "cuda:0"):
    """Load AR baseline model on specified device."""
    from transformers import AutoModelForImageTextToText, AutoModelForCausalLM, AutoProcessor, AutoTokenizer
    import torch

    # Try image-text model first; fallback to causal LM for custom AR checkpoints.
    try:
        model = AutoModelForImageTextToText.from_pretrained(
            model_path,
            torch_dtype="auto",
            trust_remote_code=True,
        ).to(device)
    except Exception:
        model = AutoModelForCausalLM.from_pretrained(
            model_path,
            torch_dtype="auto",
            trust_remote_code=True,
        ).to(device)

    # Processor for custom checkpoints should come from base model, then swap tokenizer if available.
    processor = AutoProcessor.from_pretrained("Qwen/Qwen2.5-VL-3B-Instruct", use_fast=False)
    try:
        tokenizer = AutoTokenizer.from_pretrained(model_path)
        processor.tokenizer = tokenizer
    except Exception:
        pass
    return model, processor, device


def _resolve_stop_token_ids(model, tokenizer) -> Set[int]:
    """Collect EOS/stop ids from generation config and tokenizer."""
    stop_ids: Set[int] = set()
    gen_cfg = getattr(model, "generation_config", None)
    eos_token_id = getattr(gen_cfg, "eos_token_id", None)
    if isinstance(eos_token_id, int):
        stop_ids.add(int(eos_token_id))
    elif isinstance(eos_token_id, (list, tuple, set)):
        stop_ids.update(int(x) for x in eos_token_id if x is not None)

    if getattr(tokenizer, "eos_token_id", None) is not None:
        stop_ids.add(int(tokenizer.eos_token_id))

    # Keep compatibility with existing MDM stop setting.
    stop_ids.add(151645)
    return stop_ids


def load_mdm_model(model_path: str, device: str = "cuda:1"):
    """Load Diffusion VLM (MDM) model on specified device."""
    import types
    import torch
    from transformers import AutoProcessor, AutoTokenizer
    from lmflow.models.qwen2_5_vl_mdm import Fast_dLLM_Qwen2_5_VLForConditionalGeneration
    from generation_functions import Fast_dLLM_QwenForCausalLM

    resolved_model_path = _resolve_local_mdm_model_path(model_path)
    print(f"[MDM] Loading local model from: {resolved_model_path}", file=sys.stderr)

    # Use torch_dtype="auto" like chatbot_mdm_block_causal_speculative.py
    # and .to(device) instead of device_map for consistency
    model = Fast_dLLM_Qwen2_5_VLForConditionalGeneration.from_pretrained(
        resolved_model_path,
        torch_dtype="auto",
        trust_remote_code=True,
        local_files_only=True,
    ).to(device)
    
    # Bind the speculative_block_causal_sample_cache method
    model.speculative_block_causal_sample_cache = types.MethodType(
        Fast_dLLM_QwenForCausalLM.speculative_block_causal_sample_cache, model
    )
    
    processor = AutoProcessor.from_pretrained("Qwen/Qwen2.5-VL-3B-Instruct", use_fast=False)
    try:
        tokenizer = AutoTokenizer.from_pretrained(resolved_model_path, local_files_only=True)
    except Exception as exc:
        print(
            "[MDM] Local tokenizer not found in checkpoint; "
            "falling back to base tokenizer Qwen/Qwen2.5-VL-3B-Instruct.\n"
            f"[MDM] Tokenizer load error: {exc}",
            file=sys.stderr,
        )
        tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-VL-3B-Instruct", use_fast=False)
    processor.tokenizer = tokenizer
    
    return model, processor, device


@torch.no_grad()
def run_ar_generation(model, processor, device: str, image_path: str, prompt: str, max_tokens: int, display: DualStreamDisplay):
    """Run AR baseline generation with streaming."""
    import traceback
    from transformers import TextIteratorStreamer
    from qwen_vl_utils import process_vision_info
    from threading import Thread
    
    try:
        # Prepare inputs
        messages = [{"role": "user", "content": [
            {"type": "image", "image": image_path},
            {"type": "text", "text": 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(device)
        
        # Streamer
        streamer = TextIteratorStreamer(processor.tokenizer, skip_prompt=True, skip_special_tokens=True)
        generation_kwargs = dict(inputs, max_new_tokens=max_tokens, streamer=streamer)
        
        thread = Thread(target=model.generate, kwargs=generation_kwargs)
        
        start_time = time.time()
        thread.start()
        
        generated_text = ""
        token_count = 0
        for new_text in streamer:
            generated_text += new_text
            token_count += 1
            display.update_left(generated_text)
        
        thread.join()
        elapsed = time.time() - start_time
        tps = token_count / elapsed if elapsed > 0 else 0
        display.update_left(generated_text, stats=f"{token_count} tok / {elapsed:.1f}s = {tps:.1f} tok/s")
    
    except Exception as e:
        display.update_left(f"Error: {str(e)}", stats="FAILED")
        traceback.print_exc(file=sys.stderr)


@torch.no_grad()
def run_mdm_generation(model, processor, device: str, image_path: str, prompt: str, max_tokens: int, 
                       block_size: int, display: DualStreamDisplay):
    """Run Diffusion VLM generation with streaming display.
    
    Uses speculative_block_causal_sample_cache with stream_callback for real-time display.
    """
    import traceback
    from qwen_vl_utils import process_vision_info
    
    try:
        # Prepare inputs
        messages = [{"role": "user", "content": [
            {"type": "image", "image": image_path},
            {"type": "text", "text": 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(device)
        
        mask_id = processor.tokenizer.encode("|<MASK>|")[0]
        original_input_length = inputs.input_ids.shape[1]
        
        # Show "processing..." while first forward runs
        display.update_right("Processing image...", tokens_in_block=0)
        
        start_time = time.time()
        
        # Stream callback for real-time display
        def stream_callback(decoded_text, tokens_in_block):
            display.update_right(decoded_text, tokens_in_block=tokens_in_block)
        
        # Build kwargs for speculative_block_causal_sample_cache
        mdm_kwargs = {
            "input_ids": inputs.input_ids,
            "tokenizer": processor.tokenizer,
            "block_size": block_size,
            "max_tokens": max_tokens,
            "mask_id": mask_id,
            "debug": False,
            "stream_callback": stream_callback,  # Use callback for streaming
        }
        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
        
        # Use the proven speculative_block_causal_sample_cache method (no OOM)
        generated_ids = model.speculative_block_causal_sample_cache(**mdm_kwargs)
        
        # Final stats
        gen_length = generated_ids.shape[1] - original_input_length
        elapsed = time.time() - start_time
        tps = gen_length / elapsed if elapsed > 0 else 0
        
        # Decode final result
        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,
        )[0]
        
        display.update_right(output_text, stats=f"{gen_length} tok / {elapsed:.1f}s = {tps:.1f} tok/s")
    
    except Exception as e:
        error_msg = f"Error: {str(e)}"
        display.update_right(error_msg, stats="FAILED")
        traceback.print_exc(file=sys.stderr)


@torch.no_grad()
def run_synced_comparison(ar_model, ar_processor, ar_device, 
                          mdm_model, mdm_processor, mdm_device,
                          image_path: str, prompt: str, max_tokens: int, block_size: int,
                          display: DualStreamDisplay):
    """Run AR and MDM generation in sync, step by step.
    
    Both models process the image first, then alternate generation steps.
    This shows MDM's advantage of generating multiple tokens per step.
    """
    import traceback
    from qwen_vl_utils import process_vision_info
    from transformers.cache_utils import DynamicCache
    
    try:
        # Prepare inputs for both models
        messages = [{"role": "user", "content": [
            {"type": "image", "image": image_path},
            {"type": "text", "text": prompt}
        ]}]
        
        text_ar = ar_processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
        text_mdm = mdm_processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
        image_inputs, video_inputs = process_vision_info(messages)
        
        inputs_ar = ar_processor(text=[text_ar], images=image_inputs, videos=video_inputs, padding=True, return_tensors="pt")
        inputs_ar = inputs_ar.to(ar_device)
        
        inputs_mdm = mdm_processor(text=[text_mdm], images=image_inputs, videos=video_inputs, padding=True, return_tensors="pt")
        inputs_mdm = inputs_mdm.to(mdm_device)
        
        # MDM setup
        mask_id = mdm_processor.tokenizer.encode("|<MASK>|")[0]
        stop_token = 151645
        mdm_original_length = inputs_mdm.input_ids.shape[1]
        mdm_input_ids = inputs_mdm.input_ids
        mdm_model.model.bd_size = block_size
        mdm_done = False
        mdm_stop_reason = "running"
        
        # AR setup
        ar_original_length = inputs_ar.input_ids.shape[1]
        ar_input_ids = inputs_ar.input_ids
        ar_attention_mask = getattr(inputs_ar, "attention_mask", torch.ones_like(ar_input_ids))
        ar_stop_ids = _resolve_stop_token_ids(ar_model, ar_processor.tokenizer)
        ar_done = False
        ar_stop_reason = "running"
        
        display.update_left("Processing image...")
        display.update_right("Processing image...")
        
        start_time = time.time()
        round_idx = 0
        ar_past_kv = None
        mdm_past_kv = None

        def trim_mdm_kv_cache():
            nonlocal mdm_past_kv
            new_past_key_values = []
            for layer_num in range(len(mdm_past_kv)):
                layer_past_key_values = ()
                for kv_idx in range(len(mdm_past_kv[layer_num])):
                    layer_past_key_values += (
                        mdm_past_kv[layer_num][kv_idx][:, :, :mdm_input_ids.shape[1]-1, :],
                    )
                new_past_key_values.append(layer_past_key_values)
            mdm_past_kv = DynamicCache(new_past_key_values)

        def mdm_generate_round() -> int:
            """Generate one MDM speculative round and return accepted token count."""
            nonlocal mdm_input_ids, mdm_past_kv, mdm_done, mdm_stop_reason
            if mdm_done:
                return 0

            prompt_length = mdm_input_ids.shape[1]
            current_gen_len = prompt_length - mdm_original_length
            remaining = max_tokens - current_gen_len
            if remaining <= 0:
                mdm_done = True
                mdm_stop_reason = f"max_tokens({max_tokens})"
                return 0

            x_init = mask_id * torch.ones((1, block_size - 1), device=mdm_device, dtype=torch.long)
            x_t = torch.cat([mdm_input_ids, x_init], dim=1)
            mask_idx = (x_t[:, -block_size:] == mask_id)

            # Speculative forward
            current_ids = x_t[:, -block_size:]
            logits = mdm_model.forward(
                input_ids=current_ids,
                use_cache=True,
                past_key_values=mdm_past_kv,
                update_kv_cache=False,
                eval_bd_size=block_size,
            ).logits

            logits = torch.cat([logits[:, :1, :], logits[:, :-1, :]], dim=1)
            x_1 = logits.argmax(dim=-1)
            probs = torch.softmax(logits, dim=-1)
            x1_p = torch.gather(probs, dim=-1, index=x_1.unsqueeze(-1)).squeeze(-1)
            x1_p = torch.where(mask_idx, x1_p, -torch.inf)
            unmask_idx = (x1_p > 0)

            if unmask_idx.sum() > 0:
                x_t[:, -block_size:][unmask_idx] = x_1[unmask_idx]

            # Verification forward
            current_ids = x_t[:, -block_size:]
            output = mdm_model.forward(
                input_ids=current_ids,
                use_cache=True,
                past_key_values=mdm_past_kv,
                update_kv_cache=True,
                eval_bd_size=block_size,
            )
            mdm_past_kv = output.past_key_values
            ar_block_token = output.logits.argmax(dim=-1)

            # Count accepted tokens
            accepted = 0
            for i in range(ar_block_token.shape[1] - 1):
                if (ar_block_token[:, i] == current_ids[:, i + 1]).all():
                    accepted += 1
                else:
                    break
            accepted += 1
            accepted = min(accepted, remaining)
            mdm_input_ids = torch.cat([mdm_input_ids, ar_block_token[:, :accepted]], dim=1)
            trim_mdm_kv_cache()

            mdm_gen_len = mdm_input_ids.shape[1] - mdm_original_length
            if mdm_gen_len >= max_tokens:
                mdm_done = True
                mdm_stop_reason = f"max_tokens({max_tokens})"

            if stop_token in mdm_input_ids[:, prompt_length:]:
                stop_idx = (mdm_input_ids[:, prompt_length:] == stop_token).nonzero()[0][1]
                if (mdm_input_ids[:, prompt_length:prompt_length + stop_idx] == mask_id).sum() == 0:
                    mdm_done = True
                    mdm_stop_reason = f"eos({stop_token})"

            return accepted

        def prepare_ar_inputs(first_step: bool):
            common_kwargs = {
                "past_key_values": ar_past_kv,
                "attention_mask": ar_attention_mask,
                "use_cache": True,
                "pixel_values": inputs_ar.pixel_values if (first_step and hasattr(inputs_ar, "pixel_values")) else None,
                "image_grid_thw": inputs_ar.image_grid_thw if (first_step and hasattr(inputs_ar, "image_grid_thw")) else None,
            }
            if hasattr(ar_model, "prepare_inputs_for_generation"):
                return ar_model.prepare_inputs_for_generation(ar_input_ids, **common_kwargs)

            fallback = {
                "input_ids": ar_input_ids if first_step else ar_input_ids[:, -1:],
                "past_key_values": ar_past_kv,
                "attention_mask": ar_attention_mask,
                "use_cache": True,
            }
            if first_step and hasattr(inputs_ar, "pixel_values"):
                fallback["pixel_values"] = inputs_ar.pixel_values
            if first_step and hasattr(inputs_ar, "image_grid_thw"):
                fallback["image_grid_thw"] = inputs_ar.image_grid_thw
            return fallback

        def ar_generate_tokens(target_tokens: int) -> int:
            """Generate exactly target_tokens steps for AR unless it stops earlier."""
            nonlocal ar_input_ids, ar_attention_mask, ar_past_kv, ar_done, ar_stop_reason
            generated = 0
            for _ in range(max(0, target_tokens)):
                if ar_done:
                    break

                ar_inputs = prepare_ar_inputs(first_step=False)
                ar_output = ar_model(**ar_inputs, return_dict=True)
                ar_past_kv = ar_output.past_key_values
                ar_next_token = ar_output.logits[:, -1:, :].argmax(dim=-1)
                ar_input_ids = torch.cat([ar_input_ids, ar_next_token], dim=1)
                ar_attention_mask = torch.cat([ar_attention_mask, torch.ones_like(ar_next_token)], dim=1)
                generated += 1

                next_token_id = int(ar_next_token[0, 0].item())
                ar_gen_len = ar_input_ids.shape[1] - ar_original_length
                if next_token_id in ar_stop_ids:
                    ar_done = True
                    ar_stop_reason = f"eos({next_token_id})"
                    break
                if ar_gen_len >= max_tokens:
                    ar_done = True
                    ar_stop_reason = f"max_tokens({max_tokens})"
                    break
            return generated
        
        # === First shared step: both models do prefill + first token ===
        ar_inputs = prepare_ar_inputs(first_step=True)
        ar_output = ar_model(**ar_inputs, return_dict=True)
        ar_past_kv = ar_output.past_key_values
        ar_next_token = ar_output.logits[:, -1:, :].argmax(dim=-1)
        ar_input_ids = torch.cat([ar_input_ids, ar_next_token], dim=1)
        ar_attention_mask = torch.cat([ar_attention_mask, torch.ones_like(ar_next_token)], dim=1)
        round_idx += 1

        ar_next_token_id = int(ar_next_token[0, 0].item())
        if ar_next_token_id in ar_stop_ids:
            ar_done = True
            ar_stop_reason = f"eos({ar_next_token_id})"
        elif (ar_input_ids.shape[1] - ar_original_length) >= max_tokens:
            ar_done = True
            ar_stop_reason = f"max_tokens({max_tokens})"

        ar_text = ar_processor.tokenizer.decode(ar_input_ids[0, ar_original_length:], skip_special_tokens=True)
        display.update_left(ar_text)
        
        # MDM first forward
        mdm_output = mdm_model.forward(
            input_ids=mdm_input_ids,
            pixel_values=inputs_mdm.pixel_values if hasattr(inputs_mdm, "pixel_values") else None,
            image_grid_thw=inputs_mdm.image_grid_thw if hasattr(inputs_mdm, "image_grid_thw") else None,
            use_cache=True,
            update_kv_cache=True,
        )
        mdm_past_kv = mdm_output.past_key_values
        mdm_next_token = mdm_output.logits[:, -1:, :].argmax(dim=-1)
        mdm_input_ids = torch.cat([mdm_input_ids, mdm_next_token], dim=1)
        mdm_text = mdm_processor.tokenizer.decode(mdm_input_ids[0, mdm_original_length:], skip_special_tokens=True)
        mdm_next_token_id = int(mdm_next_token[0, 0].item())
        if mdm_next_token_id == stop_token:
            mdm_done = True
            mdm_stop_reason = f"eos({stop_token})"
        elif (mdm_input_ids.shape[1] - mdm_original_length) >= max_tokens:
            mdm_done = True
            mdm_stop_reason = f"max_tokens({max_tokens})"
        display.update_right(mdm_text, tokens_in_block=1)
        
        # === Main generation loop (strictly synced by round) ===
        while True:
            if ar_done or mdm_done:
                break

            accepted = mdm_generate_round()
            ar_generate_tokens(max(1, accepted))
            round_idx += 1

            ar_text = ar_processor.tokenizer.decode(ar_input_ids[0, ar_original_length:], skip_special_tokens=True)
            mdm_text = mdm_processor.tokenizer.decode(mdm_input_ids[0, mdm_original_length:], skip_special_tokens=True)
            display.update_left(ar_text)
            display.update_right(mdm_text, tokens_in_block=max(1, accepted))

            # In synced mode, stop once one side finishes to avoid desynced tails.
            if ar_done != mdm_done:
                break
        
        # Final stats
        elapsed = time.time() - start_time
        ar_gen_len = ar_input_ids.shape[1] - ar_original_length
        mdm_gen_len = mdm_input_ids.shape[1] - mdm_original_length
        
        ar_text = ar_processor.tokenizer.decode(ar_input_ids[0, ar_original_length:], skip_special_tokens=True)
        mdm_text = mdm_processor.tokenizer.decode(mdm_input_ids[0, mdm_original_length:], skip_special_tokens=True)
        ar_tps = ar_gen_len / elapsed if elapsed > 0 else 0
        mdm_tps = mdm_gen_len / elapsed if elapsed > 0 else 0
        desync_note = " | stopped_on_first_finish" if ar_done != mdm_done else ""
        
        display.update_left(
            ar_text,
            stats=f"{ar_gen_len} tok / {elapsed:.1f}s = {ar_tps:.1f} tok/s | rounds={round_idx} | stop={ar_stop_reason}{desync_note}",
        )
        display.update_right(
            mdm_text,
            stats=f"{mdm_gen_len} tok / {elapsed:.1f}s = {mdm_tps:.1f} tok/s | rounds={round_idx} | stop={mdm_stop_reason}{desync_note}",
        )
        
    except Exception as e:
        traceback.print_exc(file=sys.stderr)
        display.update_left(f"Error: {str(e)}", stats="FAILED")
        display.update_right(f"Error: {str(e)}", stats="FAILED")


def main():
    parser = argparse.ArgumentParser(description="Compare AR vs Diffusion VLM")
    parser.add_argument("--image", default=DEFAULT_IMAGE, help="Image path")
    parser.add_argument("--prompt", default=DEFAULT_PROMPT, help="Prompt text")
    parser.add_argument("--max-tokens", type=int, default=128, help="Max tokens to generate")
    parser.add_argument("--block-size", type=int, default=16, help="Block size for Diffusion VLM decoding")
    parser.add_argument("--ar-model", default="Qwen/Qwen2.5-VL-3B-Instruct", help="AR model path or HF model id")
    parser.add_argument("--mdm-model", default="output_models/sharegpt4v_8node_bsz256_anneal_bd16_block_causal_lr_1e-6_unfreeze_vision",
                        help="MDM model path")
    parser.add_argument("--ar-device", default="cuda:0", help="Device for AR model")
    parser.add_argument("--mdm-device", default="cuda:1", help="Device for MDM model")
    parser.add_argument("--sequential", action="store_true", help="Run sequentially instead of parallel")
    parser.add_argument("--synced", action="store_true", help="Run step-by-step synced comparison (default)")
    parser.add_argument("--mdm-only", action="store_true", help="Only run MDM model (skip AR baseline)")
    parser.add_argument("--ar-only", action="store_true", help="Only run AR model (skip MDM)")
    args = parser.parse_args()
    
    # Load models based on mode
    ar_model, ar_processor, ar_device = None, None, None
    mdm_model, mdm_processor, mdm_device = None, None, None
    
    print("Loading models...")
    
    if not args.mdm_only:
        print(f"  - AR Baseline ({args.ar_model}) on {args.ar_device}")
        ar_model, ar_processor, ar_device = load_ar_model(args.ar_model, args.ar_device)
    
    if not args.ar_only:
        print(f"  - Diffusion VLM (MDM) ({args.mdm_model}) on {args.mdm_device}")
        mdm_model, mdm_processor, mdm_device = load_mdm_model(args.mdm_model, args.mdm_device)
    
    print("\nStarting comparison demo...")
    print(f"  Image: {args.image}")
    print(f"  Prompt: {args.prompt}")
    print(f"  Max tokens: {args.max_tokens}")
    print(f"  Block size: {args.block_size}")
    print("\nPress Enter to start...")
    input()
    
    # Create display
    display = DualStreamDisplay()
    
    if args.mdm_only:
        # Only MDM
        run_mdm_generation(mdm_model, mdm_processor, mdm_device, args.image, args.prompt, args.max_tokens, args.block_size, display)
    elif args.ar_only:
        # Only AR
        run_ar_generation(ar_model, ar_processor, ar_device, args.image, args.prompt, args.max_tokens, display)
    elif args.sequential:
        # Run sequentially (one after another)
        run_mdm_generation(mdm_model, mdm_processor, mdm_device, args.image, args.prompt, args.max_tokens, args.block_size, display)
        run_ar_generation(ar_model, ar_processor, ar_device, args.image, args.prompt, args.max_tokens, display)
    else:  # Default to synced mode
        # Run step-by-step synced comparison
        run_synced_comparison(
            ar_model, ar_processor, ar_device,
            mdm_model, mdm_processor, mdm_device,
            args.image, args.prompt, args.max_tokens, args.block_size,
            display
        )
    
    print("\n\nDemo complete!")


if __name__ == "__main__":
    main()
