import sys
import os
script_dir = os.path.dirname(os.path.abspath(__file__))
project_root = os.path.dirname(script_dir)
sys.path.append(project_root)

from qwen_src.qwen2_5_vl import Qwen2_5_VLForConditionalGeneration
from qwen_src.qwen2_vl import Qwen2VLForConditionalGeneration
try:
    from qwen_src.qwen3_vl.modeling_qwen3_vl_batch import Qwen3VLForConditionalGeneration
except Exception:
    Qwen3VLForConditionalGeneration = None
import torch
import matplotlib.pyplot as plt
from qwen_vl_utils import process_vision_info
from transformers import AutoProcessor
import math
from PIL import Image, ImageDraw, ImageFont
import numpy as np


import json
import os
from tqdm import tqdm

import pickle
import argparse
import multiprocessing as mp
import pandas as pd  # <-- Make sure to import pandas
from matching_utils import TextVQAAccuracyEvaluator
from check_gate_label import plot_single_sample


def infer_model_type(model_path, model_type):
    if model_type and model_type != "auto":
        return model_type
    lower = model_path.lower()
    if "qwen3" in lower:
        return "qwen3"
    if "qwen2.5" in lower or "qwen2_5" in lower or "qwen2-5" in lower:
        return "qwen2_5"
    if "qwen2" in lower:
        return "qwen2"
    return "qwen2_5"


def resolve_processor_path(model_path, model_type, processor_path=None):
    if processor_path:
        return processor_path
    lower = model_path.lower()
    public_prefix = "/root/public/Qwen"

    def to_repo_path(name):
        if "/root/public" in model_path:
            return f"{public_prefix}/{name}"
        return f"Qwen/{name}"

    if model_type == "qwen3":
        if "8b" in lower:
            return to_repo_path("Qwen3-VL-8B-Instruct")
        if "4b" in lower:
            return to_repo_path("Qwen3-VL-4B-Instruct")
    if model_type == "qwen2_5":
        if "7b" in lower:
            return to_repo_path("Qwen2.5-VL-7B-Instruct")
        if "3b" in lower:
            return to_repo_path("Qwen2.5-VL-3B-Instruct")
    if model_type == "qwen2" and "7b" in lower:
        return to_repo_path("Qwen2-VL-7B-Instruct")

    return model_path


def resolve_model_class(model_type):
    if model_type == "qwen3":
        if Qwen3VLForConditionalGeneration is None:
            raise RuntimeError("Qwen3VLForConditionalGeneration not available.")
        return Qwen3VLForConditionalGeneration
    if model_type == "qwen2":
        return Qwen2VLForConditionalGeneration
    return Qwen2_5_VLForConditionalGeneration


def set_twig_flags(model, model_type, two_stage_flag):
    if model_type == "qwen3":
        if hasattr(model, "enable_twig"):
            model.enable_twig = two_stage_flag
        if hasattr(model, "roi_enable2stage"):
            model.roi_enable2stage = two_stage_flag
        if hasattr(model, "language_model"):
            if hasattr(model.language_model, "enable_twig"):
                model.language_model.enable_twig = two_stage_flag
            if hasattr(model.language_model, "roi_enable2stage"):
                model.language_model.roi_enable2stage = two_stage_flag
        return

    if hasattr(model, "model"):
        if hasattr(model.model, "enable_twig"):
            model.model.enable_twig = two_stage_flag
        if hasattr(model.model, "roi_enable2stage"):
            model.model.roi_enable2stage = two_stage_flag
        if hasattr(model.model, "two_stage_vanilla"):
            model.model.two_stage_vanilla = two_stage_flag
    if hasattr(model, "roi_enable2stage"):
        model.roi_enable2stage = two_stage_flag
def get_KL_divergence(logits_list,):
    import torch.nn as nn
    from torch.nn import functional as F
    KL_div = nn.KLDivLoss(reduction='batchmean')
    KL_div_list = [KL_div(F.log_softmax(logits_list[-1], dim=1), F.softmax(logits_list[i], dim=1)).item() for i in range(len(logits_list))]
    return KL_div_list


from torch.utils.data import Dataset, DataLoader


