import argparse
import json
import os
import re
import sys
from typing import Dict, List

import torch
import numpy as np
from PIL import Image
from tqdm import tqdm

ROOT_DIR = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
EXPERIMENTS_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.append(ROOT_DIR)
sys.path.append(EXPERIMENTS_DIR)
sys.path.append(os.path.join(EXPERIMENTS_DIR, "halc_utils"))
sys.path.insert(0, os.path.join(EXPERIMENTS_DIR, "transformers-4.31.0", "src"))

from transformers import set_seed

from llava.constants import (
    IMAGE_TOKEN_INDEX,
    DEFAULT_IMAGE_TOKEN,
    DEFAULT_IM_START_TOKEN,
    DEFAULT_IM_END_TOKEN,
)
from llava.conversation import conv_templates, SeparatorStyle
from llava.model.builder import load_pretrained_model
from llava.mm_utils import tokenizer_image_token, get_model_name_from_path
from llava.utils import disable_torch_init
from vcd_utils.vcd_add_noise import add_diffusion_noise
from vcd_utils.vcd_sample import evolve_vcd_sampling

NUM_IMAGE_TOKENS = 576


def apply_D_to_cache(cache, D_576_576, cut_idx, start_layer_idx, end_layer_idx, key_only, value_only, key_value, dscr_lambda=1.0):
    """Apply D matrix to image-token segment of KV cache (keys/values). dscr_lambda: 0=no DSCR, 1=full DSCR."""
    if D_576_576.dim() == 4:
        D = D_576_576.squeeze(0).squeeze(0)
    elif D_576_576.dim() == 2:
        D = D_576_576
    else:
        raise ValueError(f"Unsupported D shape: {tuple(D_576_576.shape)}")

    start_layer_idx = int(start_layer_idx)
    end_layer_idx = int(end_layer_idx)
    dscr_lambda = float(max(0.0, min(1.0, dscr_lambda)))

    out = []
    D_cache = {}

    for layer_idx, (k, v) in enumerate(cache):
        if start_layer_idx <= layer_idx < end_layer_idx:
            if key_only or key_value:
                tag = (k.device, k.dtype)
                Dk = D_cache.get(tag)
                if Dk is None:
                    Dk = D.to(device=k.device, dtype=k.dtype)
                    D_cache[tag] = Dk

                seg = k[:, :, cut_idx:cut_idx + NUM_IMAGE_TOKENS, :]
                seg_refined = torch.einsum("ij,bhjd->bhid", Dk, seg)
                seg_mixed = (1.0 - dscr_lambda) * seg + dscr_lambda * seg_refined
                k = torch.cat(
                    (k[:, :, :cut_idx, :], seg_mixed, k[:, :, cut_idx + NUM_IMAGE_TOKENS:, :]),
                    dim=2
                )

            if value_only or key_value:
                tag = (v.device, v.dtype)
                Dv = D_cache.get(tag)
                if Dv is None:
                    Dv = D.to(device=v.device, dtype=v.dtype)
                    D_cache[tag] = Dv

                seg = v[:, :, cut_idx:cut_idx + NUM_IMAGE_TOKENS, :]
                seg_refined = torch.einsum("ij,bhjd->bhid", Dv, seg)
                seg_mixed = (1.0 - dscr_lambda) * seg + dscr_lambda * seg_refined
                v = torch.cat(
                    (v[:, :, :cut_idx, :], seg_mixed, v[:, :, cut_idx + NUM_IMAGE_TOKENS:, :]),
                    dim=2
                )

        out.append((k, v))

    return tuple(out)


def expand_past_key_values_for_beams(cache, num_beams: int):
    """Expand past_key_values (batch_size=1) to batch_size=num_beams for beam search.
    _expand_inputs_for_generation does not expand tuple-of-tuples caches."""
    if cache is None or num_beams <= 1:
        return cache
    out = []
    for layer_kv in cache:
        k, v = layer_kv[0], layer_kv[1]
        k_exp = k.repeat_interleave(num_beams, dim=0)
        v_exp = v.repeat_interleave(num_beams, dim=0)
        out.append((k_exp, v_exp))
    return tuple(out)


def recompute_first_logits_from_refined_cache(model, input_ids, cache):
    """Recompute logits for the last token using refined prefix cache."""
    total_keys = [kv[0] for kv in cache]
    total_values = [kv[1] for kv in cache]

    prefix_keys = [k[:, :, :-1, :].contiguous() for k in total_keys]
    prefix_values = [v[:, :, :-1, :].contiguous() for v in total_values]
    prefix_cache = tuple((k, v) for k, v in zip(prefix_keys, prefix_values))

    last_token_ids = input_ids[:, -1:]
    with torch.inference_mode():
        outputs_last = model(
            last_token_ids,
            images=None,
            use_cache=True,
            past_key_values=prefix_cache,
            return_dict=True,
        )
    return outputs_last.past_key_values, outputs_last.logits[:, -1, :]


