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

import numpy as np
import torch
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.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):
    """Match chair_eval.py: use dscr_lambda for mixing"""
    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)
                # Match chair_eval.py: use dscr_lambda for mixing
                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)
                # Match chair_eval.py: use dscr_lambda for mixing
                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 recompute_first_logits_from_refined_cache(model, input_ids, cache):
    """Copied from object_hallucination_vqa_llava4.py"""
    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 build_prompt(conv_mode: str, use_im_start_end: bool) -> str:
    qs = "Please describe this image in detail."
    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)
    conv.append_message(conv.roles[1], None)
    return conv.get_prompt()


def load_coco_images(ann_file: str, seed: int, num_samples: int) -> List[Dict]:
    with open(ann_file, "r") as f:
        anns = json.load(f)
    images = anns.get("images", [])
    rng = random.Random(seed)
    rng.shuffle(images)
    if num_samples > 0:
        images = images[:num_samples]
    return images


def main() -> None:
    parser = argparse.ArgumentParser("CHAIR (LLaVA 1.5) VCD/OPERA")
    parser.add_argument("--model-path", type=str, default="liuhaotian/llava-v1.5-13b")
    parser.add_argument("--model-base", type=str, default=None)
    parser.add_argument("--data-path", type=str, default="/path/to/data")
    parser.add_argument("--coco-ann-file", type=str, required=True)
    parser.add_argument("--answers-file", type=str, required=True)
    parser.add_argument("--conv-mode", type=str, default="vicuna_v1")
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--num-samples", type=int, default=500)
    parser.add_argument("--max-new-tokens", type=int, default=256)
    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("--use-vcd", action="store_true", default=False)
    parser.add_argument("--noise-step", type=int, default=500)
    parser.add_argument("--cd-alpha", type=float, default=1.0)
    parser.add_argument("--cd-beta", type=float, default=0.1)

    # OPERA
    parser.add_argument("--use-opera", action="store_true", default=False)
    parser.add_argument("--beam", type=int, default=5)
    parser.add_argument("--scale-factor", "--scale_factor", dest="scale_factor", type=float, default=50.0)
    parser.add_argument("--threshold", type=int, default=15)
    parser.add_argument("--num-attn-candidates", "--num_attn_candidates", dest="num_attn_candidates", type=int, default=5)
    parser.add_argument("--penalty-weights", "--penalty_weights", dest="penalty_weights", type=float, default=1.0)

    # DSCR
    parser.add_argument("--use-dscr", action="store_true", default=False)
    parser.add_argument("--depth-folder", type=str, default="/path/to/data")
    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=0.01)
    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)

    args = parser.parse_args()
    set_seed(args.seed)
    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 = next(model.parameters()).device
    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 576
    )

    if args.use_vcd:
        evolve_vcd_sampling()

    # Print configuration for debugging
    print("="*80)
    print("[CHAIR LLaVA 1.5 CONFIG]")
    print(f"  Model: {model_name}")
    print(f"  Method: {'VCD' if args.use_vcd else 'OPERA' if args.use_opera else 'BASELINE'}")
    if args.use_vcd:
        print(f"  VCD: noise_step={args.noise_step} cd_alpha={args.cd_alpha} cd_beta={args.cd_beta}")
    if args.use_opera:
        print(f"  OPERA: beam={args.beam} scale_factor={args.scale_factor} threshold={args.threshold} "
              f"num_attn_candidates={args.num_attn_candidates} penalty_weights={args.penalty_weights}")
    if args.use_dscr:
        dscr_key_value = args.dscr_key_value or (not args.dscr_key_only and not args.dscr_value_only)
        print(f"  DSCR: alpha={args.dscr_alpha} beta={args.dscr_beta} sigma={args.dscr_sigma} "
              f"keep_ratio={args.dscr_keep_ratio} lambda={args.dscr_lambda} "
              f"start_layer={args.dscr_start_layer} end_layer={args.dscr_end_layer} "
              f"key_only={args.dscr_key_only} value_only={args.dscr_value_only} key_value={dscr_key_value}")
    print(f"  Temperature: {args.temperature}, top_p: {args.top_p}, top_k: {args.top_k}")
    print(f"  Max new tokens: {args.max_new_tokens}, Seed: {args.seed}")
    print("="*80)

    images = load_coco_images(args.coco_ann_file, args.seed, args.num_samples)

    out_dir = os.path.dirname(args.answers_file)
    if out_dir:
        os.makedirs(out_dir, exist_ok=True)

    prompt = build_prompt(args.conv_mode, model.config.mm_use_im_start_end)

    with open(args.answers_file, "w") as f:
        for info in tqdm(images):
            image_id = info["id"]
            image_file = info["file_name"]
            img_path = os.path.join(args.data_path, image_file)
            image = Image.open(img_path).convert("RGB")
            image_tensor = image_processor.preprocess(image, return_tensors="pt")["pixel_values"][0]
            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())

            image_bchw = image_tensor.unsqueeze(0).to(device=device, dtype=model.dtype)
            image_bchw_cd = None
            if args.use_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)

            cache_clean = None
            if args.use_dscr:
                # Build D matrix (match chair_eval.py: use addition, dscr_lambda, self_keep)
                import torch.nn.functional as F
                depth_path = os.path.join(args.depth_folder, os.path.splitext(image_file)[0] + ".npy")
                depth_np = np.load(depth_path, allow_pickle=True)
                depth_tensor = torch.tensor(depth_np, dtype=torch.float32).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)
                
                # Match chair_eval.py: use addition instead of multiplication
                depth_diff = 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)
                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))
                
                # Match chair_eval.py: use addition
                gaussian_weight = (gaussian_weight_depth ** args.dscr_alpha) + (gaussian_weight_position ** args.dscr_beta)
                
                # Match chair_eval.py: self_keep handling
                self_keep = float(args.dscr_keep_ratio)
                self_keep = max(0.0, min(1.0, self_keep))
                if self_keep < 1.0:
                    eye = torch.eye(NUM_IMAGE_TOKENS, device=gaussian_weight.device, dtype=gaussian_weight.dtype)
                    gaussian_weight = gaussian_weight * (1.0 - eye) + self_keep * eye
                
                row_sums = gaussian_weight.sum(dim=-1, keepdim=True)
                zero_row_mask = (row_sums < 1e-12)
                if zero_row_mask.any():
                    gaussian_weight[zero_row_mask.squeeze(-1)] = 1.0
                    row_sums = gaussian_weight.sum(dim=-1, keepdim=True)
                
                D = gaussian_weight / (row_sums + 1e-6)
                
                # Forward to get initial cache
                attention_mask_clean = torch.ones_like(input_ids, device=input_ids.device)
                with torch.inference_mode():
                    kv_out = model(
                        input_ids=input_ids,
                        images=image_bchw,
                        attention_mask=attention_mask_clean,
                        use_cache=True,
                        return_dict=True,
                    )
                
                # Apply D to cache with dscr_lambda mixing (match chair_eval.py)
                dscr_key_value = args.dscr_key_value or (
                    not args.dscr_key_only and not args.dscr_value_only
                )
                cache_clean = apply_D_to_cache(
                    cache=kv_out.past_key_values,
                    D_576_576=D,
                    cut_idx=cut_idx,
                    start_layer_idx=int(args.dscr_start_layer),
                    end_layer_idx=int(args.dscr_end_layer),
                    key_only=bool(args.dscr_key_only),
                    value_only=bool(args.dscr_value_only),
                    key_value=bool(dscr_key_value),
                    dscr_lambda=float(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,
                )

            is_baseline = (not args.use_vcd) and (not args.use_opera) and (not args.use_dscr)
            is_dscr_only = args.use_dscr and (not args.use_vcd) and (not args.use_opera)
            
            # DSCR only: match chair_eval.py behavior (do_sample=False, no temperature)
            if is_dscr_only:
                do_sample = False
                gen_kwargs = dict(
                    input_ids=input_ids,
                    images=image_bchw,
                    do_sample=False,
                    max_new_tokens=args.max_new_tokens,
                    use_cache=True,
                    past_key_values=cache_clean,
                )
            else:
                do_sample = True if (is_baseline or args.use_vcd) else False
                gen_kwargs = dict(
                    input_ids=input_ids,
                    images=image_bchw,
                    do_sample=do_sample,
                    temperature=args.temperature,
                    top_p=args.top_p,
                    top_k=args.top_k,
                    max_new_tokens=args.max_new_tokens,
                )
                if cache_clean is not None:
                    gen_kwargs.update(
                        dict(
                            use_cache=True,
                            past_key_values=cache_clean,
                        )
                    )

            if args.use_vcd:
                gen_kwargs.update(
                    dict(
                        images_cd=image_bchw_cd,
                        cd_alpha=args.cd_alpha,
                        cd_beta=args.cd_beta,
                    )
                )

            if args.use_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,
                }
                gen_kwargs.update(
                    dict(
                        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,
                        do_sample=False,
                    )
                )

            with torch.inference_mode():
                output_ids = model.generate(**gen_kwargs)

            decoded = tokenizer.decode(
                output_ids[0, input_ids.shape[1]:], skip_special_tokens=True
            ).strip()

            json.dump({"image_id": image_id, "caption": decoded}, f)
            f.write("\n")


if __name__ == "__main__":
    main()