def worker_wrapper(chunk, gpu_id, args, return_dict):
    """
    Wraps the main worker function to handle return values via shared dictionary.
    """

    result = worker_process(chunk, gpu_id, args)
    # Store result in the shared dictionary using gpu_id as key
    return_dict[gpu_id] = result



class QwenEvalDataset(Dataset):
    def __init__(self, questions_chunk, args):
        self.data = questions_chunk
        self.args = args
        self.image_folder = args.image_folder
        self.use_short_prompt = getattr(args, 'use_short_prompt', False)


    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        line = self.data[idx]
        image_file_str = ''
        ground_truth = None
        qs_text = ""
        image = None

        try:

            qs_text = line["text"]
            qs_text = qs_text
            image_file = line["image"]
            image_path = os.path.join(self.image_folder, image_file)
            image = Image.open(image_path).convert('RGB')
            image_file_str = image_file
            ground_truth = line.get("answer") if "answer" in line else line.get("answers", None)

            if "<image>" in qs_text:
                qs_text = qs_text.replace("<image>", "")  # Remove it from the text body

            qs = "<image>\n" + qs_text

            if self.use_short_prompt:
                qs = qs + "\nAnswer the question using a single word or phrase."

            # Return a dict containing everything needed for inference
            return {
                'image': image,
                'prompt_text': qs,
                'meta': {
                    'image_id': image_file_str,
                    'prompt': qs_text,
                    'ground_truth': ground_truth,
                    'line_number': line['original_line_number'],
                    'id': line.get('question_id', line.get('original_line_number', 'unknown'))
                },
                'valid': True
            }

        except Exception as e:
            print(f"Error loading item {idx}: {e}")
            # Return dummy invalid item
            return {'valid': False, 'meta': {'line_number': line.get('original_line_number', -1)}}


class SingleResCollator:
    def __init__(self, processor_path, patch_size):
        self.processor_path = processor_path
        self.processor = None
        # Define the requested limits
        self.min_pixels = 128 * patch_size * patch_size
        self.max_pixels = 576 * patch_size * patch_size

    def __call__(self, batch):
        if self.processor is None:
            # Initialize processor once per worker
            self.processor = AutoProcessor.from_pretrained(
                self.processor_path,
                trust_remote_code=True,
                padding_side='left',
                use_fast=True
            )
            # Set the resolution constraints globally for this processor
            self.processor.image_processor.min_pixels = self.min_pixels
            self.processor.image_processor.max_pixels = self.max_pixels

        batch = [item for item in batch if item.get('valid', False)]
        if not batch:
            return None

        # 1. Extract raw data
        raw_images = [item['image'] for item in batch]
        raw_prompts = [item['prompt_text'] for item in batch]
        metas = [item['meta'] for item in batch]

        # 2. Prepare text and images
        current_texts = []
        for prompt_str, img in zip(raw_prompts, raw_images):
            # Qwen-VL-Chat template construction
            messages = [{"role": "user", "content": [{"type": "text", "text": prompt_str}, {"image": img}]}]
            text = self.processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)

            # Specific replacement logic for Qwen2.5-VL
            text = text.replace('<|vision_start|><|image_pad|><|vision_end|>', '') \
                .replace('<image>', '<image><|vision_start|><|image_pad|><|vision_end|>')
            current_texts.append(text)

        try:
            # 3. Process Batch (Single Resolution Setting)
            inputs = self.processor(
                text=current_texts,
                images=raw_images,
                padding=True,
                return_tensors="pt"
            )
            # Attach extra data needed later if necessary (or just return inputs)
            inputs.data['src_images'] = raw_images

            return {
                'metas': metas,
                'inputs': inputs
            }

        except Exception as e:
            print(f"Error processing batch: {e}")
            return None