def top_k_top_p_filtering(logits: torch.Tensor, top_k: int = 0, top_p: float = 1.0) -> torch.Tensor:
    """Filter logits using top-k and/or nucleus (top-p) sampling."""
    if top_k > 0:
        v, _ = torch.topk(logits, top_k)
        kth = v[:, -1].unsqueeze(-1)
        logits = torch.where(logits < kth, torch.full_like(logits, -float("inf")), logits)

    if 0.0 < top_p < 1.0:
        sorted_logits, sorted_idx = torch.sort(logits, descending=True, dim=-1)
        probs = torch.softmax(sorted_logits, dim=-1)
        cum = torch.cumsum(probs, dim=-1)

        mask = cum > top_p
        mask[:, 1:] = mask[:, :-1].clone()
        mask[:, 0] = False

        sorted_logits = torch.where(mask, torch.full_like(sorted_logits, -float("inf")), sorted_logits)
        logits = torch.zeros_like(logits).scatter(-1, sorted_idx, sorted_logits)

    return logits


def sample_from_logits(
    logits: torch.Tensor,
    temperature: float,
    top_k: int,
    top_p: float,
    do_sample: bool,
) -> torch.Tensor:
    """Sample next token from logits (greedy or stochastic)."""
    if temperature is None or temperature <= 0:
        do_sample = False
        temperature = 1.0

    if temperature != 1.0:
        logits = logits / temperature

    logits = top_k_top_p_filtering(logits, top_k=top_k, top_p=top_p)

    if do_sample:
        probs = torch.softmax(logits, dim=-1)
        return torch.multinomial(probs, num_samples=1).squeeze(-1)
    else:
        return torch.argmax(logits, dim=-1)


def decode_baseline_with_first_logits(
    model,
    tokenizer,
    input_ids: torch.Tensor,
    cache,
    first_logits: torch.Tensor,
    max_new_tokens: int,
    temperature: float,
    top_p: float,
    top_k: int,
    do_sample: bool,
) -> torch.Tensor:
    """Manual autoregressive decoding starting from given first logits and cache."""
    eos = tokenizer.eos_token_id
    seq = input_ids
    logits = first_logits

    for _ in range(int(max_new_tokens)):
        next_id = sample_from_logits(
            logits,
            temperature=float(temperature),
            top_k=int(top_k) if top_k is not None else 0,
            top_p=float(top_p),
            do_sample=bool(do_sample),
        )
        next_id = next_id.view(1, 1)
        seq = torch.cat([seq, next_id.to(seq.device)], dim=1)
        if eos is not None and int(next_id.item()) == int(eos):
            break

        with torch.inference_mode():
            out = model(
                input_ids=next_id,
                past_key_values=cache,
                use_cache=True,
                return_dict=True,
            )
        cache = out.past_key_values
        logits = out.logits[:, -1, :]

    return seq


def load_questions(path: str) -> List[Dict]:
    with open(path, "r", encoding="utf-8") as f:
        first = f.read(1)
        f.seek(0)
        if first == "[":
            return json.load(f)
        return [json.loads(line) for line in f if line.strip()]


def discover_sets(gt_dir: str) -> List[str]:
    sets = []
    for f in os.listdir(gt_dir):
        if f.endswith(".json") or f.endswith(".jsonl"):
            sets.append(f.rsplit(".", 1)[0])
    return sorted(list(set(sets)))


def pick_img_dir(img_root: str, dataset: str) -> str:
    if os.path.isdir(os.path.join(img_root, dataset)):
        return os.path.join(img_root, dataset)
    if os.path.isdir(os.path.join(img_root, dataset.lower())):
        return os.path.join(img_root, dataset.lower())
    if dataset.lower() == "posters":
        if os.path.isdir(os.path.join(img_root, "posters")):
            return os.path.join(img_root, "posters")
        if os.path.isdir(os.path.join(img_root, "Posters")):
            return os.path.join(img_root, "Posters")
    return img_root


def pick_depth_dir(depth_root: str, dataset: str) -> str:
    return pick_img_dir(depth_root, dataset)


def get_device(model) -> torch.device:
    try:
        return next(model.parameters()).device
    except StopIteration:
        return torch.device("cuda" if torch.cuda.is_available() else "cpu")


def build_prompt(qs: str, conv_mode: str, fmt: str, use_im_start_end: bool) -> str:
    if use_im_start_end:
        qs_mm = DEFAULT_IM_START_TOKEN + DEFAULT_IMAGE_TOKEN + DEFAULT_IM_END_TOKEN + "\n" + qs
    else:
        qs_mm = DEFAULT_IMAGE_TOKEN + "\n" + qs
    conv = conv_templates[conv_mode].copy()
    conv.append_message(conv.roles[0], qs_mm + fmt)
    conv.append_message(conv.roles[1], None)
    return conv.get_prompt(), conv.sep if conv.sep_style != SeparatorStyle.TWO else conv.sep2


def normalize_format(fmt: str) -> str:
    if fmt == "yn_format":
        return " Please only answer yes or no."
    if fmt == "ow_format":
        return " Please answer this question with one word."
    return ""


def normalize_yn_answer(text: str) -> str:
    cleaned = text.strip().lower()
    match = re.search(r"\b(yes|no)\b", cleaned)
    return match.group(1) if match else cleaned


