from typing import OrderedDict
from tqdm import tqdm
import src.custom_utils as custom_utils
import argparse
import torch.distributed as dist
import torch
import pandas as pd
from functools import partial
from pathlib import Path
import src.rag.prompts as prompts
from PIL import Image
from transformers import (
    AutoProcessor,
    Qwen2_5_VLForConditionalGeneration
)

logger = custom_utils.get_logger(__name__)


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('--model_checkpoint', type=str)

    parser.add_argument('--ranking_path', type=str)
    parser.add_argument('--top_k', type=int)
    parser.add_argument('--tokenizer_buffer_length', type=int)

    parser.add_argument('--dataset_name', type=str,
                        choices=['evqa', 'infoseek'], required=True)
    parser.add_argument('--query_path', type=str)
    parser.add_argument('--nrows_query', type=int)
    parser.add_argument('--image_root', type=str, required=True)
    parser.add_argument('--infoseek_few_shots', action='store_true')

    parser.add_argument('--dataloader_num_workers', type=int, default=0)

    parser.add_argument('--output_path', type=str, required=True)

    parser.add_argument('--generation_num_beams', type=int, default=1)
    parser.add_argument('--temperature', type=float, default=1.0)

    args = parser.parse_args()
    return args


class RAGDataset(torch.utils.data.Dataset):
    def __init__(
        self,
        args,
        processor  # Changed from tokenizer and image_processor
    ):
        super().__init__()
        self.args = args
        self.processor = processor

        self.data = pd.read_json(
            args.query_path, lines=True, nrows=args.nrows_query)
        image_root = Path(args.image_root)
        self.data.image_path = self.data.image_path.apply(
            lambda x: str(image_root.joinpath(x)))
        logger.info(f"Loaded {len(self.data)} queries from {args.query_path}")

        if args.ranking_path:
            self.ranking = pd.read_json(args.ranking_path, lines=True)
            self.ranking = {qid: group for qid,
                            group in self.ranking.groupby('qid')}
            if args.top_k:
                self.ranking = {qid: group.head(
                    args.top_k) for qid, group in self.ranking.items()}
            logger.info(
                f"Loaded {len(self.ranking)} ranking entries from {args.ranking_path}")
        else:
            logger.info(
                'No ranking provided, generating without external knowledge')
            self.ranking = None

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

    def __getitem__(self, idx):
        sample = self.data.iloc[idx]

        question = sample.question
        if self.ranking is not None:
            context = self.ranking.get(sample.data_id, pd.DataFrame())
            if context.empty:
                logger.warning(
                    f"No context found for data_id {sample.data_id}. Using empty context.")
            context = context.passage_text.str.cat(sep=prompts.CONTEXT_SEP)
            if self.args.tokenizer_buffer_length:
                # Truncate context if needed
                tokenizer = self.processor.tokenizer
                context = ''.join(
                    tokenizer.batch_decode(
                        tokenizer(context, truncation=True,
                                  max_length=tokenizer.model_max_length - self.args.tokenizer_buffer_length
                                  )
                        .input_ids,
                        skip_special_tokens=True
                    )
                )
                if not context:
                    logger.warning(
                        f"Context for data_id {sample.data_id} is empty after truncation. Using empty context.")
            question = prompts.CONTEXT_VQA_PROMPT.format(
                question=question,
                context=context,
                answer=''
            )
        else:
            question = prompts.VQA_PROMPT.format(question=question, answer='')

        if self.args.infoseek_few_shots:
            fs_prompt = '\n\n\n'.join(
                prompts.CONTEXT_VQA_PROMPT.format(
                    context=prompts.CONTEXT_SEP.join(shot['context']),
                    question=shot['question'],
                    answer=shot['answer']
                ) for shot in prompts.INFOSEEK_FEW_SHOTS
            )
            question = fs_prompt + '\n\n\n' + question

        image = Image.open(sample.image_path).convert('RGB')

        messages = [
            {
                'role': 'user',
                'content': [
                    # Just indicate image position, don't pass the image here
                    {'type': 'image'},
                    {'type': 'text', 'text': question}
                ]
            }
        ]

        extra_data = {}
        if args.dataset_name == 'evqa':
            extra_data['question'] = sample.question
            extra_data['reference'] = sample.answer
            extra_data['question_type'] = sample.question_type

        return dict(
            data_id=sample.data_id,
            image=image,
            messages=messages,  # Return messages instead of prompt
            **extra_data
        )