def worker_process(questions_chunk, gpu_id, args):
    device = f'cuda:{gpu_id}'
    answer_analysis_dict = {}  # Renamed from ms_answer_analysis_dict
    max_new_tokens = 256

    # --- SETUP DATALOADER ---
    batch_size = getattr(args, 'batch_size', 8)
    num_workers = getattr(args, 'num_workers', 4)

    dataset = QwenEvalDataset(questions_chunk, args)

    # Initialize the NEW single resolution collator
    collator = SingleResCollator(args.processor_path, args.patch_size)

    dataloader = DataLoader(
        dataset,
        batch_size=batch_size,
        shuffle=False,
        num_workers=num_workers,
        collate_fn=collator,
        pin_memory=True
    )

    if len(dataset) == 0: return answer_analysis_dict

    try:
        # Load Model
        print(f"Worker {gpu_id}: Loading model...")
        model_class = resolve_model_class(args.model_type)
        model = model_class.from_pretrained(
            args.model_path,
            torch_dtype=torch.bfloat16,
            trust_remote_code=True,
            device_map=device
        )
        model.eval()

        # Flags (keeping your original config)
        two_stage_flag = False
        set_twig_flags(model, args.model_type, two_stage_flag)

        # Tokenizer for decoding
        tokenizer = AutoProcessor.from_pretrained(args.processor_path, trust_remote_code=True).tokenizer

        print(f"Worker {gpu_id}: Starting inference...")

        # --- MAIN INFERENCE LOOP ---
        progress_bar = tqdm(dataloader, desc=f"GPU {gpu_id}", position=gpu_id)

        for batch_data in progress_bar:
            if batch_data is None: continue

            metas = batch_data['metas']
            inputs = batch_data['inputs']  # Now getting direct inputs, not inputs_by_res

            try:
                # Move to GPU
                inputs = inputs.to(device)

                with torch.no_grad():
                    output_ids = model.generate(
                        **inputs,
                        max_new_tokens=max_new_tokens,
                        output_attentions=False,
                        pad_token_id=tokenizer.eos_token_id
                    )

                # Decode
                input_ids_len = inputs.input_ids.shape[1]
                generated_ids = output_ids[:, input_ids_len:]
                batch_answers = tokenizer.batch_decode(
                    generated_ids, skip_special_tokens=True, clean_up_tokenization_spaces=True
                )

                # --- SAVE RESULTS ---
                for idx, (meta, answer) in enumerate(zip(metas, batch_answers)):
                    line_number = meta['line_number']

                    tmp_sample = {
                        'image': meta['image_id'],
                        'prompt': meta['prompt'],
                        'text': meta['ground_truth'],
                        'id': meta['id'],
                        'response': answer  # Storing single response
                    }
                    answer_analysis_dict[line_number] = tmp_sample

                # Cleanup GPU memory for this batch
                del inputs, output_ids, generated_ids

            except RuntimeError as e:
                msg = "ERROR_OOM" if "out of memory" in str(e) else f"ERROR_RUNTIME: {e}"
                print(f"Worker {gpu_id}: {msg}")
                torch.cuda.empty_cache()
                # Log errors for this batch
                for meta in metas:
                    line_number = meta['line_number']
                    answer_analysis_dict[line_number] = {
                        'id': meta['id'],
                        'response': msg
                    }

    except Exception as e:
        print(f"Worker {gpu_id}: Critical Error: {e}")
        import traceback
        traceback.print_exc()
    finally:
        torch.cuda.empty_cache()

    return answer_analysis_dict


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Parallel KL Divergence Evaluation")
    parser.add_argument('--model-path', type=str, required=True,
                        help="Path to the model and processor")
    parser.add_argument('--model-type', type=str, default="auto",
                        choices=["auto", "qwen2_5", "qwen2", "qwen3"],
                        help="Model family for loading (auto detects from model path).")
    parser.add_argument('--processor-path', type=str, default="",
                        help="Override processor path if auto detection is incorrect.")
    parser.add_argument('--question-file', type=str, default=None,
                        help="Path to the JSONL question file (E.g., TextVQA)")
    parser.add_argument('--image-folder', type=str, default=None,
                        help="Path to the image directory (for JSONL file)")
    parser.add_argument('--dst-save-path', type=str, required=True,
                        help="Path to save the output .pkl file")
    parser.add_argument('--gpu-ids', type=str, default='0,1',
                        help="Comma-separated list of GPU IDs to use (e.g., '0,1,2,3')")
    # --- ADDED: New HF Dataset args ---
    parser.add_argument('--hf_dataset', type=str, default=None,
                        help="Name of the Hugging Face dataset (e.g., 'craigwu/vstar_bench')")
    parser.add_argument('--dataset_split', type=str, default='test',
                        help="Split to use from the Hugging Face dataset (e.g., 'test')")
    parser.add_argument('--is_check_pred', type=bool, default=False,
                        help="Whether to check high-res predictions")
    ### --- END ADDITION ---
    parser.add_argument('--use-short-prompt', action='store_true',
                        help="Whether to use short prompt format")
    args = parser.parse_args()
    args.model_type = infer_model_type(args.model_path, args.model_type)
    args.processor_path = resolve_processor_path(
        args.model_path, args.model_type, args.processor_path
    )
    args.patch_size = 32 if args.model_type == "qwen3" else 28


    # --- 1. Define a Wrapper Function ---
    # This helper must be defined at the top level (global scope) so it can be pickled.
    # --- 2. The Main Execution Block ---
    # Ensure start method is spawn (Crucial for PyTorch+CUDA)
    try:
        mp.set_start_method('spawn', force=True)
    except RuntimeError:
        pass

    # 1. Prepare Data and Split into Chunks
    gpu_ids = [int(gid) for gid in args.gpu_ids.split(',')]
    num_gpus = len(gpu_ids)
    questions = []

    if args.hf_dataset:
        print(f"Loading Hugging Face dataset {args.hf_dataset}, split {args.dataset_split}")
        from datasets import load_dataset

        dataset = load_dataset(args.hf_dataset, split=args.dataset_split)
        print("Casting 'image' column to load from cache...")
        questions_list = list(dataset)

        new_questions = []
        for i, tmp_ques in enumerate(questions_list):
            tmp_ques['original_line_number'] = i
            new_questions.append(tmp_ques)
        questions = new_questions
        print(f"Loaded {len(questions)} preprocessed questions from Hugging Face.")

    elif args.question_file:
        print("Loading and preprocessing questions from JSONL file...")
        if not args.image_folder:
            raise ValueError("--image-folder is required when using --question-file")

        questions_list = [json.loads(q) for q in open(os.path.expanduser(args.question_file), "r")]

        new_questions = []
        for i, tmp_ques in enumerate(questions_list):
            qs = tmp_ques['text'] if 'text' in tmp_ques else tmp_ques['prompt']
            if 'Reference OCR' in qs:
                qs = '\n'.join([ele for ele in qs.split('\n') if 'Reference OCR' not in ele])
            tmp_ques['text'] = qs
            tmp_ques['original_line_number'] = i
            new_questions.append(tmp_ques)
        questions = new_questions
        print(f"Loaded {len(questions)} preprocessed questions from local file.")

    else:
        raise ValueError("You must provide either --question-file or --hf_dataset")

    # Create chunks
    chunk_size = math.ceil(len(questions) / num_gpus)
    chunks = [questions[i:i + chunk_size] for i in range(0, len(questions), chunk_size)]

    # 2. Run processes in parallel using mp.Process (Manual Management)
    print(f"Starting parallel processing on {num_gpus} GPUs: {gpu_ids}")

    # Create a Manager to hold the results safely across processes
    manager = mp.Manager()
    return_dict = manager.dict()
    processes = []

    for i in range(num_gpus):
        chunk_data = chunks[i] if i < len(chunks) else []

        # We pass the shared 'return_dict' to the wrapper
        p = mp.Process(
            target=worker_wrapper,
            args=(chunk_data, gpu_ids[i], args, return_dict)
        )
        p.start()
        processes.append(p)

    # Wait for all processes to finish
    for p in processes:
        p.join()

    # 3. Gather and Save Results
    print("All chunks processed. Gathering results...")
    final_dict = {}

    # Extract results from the shared return_dict
    for gpu_id, sub_dict in return_dict.items():
        final_dict.update(sub_dict)

    #check save directory
    os.makedirs(os.path.dirname(os.path.expanduser(args.dst_save_path)), exist_ok=True)
    with open(os.path.expanduser(args.dst_save_path), 'wb') as f:
        pickle.dump(final_dict, f)

    print(f"✅ Successfully processed and saved {len(final_dict)} answers to {args.dst_save_path}.")