def build_yn_prefix_allowed_tokens_fn(tokenizer, input_len: int):
    candidates = ["yes", "no", "Yes", "No", " yes", " no"]
    allowed = set()
    for text in candidates:
        token_ids = tokenizer.encode(text, add_special_tokens=False)
        if token_ids:
            allowed.add(token_ids[0])
    if not allowed:
        return None

    eos_id = getattr(tokenizer, "eos_token_id", None)
    if eos_id is None:
        eos_id = getattr(tokenizer, "eod_id", None)

    def _fn(batch_id, input_ids):
        cur_len = input_ids.shape[-1] if hasattr(input_ids, "shape") else len(input_ids)
        if cur_len == input_len:
            return list(allowed)
        return [eos_id] if eos_id is not None else list(allowed)

    return _fn


def main() -> None:
    parser = argparse.ArgumentParser("MME (LLaVA 1.6)")
    parser.add_argument("--model-path", type=str, default="liuhaotian/llava-v1.6-vicuna-13b")
    parser.add_argument("--model-base", type=str, default=None)
    parser.add_argument("--gt-dir", type=str, required=True)
    parser.add_argument("--image-root", type=str, required=True)
    parser.add_argument("--out-root", type=str, required=True)
    parser.add_argument("--depth-root", type=str, default=None)
    parser.add_argument("--datasets", nargs="*", default=None)

    parser.add_argument("--method", type=str, default="baseline",
                        choices=["baseline", "vcd", "opera", "halc", "damo", "agla"])
    # Match mme_llava15.py defaults for consistency
    parser.add_argument("--conv-mode", type=str, default="vicuna_v1")
    parser.add_argument("--format", type=str, default="no_format",
                        choices=["no_format", "ow_format", "yn_format"])
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--max-new-tokens", type=int, default=2)
    parser.add_argument("--temperature", type=float, default=1.0)
    parser.add_argument("--top-p", dest="top_p", type=float, default=1.0)
    parser.add_argument("--top-k", dest="top_k", type=int, default=None)

    # VCD
    parser.add_argument("--noise-step", type=int, default=500)
    parser.add_argument("--cd-alpha", type=float, default=0.5)
    parser.add_argument("--cd-beta", type=float, default=0.1)

    # OPERA
    parser.add_argument("--beam", type=int, default=5)
    parser.add_argument("--scale-factor", type=float, default=50.0)
    parser.add_argument("--threshold", type=int, default=15)
    parser.add_argument("--num-attn-candidates", type=int, default=5)
    parser.add_argument("--penalty-weights", type=float, default=0.6)

    # HALC
    parser.add_argument("--halc-detector", type=str, default="dino", choices=["dino", "owlv2"])
    parser.add_argument("--halc-box-threshold", type=float, default=0.45)
    parser.add_argument("--halc-k-candidate-num", type=int, default=4)
    parser.add_argument("--halc-context-window", type=int, default=4)
    parser.add_argument("--halc-expand-ratio", type=float, default=0.6)
    parser.add_argument("--halc-context-domain", type=str, default="upper", choices=["upper", "lower"])
    parser.add_argument("--halc-contrast-weight", type=float, default=0.05)
    parser.add_argument("--halc-score-type", type=str, default="BLIP",
                        choices=["CLIP", "BLIP", "Random", "Perplexity", "HPSv2"])
    parser.add_argument("--halc-debugger", type=int, default=0, choices=[0, 1, 2])
    parser.add_argument("--halc-mature-layer", type=int, default=None)
    parser.add_argument("--halc-base-layer", type=int, default=0)
    parser.add_argument("--halc-candidate-layers", type=int, nargs="*", default=None)
    parser.add_argument("--halc-relative-top", type=float, default=0.1)
    parser.add_argument("--halc-beam-search", action="store_true", default=False)
    parser.add_argument("--halc-num-beams", type=int, default=1)

    # DAMO
    parser.add_argument("--tau", type=float, default=-0.3)
    parser.add_argument("--beta-1", dest="beta_1", type=float, default=0.05)
    parser.add_argument("--beta-2", dest="beta_2", type=float, default=0.20)
    parser.add_argument("--alpha", type=float, default=0.7)
    parser.add_argument("--damo-start-layer", dest="damo_start_layer", type=int, default=20,
                        help="DAMO momentum starts at this layer (default 20 for 13B/40 layers, 50%%)")

    # AGLA
    parser.add_argument("--agla-alpha", type=float, default=2.0)
    parser.add_argument("--agla-beta", type=float, default=0.3)

    # DSCR
    parser.add_argument("--use-dscr", action="store_true", default=False)
    parser.add_argument("--dscr-alpha", type=float, default=0.6)
    parser.add_argument("--dscr-beta", type=float, default=0.8)
    parser.add_argument("--dscr-sigma", type=float, default=0.6)
    parser.add_argument("--dscr-keep-ratio", type=float, default=1.0)
    parser.add_argument("--dscr-lambda", type=float, default=1.0)
    parser.add_argument("--dscr-start-layer", type=int, default=10)
    parser.add_argument("--dscr-end-layer", type=int, default=40)
    parser.add_argument("--dscr-key-only", action="store_true", default=False)
    parser.add_argument("--dscr-value-only", action="store_true", default=False)
    parser.add_argument("--dscr-key-value", action="store_true", default=False)

    parser.add_argument("--run-name", type=str, default=None)
    parser.add_argument("--debug-first-step", action="store_true", default=False,
                        help="Debug: print first step logits comparison for AGLA")

    args = parser.parse_args()
    set_seed(args.seed)

    if args.method in ("halc", "opera"):
        args.format = "yn_format"

    fmt = normalize_format(args.format)
    disable_torch_init()
    model_path = os.path.expanduser(args.model_path)
    model_name = get_model_name_from_path(model_path)
    tokenizer, model, image_processor, _ = load_pretrained_model(
        model_path, args.model_base, model_name
    )
    device = get_device(model)
    vision_tower = model.get_vision_tower() if hasattr(model, "get_vision_tower") else None
    image_token_len = (
        vision_tower.num_patches
        if vision_tower is not None and hasattr(vision_tower, "num_patches")
        else NUM_IMAGE_TOKENS
    )

    if args.method == "vcd":
        evolve_vcd_sampling()
    if args.method == "agla":
        from agla_utils.agla_sample import evolve_agla_sampling
        from agla_utils.augmentation import augmentation
        from lavis.models import load_model_and_preprocess
        from torchvision import transforms

        evolve_agla_sampling()
        blip_device = "cuda" if torch.cuda.is_available() else "cpu"
        blip_model, blip_vis_processors, blip_text_processors = load_model_and_preprocess(
            "blip_image_text_matching", "large", device=blip_device, is_eval=True
        )
        blip_loader = transforms.Compose([transforms.ToTensor()])

    if args.use_dscr and not args.depth_root:
        raise ValueError("--depth-root is required when --use-dscr is set")

    if args.datasets:
        datasets = args.datasets
    else:
        datasets = discover_sets(args.gt_dir)

    run_name = args.run_name or f"llava16_{args.method}_seed{args.seed}"
    os.makedirs(args.out_root, exist_ok=True)

    debug_printed = False

    def log_once(msg: str) -> None:
        nonlocal debug_printed
        if not debug_printed:
            print(msg)
            debug_printed = True

    for dataset in datasets:
        qfile_json = os.path.join(args.gt_dir, f"{dataset}.json")
        qfile_jsonl = os.path.join(args.gt_dir, f"{dataset}.jsonl")
        qfile = qfile_json if os.path.isfile(qfile_json) else qfile_jsonl
        if not os.path.isfile(qfile):
            print(f"[skip] missing questions for {dataset}")
            continue

        questions = load_questions(qfile)
        img_dir = pick_img_dir(args.image_root, dataset)
        depth_dir = pick_depth_dir(args.depth_root, dataset) if args.use_dscr else None
        answers_file = os.path.join(args.out_root, f"{run_name}_{dataset}.jsonl")

        with open(answers_file, "w", encoding="utf-8") as fout:
            for q in tqdm(questions, desc=f"{dataset}"):
                qid = q.get("question_id", q.get("id"))
                image_file = q.get("image", "")
                text = q.get("text", q.get("question", ""))

                prompt, stop_str = build_prompt(
                    text, args.conv_mode, fmt, model.config.mm_use_im_start_end
                )
                input_ids = tokenizer_image_token(
                    prompt, tokenizer, IMAGE_TOKEN_INDEX, return_tensors="pt"
                ).unsqueeze(0).to(device)

                pos = torch.where(input_ids[0] == IMAGE_TOKEN_INDEX)[0]
                if pos.numel() == 0:
                    raise RuntimeError("IMAGE_TOKEN_INDEX not found in input_ids")
                cut_idx = int(pos[0].item())

                img_path = os.path.join(img_dir, image_file)
                image = Image.open(img_path).convert("RGB")
                dscr_depth = None
                if args.use_dscr:
                    depth_path = os.path.join(depth_dir, os.path.splitext(image_file)[0] + ".npy")
                    depth_np = np.load(depth_path, allow_pickle=True)
                    dscr_depth = torch.tensor(depth_np, dtype=torch.float32)
                image_tensor = image_processor.preprocess(image, return_tensors="pt")["pixel_values"][0]
                image_bchw = image_tensor.unsqueeze(0).to(device=device, dtype=model.dtype)

                image_bchw_cd = None
                if args.method == "vcd":
                    image_tensor_cd = add_diffusion_noise(image_tensor, int(args.noise_step))
                    image_bchw_cd = image_tensor_cd.unsqueeze(0).to(device=device, dtype=model.dtype)
                elif args.method == "agla":
                    tensor_image = blip_loader(image.resize((384, 384)))
                    blip_image = blip_vis_processors["eval"](image).unsqueeze(0).to(blip_device)
                    blip_question = blip_text_processors["eval"](text)
                    tokenized_text = blip_model.tokenizer(
                        blip_question, padding="longest", truncation=True, return_tensors="pt"
                    ).to(blip_device)
                    augmented_image = augmentation(
                        blip_image, blip_question, tensor_image, blip_model, tokenized_text, image
                    )
                    image_tensor_cd = image_processor.preprocess(
                        augmented_image, return_tensors="pt"
                    )["pixel_values"][0]
                    image_bchw_cd = image_tensor_cd.unsqueeze(0).to(device=device, dtype=model.dtype)

                dscr_key_value = False
                dscr_end_layer_eff = args.dscr_end_layer
                dscr_key_only_eff = args.dscr_key_only
                dscr_value_only_eff = args.dscr_value_only
                dscr_key_value_eff = False
                if args.use_dscr:
                    dscr_key_value = args.dscr_key_value or (
                        not args.dscr_key_only and not args.dscr_value_only
                    )
                    dscr_end_layer_eff = args.dscr_end_layer
                    dscr_key_only_eff = args.dscr_key_only
                    dscr_value_only_eff = args.dscr_value_only
                    dscr_key_value_eff = dscr_key_value
                    if args.method == "opera":
                        num_hidden_layers = getattr(model.config, "num_hidden_layers", 32)
                        dscr_end_layer_eff = min(args.dscr_end_layer, num_hidden_layers - 1)

                dscr_note = ""
                if args.use_dscr:
                    dscr_note = (
                        f" | DSCR img_start={cut_idx} img_len={image_token_len} "
                        f"a={args.dscr_alpha} b={args.dscr_beta} s={args.dscr_sigma} "
                        f"k={args.dscr_keep_ratio} l={args.dscr_lambda} "
                        f"layers={args.dscr_start_layer}-{dscr_end_layer_eff} "
                        f"key_only={dscr_key_only_eff} value_only={dscr_value_only_eff} "
                        f"key_value={dscr_key_value_eff}"
                    )

                cache_clean = None
                cache_noisy = None
                cache_augmented = None
                first_logits_clean = None
                first_logits_noisy = None
                first_logits_augmented = None
                kv_outputs = None

                if args.method == "vcd" and args.use_dscr and image_bchw_cd is not None:
                    attention_mask_clean = torch.ones_like(input_ids, device=input_ids.device)
                    with torch.inference_mode():
                        out_noisy = model(
                            input_ids=input_ids,
                            images=image_bchw_cd,
                            attention_mask=attention_mask_clean,
                            use_cache=True,
                            return_dict=True,
                        )
                    cache_noisy = out_noisy.past_key_values
                    first_logits_noisy = out_noisy.logits[:, -1, :]

                if args.method == "agla" and args.use_dscr and image_bchw_cd is not None:
                    attention_mask_clean = torch.ones_like(input_ids, device=input_ids.device)
                    with torch.inference_mode():
                        out_augmented = model(
                            input_ids=input_ids,
                            images=image_bchw_cd,
                            attention_mask=attention_mask_clean,
                            use_cache=True,
                            return_dict=True,
                        )
                    cache_augmented = out_augmented.past_key_values
                    first_logits_augmented = out_augmented.logits[:, -1, :]
                
                # Forward to get initial cache (always, like llava4)
                attention_mask_clean = torch.ones_like(input_ids, device=input_ids.device)
                with torch.inference_mode():
                    kv_outputs = model(
                        input_ids=input_ids,
                        images=image_bchw,
                        attention_mask=attention_mask_clean,
                        use_cache=True,
                        return_dict=True,
                    )
                
                if args.use_dscr:
                    # Build D matrix (copied from llava4, with device specification like llava16_2)
                    import torch.nn.functional as F
                    depth_np = dscr_depth.cpu().numpy() if isinstance(dscr_depth, torch.Tensor) else dscr_depth
                    depth_tensor = torch.tensor(depth_np, dtype=torch.float32, device=device).unsqueeze(0).unsqueeze(1)
                    depth_patch = F.interpolate(depth_tensor, size=(24, 24), mode="bilinear", align_corners=False)
                    depth_patch = torch.clamp(1.0e-6 + 1.0 / depth_patch, 0.001, 1000).view(-1)
                    
                    dmin = depth_patch.min()
                    dmax = depth_patch.max()
                    depth_patch = (depth_patch - dmin) / (dmax - dmin + 1e-6)
                    
                    depth_diff = torch.abs(depth_patch.unsqueeze(0) - depth_patch.unsqueeze(1))
                    gaussian_weight_depth = torch.exp(-(depth_diff ** 2) / (2 * (args.dscr_sigma ** 2) + 1e-12))
                    
                    pixel_positions = torch.tensor([[i // 24, i % 24] for i in range(NUM_IMAGE_TOKENS)], 
                                                   dtype=torch.float32, device=device)
                    pixel_positions = pixel_positions / 23.0
                    position_diff = torch.cdist(pixel_positions, pixel_positions, p=2)
                    gaussian_weight_position = torch.exp(-(position_diff ** 2) / (2 * (args.dscr_sigma ** 2) + 1e-12))

                    gaussian_weight = (gaussian_weight_depth ** args.dscr_alpha) + (gaussian_weight_position ** args.dscr_beta)
                    D = gaussian_weight / (gaussian_weight.sum(dim=-1, keepdim=True) + 1e-6)

                    #######
                    # Apply D to cache
                    cache_clean = apply_D_to_cache(
                        cache=kv_outputs.past_key_values,
                        D_576_576=D,
                        cut_idx=cut_idx,
                        start_layer_idx=int(args.dscr_start_layer),
                        end_layer_idx=int(dscr_end_layer_eff),
                        key_only=bool(dscr_key_only_eff),
                        value_only=bool(dscr_value_only_eff),
                        key_value=bool(dscr_key_value_eff),
                        dscr_lambda=args.dscr_lambda,
                    )
                    
                    # Recompute first logits
                    cache_clean, first_logits_clean = recompute_first_logits_from_refined_cache(
                        model=model,
                        input_ids=input_ids,
                        cache=cache_clean,
                    )

                    if args.method == "vcd" and cache_noisy is not None:
                        cache_noisy_dscr = apply_D_to_cache(
                            cache=cache_noisy,
                            D_576_576=D,
                            cut_idx=cut_idx,
                            start_layer_idx=int(args.dscr_start_layer),
                            end_layer_idx=int(dscr_end_layer_eff),
                            key_only=bool(dscr_key_only_eff),
                            value_only=bool(dscr_value_only_eff),
                            key_value=bool(dscr_key_value_eff),
                            dscr_lambda=args.dscr_lambda,
                        )
                        cache_noisy_dscr, _ = recompute_first_logits_from_refined_cache(
                            model=model,
                            input_ids=input_ids,
                            cache=cache_noisy_dscr,
                        )
                        model._dscr_past_key_values_cd = cache_noisy_dscr

                gen_dscr_kwargs = {}
                cache_augmented_for_agla = None
                if cache_clean is not None:
                    gen_dscr_kwargs = {"past_key_values": cache_clean, "use_cache": True}
                    if args.method == "agla" and cache_augmented is not None:
                        cache_augmented_for_agla = cache_augmented
                elif kv_outputs is not None:
                    # Baseline: use kv_outputs.past_key_values like llava4
                    gen_dscr_kwargs = {"past_key_values": kv_outputs.past_key_values, "use_cache": True}

                with torch.inference_mode():
                    if args.method == "opera":
                        key_position = {
                            "image_start": cut_idx,
                            "image_end": cut_idx + image_token_len - 1,
                            "response_start": input_ids.shape[1] + image_token_len - 1,
                        }
                        # OPERA+DSCR: trim cache by 1 to avoid token duplication.
                        # recompute_first_logits_from_refined_cache returns a cache with
                        # seq_len entries (including the last prompt token). If we pass
                        # that directly, prepare_inputs_for_generation truncates input_ids
                        # to the last token and processes it again → duplicate.
                        # By trimming to seq_len-1, the first forward inside OPERA
                        # re-processes only the last token with the DSCR-refined prefix.
                        opera_kwargs = {}
                        if gen_dscr_kwargs.get("past_key_values") is not None:
                            trimmed_cache = tuple(
                                (k[:, :, :-1, :].contiguous(), v[:, :, :-1, :].contiguous())
                                for k, v in gen_dscr_kwargs["past_key_values"]
                            )
                            opera_kwargs = {"past_key_values": trimmed_cache, "use_cache": True}
                        else:
                            opera_kwargs = dict(gen_dscr_kwargs)

                        # OPERA+beam: _expand_inputs_for_generation does not expand
                        # tuple-of-tuples caches, so expand to num_beams here.
                        if opera_kwargs.get("past_key_values") is not None and args.beam > 1:
                            opera_kwargs["past_key_values"] = expand_past_key_values_for_beams(
                                opera_kwargs["past_key_values"], args.beam
                            )

                        # Signal opera_beam_search to drop pre-filled cache from
                        # model_kwargs_ori so that rollback recomputes from images.
                        # Only set when DSCR cache is actually passed.
                        if args.use_dscr and opera_kwargs.get("past_key_values") is not None:
                            model._dscr_opera_active = True

                        # OPERA: keep original settings even with DSCR
                        do_sample_opera = False
                        temp_opera = float(args.temperature)
                        # Debug: log cache info for OPERA+DSCR
                        cache_info = ""
                        if args.use_dscr and cache_clean is not None:
                            cache_info = f" | cache_shape={cache_clean[0][0].shape if len(cache_clean) > 0 else 'N/A'} cache_len={len(cache_clean)}"
                        log_once(
                            f"[OPERA] do_sample={do_sample_opera} temperature={temp_opera} "
                            f"key_position={key_position} beam={args.beam} "
                            f"scale_factor={args.scale_factor} threshold={args.threshold} "
                            f"num_attn_candidates={args.num_attn_candidates} penalty_weights={args.penalty_weights}"
                            f"{dscr_note}{cache_info}"
                        )
                        try:
                            output_ids = model.generate(
                                input_ids,
                                images=image_bchw,
                                do_sample=do_sample_opera,
                                temperature=temp_opera,
                                top_p=float(args.top_p),
                                top_k=args.top_k,
                                max_new_tokens=args.max_new_tokens,
                                key_position=key_position,
                                opera_decoding=True,
                                num_beams=args.beam,
                                output_attentions=True,
                                scale_factor=args.scale_factor,
                                threshold=args.threshold,
                                num_attn_candidates=args.num_attn_candidates,
                                penalty_weights=args.penalty_weights,
                                **opera_kwargs,
                            )
                        finally:
                            # Clean up DSCR flag
                            if hasattr(model, '_dscr_opera_active'):
                                delattr(model, '_dscr_opera_active')
                    elif args.method == "halc":
                        from decoder_zoo.HALC.context_density.halc import halc_assistant

                        halc_params = {
                            "detector": args.halc_detector,
                            "box_threshold": args.halc_box_threshold,
                            "debugger": args.halc_debugger,
                            "k_candidate_num": args.halc_k_candidate_num,
                            "LVLM_backbone": "llava-1.6",
                            "score_type": args.halc_score_type,
                            "context_window": args.halc_context_window,
                            "expand_ratio": args.halc_expand_ratio,
                            "context_domain": args.halc_context_domain,
                            "contrast_weight": args.halc_contrast_weight,
                            "tokenizer": tokenizer,
                        }
                        vis_processor = (
                            lambda img: image_processor.preprocess(img, return_tensors="pt")["pixel_values"][0]
                        )
                        halc = halc_assistant(
                            model=model,
                            vis_processor=vis_processor,
                            device=device,
                            halc_params=halc_params,
                            max_new_tokens=args.max_new_tokens,
                        )
                        # Pass depth and DSCR params to HALC so cropped-region
                        # context evaluation can also apply DSCR (like mme_llava15.py).
                        dscr_params_for_halc = None
                        depth_np_for_halc = None
                        if args.use_dscr and dscr_depth is not None:
                            depth_np_for_halc = dscr_depth.cpu().numpy() if isinstance(dscr_depth, torch.Tensor) else dscr_depth
                            dscr_params_for_halc = {
                                "alpha": args.dscr_alpha,
                                "beta": args.dscr_beta,
                                "sigma": args.dscr_sigma,
                                "keep_ratio": args.dscr_keep_ratio,
                                "lambda": args.dscr_lambda,
                                "start_layer": args.dscr_start_layer,
                                "end_layer": dscr_end_layer_eff,
                                "key_only": dscr_key_only_eff,
                                "value_only": dscr_value_only_eff,
                                "key_value": dscr_key_value_eff,
                            }
                        halc.update_input(
                            img_path, prompt,
                            depth_np=depth_np_for_halc,
                            dscr_params=dscr_params_for_halc,
                        )
                        mature_layer = (
                            args.halc_mature_layer
                            if args.halc_mature_layer is not None
                            else int(getattr(model.config, "num_hidden_layers", 32)) - 1
                        )
                        candidate_layers = args.halc_candidate_layers
                        base_layer = None if candidate_layers else args.halc_base_layer
                        # HALC: keep original settings even with DSCR
                        do_sample_halc = False
                        temp_halc = float(args.temperature)
                        log_once(
                            f"[HALC] do_sample={do_sample_halc} temperature={temp_halc} "
                            f"mature_layer={mature_layer} base_layer={base_layer} "
                            f"candidate_layers={candidate_layers} relative_top={args.halc_relative_top} "
                            f"detector={args.halc_detector} expand_ratio={args.halc_expand_ratio} "
                            f"context_window={args.halc_context_window} context_domain={args.halc_context_domain} "
                            f"contrast_weight={args.halc_contrast_weight} score_type={args.halc_score_type} "
                            f"box_threshold={args.halc_box_threshold} beam_search={args.halc_beam_search} "
                            f"num_beams={args.halc_num_beams}"
                            f"{dscr_note}"
                        )
                        input_len = input_ids.shape[1]
                        yn_prefix_fn = (
                            build_yn_prefix_allowed_tokens_fn(tokenizer, input_len)
                            if args.format == "yn_format"
                            else None
                        )
                        # HALC+DSCR: trim cache by 1 to fix token duplication
                        # (same approach as OPERA+DSCR). Only when cache_clean
                        # is present; non-DSCR HALC keeps the original cache.
                        halc_kwargs = dict(gen_dscr_kwargs)
                        if cache_clean is not None and halc_kwargs.get("past_key_values") is not None:
                            trimmed_cache = tuple(
                                (k[:, :, :-1, :].contiguous(), v[:, :, :-1, :].contiguous())
                                for k, v in halc_kwargs["past_key_values"]
                            )
                            halc_kwargs["past_key_values"] = trimmed_cache

                        output_ids = model.generate(
                            input_ids,
                            images=image_bchw,
                            do_sample=do_sample_halc,
                            temperature=temp_halc,
                            top_p=float(args.top_p),
                            top_k=args.top_k,
                            max_new_tokens=args.max_new_tokens,
                            prefix_allowed_tokens_fn=yn_prefix_fn,
                            halc_decoding=True,
                            dola_decoding=False,
                            beam_search=bool(args.halc_beam_search),
                            num_beams=args.halc_num_beams,
                            halc_assistant=halc,
                            mature_layer=mature_layer,
                            base_layer=base_layer,
                            candidate_premature_layers=candidate_layers,
                            relative_top=args.halc_relative_top,
                            output_attentions=True,
                            output_hidden_states=True,
                            **halc_kwargs,
                        )
                        if isinstance(output_ids, tuple):
                            output_ids = output_ids[0]
                    elif args.method == "damo":
                        # DAMO: keep original settings even with DSCR
                        do_sample_damo = (args.temperature > 0)
                        temp_damo = float(args.temperature)
                        log_once(
                            f"[DAMO] do_sample={do_sample_damo} temperature={temp_damo} "
                            f"tau={args.tau} beta_1={args.beta_1} beta_2={args.beta_2} "
                            f"alpha={args.alpha} damo_start_layer={args.damo_start_layer}{dscr_note}"
                        )
                        output_ids = model.generate(
                            input_ids,
                            images=image_bchw,
                            do_sample=do_sample_damo,
                            temperature=temp_damo,
                            top_p=float(args.top_p),
                            top_k=args.top_k,
                            max_new_tokens=args.max_new_tokens,
                            output_hidden_states=True,
                            tau=args.tau,
                            beta_1=args.beta_1,
                            beta_2=args.beta_2,
                            alpha=args.alpha,
                            damo_start_layer=args.damo_start_layer,
                            **gen_dscr_kwargs,
                        )
                    elif args.method == "agla":
                        # AGLA: keep original settings even with DSCR
                        do_sample_agla = True
                        temp_agla = float(args.temperature)
                        if image_bchw_cd is not None:
                            log_once(
                                f"[AGLA] do_sample={do_sample_agla} temperature={temp_agla} "
                                f"cd_alpha={args.agla_alpha} cd_beta={args.agla_beta} "
                                f"aug_shape={tuple(image_bchw_cd.shape)}{dscr_note}"
                            )
                        if cache_augmented_for_agla is not None:
                            model._agla_cache_augmented = cache_augmented_for_agla
                        # Set debug flag if requested
                        if args.debug_first_step:
                            model._agla_debug_first_step = True
                            model._agla_debug_tokenizer = tokenizer
                        try:
                            output_ids = model.generate(
                                input_ids,
                                images=image_bchw,
                                images_cd=image_bchw_cd,
                                cd_alpha=args.agla_alpha,
                                cd_beta=args.agla_beta,
                                do_sample=do_sample_agla,
                                temperature=temp_agla,
                                top_p=float(args.top_p),
                                top_k=args.top_k,
                                max_new_tokens=args.max_new_tokens,
                                **gen_dscr_kwargs,
                            )
                        finally:
                            # Clean up temporary attributes
                            if hasattr(model, '_agla_cache_augmented'):
                                delattr(model, '_agla_cache_augmented')
                            if hasattr(model, '_agla_debug_first_step'):
                                delattr(model, '_agla_debug_first_step')
                            if hasattr(model, '_agla_debug_tokenizer'):
                                delattr(model, '_agla_debug_tokenizer')
                    elif args.method == "vcd":
                        # VCD: keep original settings even with DSCR
                        do_sample_vcd = True
                        temp_vcd = float(args.temperature)
                        if image_bchw_cd is not None:
                            log_once(
                                f"[VCD] do_sample={do_sample_vcd} temperature={temp_vcd} "
                                f"noise_step={args.noise_step} cd_alpha={args.cd_alpha} "
                                f"cd_beta={args.cd_beta} images_cd_mean={image_bchw_cd.mean().item():.4f} "
                                f"std={image_bchw_cd.std().item():.4f}{dscr_note}"
                            )
                        output_ids = model.generate(
                            input_ids,
                            images=image_bchw,
                            images_cd=image_bchw_cd,
                            cd_alpha=args.cd_alpha,
                            cd_beta=args.cd_beta,
                            do_sample=do_sample_vcd,
                            temperature=temp_vcd,
                            top_p=float(args.top_p),
                            top_k=args.top_k,
                            max_new_tokens=args.max_new_tokens,
                            **gen_dscr_kwargs,
                        )
                    else:
                        # Baseline: use args.temperature. If temperature=0.0, use greedy (do_sample=False)
                        # DSCR: always greedy. Use manual decoding with first_logits_clean like object_hallucination_vqa_llava16_2.py
                        if cache_clean is not None:
                            log_once(f"[BASELINE+DSCR] manual decoding with first_logits_clean do_sample=False temperature=0.0{dscr_note}")
                            output_ids = decode_baseline_with_first_logits(
                                model=model,
                                tokenizer=tokenizer,
                                input_ids=input_ids,
                                cache=cache_clean,
                                first_logits=first_logits_clean,
                                max_new_tokens=args.max_new_tokens,
                                temperature=0.0,
                                top_p=1.0,
                                top_k=0,
                                do_sample=False,
                            )
                        else:
                            # If temperature=0.0, use greedy decoding
                            temp_baseline = float(args.temperature)
                            do_sample_baseline = temp_baseline > 0.0
                            if not do_sample_baseline:
                                temp_baseline = 0.0
                            log_once(f"[BASELINE] do_sample={do_sample_baseline} temperature={temp_baseline}{dscr_note}")
                            output_ids = model.generate(
                                input_ids,
                                images=image_bchw,
                                do_sample=do_sample_baseline,
                                temperature=temp_baseline,
                                top_p=float(args.top_p),
                                top_k=args.top_k,
                                max_new_tokens=args.max_new_tokens,
                                **gen_dscr_kwargs,
                            )

                input_len = input_ids.shape[1]
                decode_ids = output_ids
                if output_ids.shape[1] > input_len:
                    decode_ids = output_ids[:, input_len:]
                output = tokenizer.batch_decode(decode_ids, skip_special_tokens=True)[0].strip()
                if args.format == "yn_format":
                    output = normalize_yn_answer(output)
                if output.endswith(stop_str):
                    output = output[:-len(stop_str)].strip()

                fout.write(
                    json.dumps(
                        {
                            "question_id": qid,
                            "prompt": prompt,
                            "text": output,
                            "model_id": model_name,
                            "image": image_file,
                            "metadata": {},
                        }
                    )
                    + "\n"
                )


if __name__ == "__main__":
    main()
