from typing import OrderedDict
from tqdm import tqdm
from llava_more.src.llava.constants import DEFAULT_IMAGE_TOKEN, IMAGE_TOKEN_INDEX
import src.custom_utils as custom_utils
from llava_more.src.llava.model.builder import load_pretrained_model
import argparse
import torch.distributed as dist
import torch
import pandas as pd
from functools import partial
from pathlib import Path
from llava_more.src.llava.conversation import conv_templates
import src.rag.prompts as prompts
from PIL import Image

logger = custom_utils.get_logger(__name__)


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('--model_checkpoint', type=str, required=True)
    parser.add_argument('--model_name', type=str, required=True)
    parser.add_argument('--conv_template', type=str, default='llama_3_1', required=True)

    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, required=True)
    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('--batch_size', type=int, default=1)

    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


def expand2square(pil_img, background_color):
    width, height = pil_img.size
    if width == height:
        return pil_img
    elif width > height:
        result = Image.new(pil_img.mode, (width, width), background_color)
        result.paste(pil_img, (0, (width - height) // 2))
        return result
    else:
        result = Image.new(pil_img.mode, (height, height), background_color)
        result.paste(pil_img, ((height - width) // 2, 0))
        return result
    
    
class RAGDataset(torch.utils.data.Dataset):
    def __init__(
        self,
        args,
        tokenizer,
        image_processor
    ):
        super().__init__()
        self.args = args
        self.tokenizer = tokenizer

        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}")

        conv = conv_templates[args.conv_template].copy()
        if args.conv_template == 'llama_3_1':
            conv.tokenizer = tokenizer
        self.conv = conv

        self.img_mean = None
        if args.pad_image:
            self.img_mean = tuple(int(x*255)
                                  for x in image_processor.image_mean)

        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 = f"{DEFAULT_IMAGE_TOKEN}\n{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:
                context = ''.join(
                    self.tokenizer.batch_decode(
                        self.tokenizer(context, truncation=True, 
                            max_length=self.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

        conv = self.conv.copy()
        conv.append_message(role='user', message=question)
        prompt = conv.get_prompt()

        if self.args.conv_template == 'llama_3_1' and prompt.startswith('<|begin_of_text|>'):
            # bos token will be added by the tokenizer
            prompt = prompt[len('<|begin_of_text|>'):]

        image = Image.open(sample.image_path).convert('RGB')
        if self.args.pad_image:
            image = expand2square(image, self.img_mean)

        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,
            prompt=prompt,
            **extra_data
        )


def collate_fn(batch, tokenizer=None, image_processor=None, dtype=None, image_token_id=None):
    ret = {}
    for key in batch[0]:
        if key == 'data_id':
            ret['data_ids'] = [item[key] for item in batch]
        elif key == 'image':
            ret['pixel_values'] = image_processor(
                images=[item[key] for item in batch], return_tensors='pt').to(dtype).pixel_values
        elif key == 'prompt':
            ret.update(tokenizer([sample[key] for sample in batch],
                       return_tensors='pt', padding=True, truncation=False))
            if image_token_id is not None:
                ret['input_ids'][ret['input_ids'] ==
                                 image_token_id] = IMAGE_TOKEN_INDEX
        else:
            ret[key] = [item[key] for item in batch]
    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

    tokenizer, model, image_processor, context_len = load_pretrained_model(
        model_path=args.model_checkpoint,
        model_base=None,
        model_name=args.model_name,
        device_map='auto',
        device=device,
        torch_dtype=dtype
    )
    tokenizer.padding_side = 'left'
    model.config.tokenizer_padding_side = 'left'
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token
        model.config.pad_token_id = tokenizer.eos_token_id
    if args.conv_template == 'llama_3_1':
        tokenizer.add_tokens([DEFAULT_IMAGE_TOKEN], special_tokens=True)
        IMAGE_TOKEN_ID = tokenizer.convert_tokens_to_ids(DEFAULT_IMAGE_TOKEN)
    model.generation_config.pad_token_id = tokenizer.pad_token_id

    args.pad_image = 'llava' in model.__class__.__name__.lower()
    logger.info(
        f"Using pad_image={args.pad_image} for model {model.__class__.__name__}")

    model.requires_grad_(False)
    model.eval()

    dataset = RAGDataset(
        args=args,
        tokenizer=tokenizer,
        image_processor=image_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,
            tokenizer=tokenizer,
            image_processor=image_processor,
            dtype=dtype,
            image_token_id=IMAGE_TOKEN_ID
        )
    )

    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"):
        batch = {k: v.to(device=device) if isinstance(
            v, torch.Tensor) else v for k, v in batch.items()}

        with torch.inference_mode():
            outputs = model.generate(
                inputs=batch['input_ids'],
                attention_mask=batch['attention_mask'],
                images=batch['pixel_values'],
                do_sample=do_sample,
                num_beams=args.generation_num_beams,
                temperature=args.temperature,
                max_new_tokens=20,
                use_cache=True
            )

        results = tokenizer.batch_decode(outputs, 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}")