'''
CUDA_VISIBLE_DEVICES=2,3 python analysis/making_base_model_response.py \
    --model-path Qwen/Qwen2.5-VL-7B-Instruct \
    --question-file /home/yuheng/datasets/textvqa/converted_llava_style_train.jsonl \
    --image-folder ../llava/playground/data/eval/textvqa/train_images \
    --dst-save-path /home/yuheng/datasets/SFT_SDRPN_PostTraining/textvqa_qwen2.5vl-7B.pkl \
    --gpu-ids 0,1

CUDA_VISIBLE_DEVICES=2,3 python analysis/making_base_model_response.py \
    --model-path Qwen/Qwen2.5-VL-7B-Instruct \
    --question-file /home/yuheng/datasets/vstar_coco_spatial_relation_data.jsonl \
    --image-folder /home/yuheng/datasets \
    --dst-save-path /home/yuheng/datasets/SFT_SDRPN_PostTraining/vstar_coco_spatial_qwen2.5vl-7B.pkl \
    --gpu-ids 0,1




python analysis/making_base_model_response.py \
    --model-path /root/epfs/sproject/code/Qwen2.5-VL/output/qwen2_5vl-7b-roi-K18T3-185k-v1bf16-denoisedv2-twiginit-filled \
    --question-file /root/epfs/sproject/datasets/textvqa/converted_llava_style_train.jsonl \
    --image-folder /root/epfs/sproject/datasets/textvqa/train_images \
    --dst-save-path /root/epfs/sproject/datasets/SFT_SDRPN_PostTraining/textvqa_qwen2.5vl-7B-short-prompt.pkl \
    --use-short-prompt \
    --gpu-ids 0,1
    
python analysis/making_base_model_response.py \
    --model-path /root/epfs/sproject/code/Qwen2.5-VL/output/qwen2_5vl-7b-roi-K18T3-185k-v1bf16-denoisedv2-twiginit-filled \
    --question-file /root/epfs/sproject/datasets/chartqa_28k_qa.jsonl \
    --image-folder /root/epfs/sproject/datasets \
    --dst-save-path /root/epfs/sproject/datasets/SFT_SDRPN_PostTraining/chartqa_qwen2.5vl-7B-short-prompt.pkl \\
    --use-short-prompt \
    --gpu-ids 0,1

python analysis/making_base_model_response.py \
    --model-path /root/epfs/sproject/code/Qwen2.5-VL/output/qwen2_5vl-7b-roi-K18T3-185k-v1bf16-denoisedv2-twiginit-filled \
    --question-file /root/epfs/sproject/datasets/visual_cot_infovqa_subset15k.jsonl \
    --image-folder /root/epfs/sproject/datasets \
    --dst-save-path /root/epfs/sproject/datasets/SFT_SDRPN_PostTraining/visual_cot_infovqa_qwen2.5vl-7B-short-prompt.pkl \
    --use-short-prompt \
    --gpu-ids 0,1
    
python analysis/making_base_model_response.py \
    --model-path /root/epfs/sproject/code/Qwen2.5-VL/output/qwen2_5vl-7b-roi-K18T3-185k-v1bf16-denoisedv2-twiginit-filled \
    --question-file /root/epfs/sproject/datasets/visual_cot_docvqa_subset33k.jsonl \
    --image-folder /root/epfs/sproject/datasets \
    --dst-save-path /root/epfs/sproject/datasets/SFT_SDRPN_PostTraining/visual_cot_docvqa_qwen2.5vl-7B-short-prompt.pkl \
    --use-short-prompt \
    --gpu-ids 0,1
    
python analysis/making_base_model_response.py \
    --model-path /root/epfs/sproject/code/Qwen2.5-VL/output/qwen2_5vl-7b-roi-K18T3-185k-v1bf16-denoisedv2-twiginit-filled \
    --question-file /root/epfs/sproject/datasets/vstar_coco_spatial_relation_data.jsonl \
    --image-folder /root/epfs/sproject/datasets/ \
    --dst-save-path /root/epfs/sproject/datasets/SFT_SDRPN_PostTraining/vstar_coco_spatial_qwen2.5vl-7B-short-prompt.pkl \
    --use-short-prompt \
    --gpu-ids 0,1
'''