def collate_fn(batch, processor=None, dtype=None):
    ret = {}

    # Collect data_ids and other metadata
    for key in batch[0]:
        if key == 'data_id':
            ret['data_ids'] = [item[key] for item in batch]
        elif key in ['image', 'messages']:
            continue  # Handle these specially below
        else:
            ret[key] = [item[key] for item in batch]

    # Process images and messages together using the processor
    images = [item['image'] for item in batch]
    messages_list = [item['messages'] for item in batch]

    # Apply chat template and process
    texts = []
    for messages in messages_list:
        text = processor.apply_chat_template(
            messages, tokenize=False, add_generation_prompt=True)
        texts.append(text)

    # Process images and text together
    inputs = processor(
        text=texts,
        images=images,
        return_tensors='pt',
        padding=True
    )

    # Convert to appropriate dtype
    if dtype is not None:
        if 'pixel_values' in inputs:
            inputs['pixel_values'] = inputs['pixel_values'].to(dtype)

    ret['inputs'] = inputs
    return ret


if __name__ == '__main__':
    args = parse_args()

    if torch.cuda.is_available():
        dist.init_process_group(backend='nccl')
        torch.cuda.set_device(custom_utils.get_local_rank())
        logger.info(
            f"Running on CUDA device {torch.cuda.current_device()}. Local rank: {custom_utils.get_local_rank()}. Global rank: {dist.get_rank() if dist.is_initialized() else 0}. World size: {dist.get_world_size()}")
    else:
        logger.warning("CUDA is not available. Running on CPU")

    device = custom_utils.get_local_rank()
    rank = dist.get_rank() if dist.is_initialized() else 0
    args.device = device
    dtype = torch.float16
    args.dtype = dtype

    args.batch_size = 1

    # Load model and processor
    model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
        args.model_checkpoint,
        torch_dtype=dtype,
        device_map=device
    )

    processor = AutoProcessor.from_pretrained(args.model_checkpoint)

    if processor.tokenizer.pad_token is None:
        processor.tokenizer.pad_token = processor.tokenizer.eos_token
        model.config.pad_token_id = processor.tokenizer.eos_token_id

    if args.dataset_name == 'infoseek':
        processor.image_processor.max_pixels = 336 ** 2
        logger.info(
            f"Set max_pixels to {processor.image_processor.max_pixels} for infoseek dataset")

    model.requires_grad_(False)
    model.eval()

    dataset = RAGDataset(
        args=args,
        processor=processor
    )

    dataloader = torch.utils.data.DataLoader(
        dataset,
        batch_size=args.batch_size,
        shuffle=False,
        num_workers=args.dataloader_num_workers,
        drop_last=False,
        sampler=torch.utils.data.DistributedSampler(
            dataset, shuffle=False, drop_last=False),
        collate_fn=partial(
            collate_fn,
            processor=processor,
            dtype=dtype,
        )
    )

    local_results = OrderedDict()
    do_sample = args.generation_num_beams > 1
    for i, batch in tqdm(enumerate(dataloader), total=len(dataloader), desc=f"[Rank {rank}] Generating responses"):
        inputs = batch['inputs']
        inputs = {k: v.to(device=device) if isinstance(v, torch.Tensor) else v
                  for k, v in inputs.items()}

        with torch.inference_mode():
            outputs = model.generate(
                **inputs,
                do_sample=do_sample,
                num_beams=args.generation_num_beams,
                temperature=args.temperature,
                max_new_tokens=30,
                use_cache=True
            )

        generated_ids = outputs[:, inputs['input_ids'].shape[-1]:]
        results = processor.batch_decode(
            generated_ids, skip_special_tokens=True)

        if args.dataset_name == 'infoseek':
            local_results.update(zip(batch['data_ids'], results))
        elif args.dataset_name == 'evqa':
            local_results.update(zip(batch['data_ids'], [dict(
                question=batch['question'][i],
                answers=results[i],
                reference=batch['reference'][i],
                question_type=batch['question_type'][i]
            ) for i in range(len(batch['data_ids']))]))

    del model
    torch.cuda.empty_cache()

    logger.info(
        f"[Rank {dist.get_rank()}] Local answers: {len(local_results)}")
    all_results = ([None] * dist.get_world_size()
                   ) if custom_utils.is_main_process() else None
    dist.gather_object(local_results, all_results, dst=0)

    if custom_utils.is_main_process():
        all_results = OrderedDict(
            res for local_res in all_results for res in local_res.items())
        logger.info(
            f"[Rank {dist.get_rank()}] Unique answers: {len(all_results)}")
        output_path = Path(args.output_path)
        output_path.parent.mkdir(parents=True, exist_ok=True, mode=0o775)

        if args.dataset_name == 'infoseek':
            all_results = [dict(data_id=k, prediction=v)
                           for k, v in all_results.items()]
        elif args.dataset_name == 'evqa':
            all_results = [dict(data_id=k, **v)
                           for k, v in all_results.items()]

        all_results = pd.DataFrame(all_results)
        all_results.to_json(output_path, orient='records', lines=True)
        logger.info(f"{len(all_results)} answers saved to {output_path}")
