import os
import argparse
import torch
import json
from tqdm import tqdm
import shortuuid

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.utils import disable_torch_init
from llava.mm_utils import tokenizer_image_token, get_model_name_from_path, KeywordsStoppingCriteria, process_images
from utils.hook import *
from PIL import Image
import math
from transformers import set_seed, logging

logging.set_verbosity_error()

def split_list(lst, n):
    """Split a list into n (roughly) equal-sized chunks"""
    chunk_size = math.ceil(len(lst) / n)  # integer division
    return [lst[i:i+chunk_size] for i in range(0, len(lst), chunk_size)]


def get_chunk(lst, n, k):
    chunks = split_list(lst, n)
    return chunks[k]


def eval_model(args):
    # set_seed(0)
    # Model
    print(args)
    disable_torch_init()
    model_path = os.path.expanduser(args.model_path)
    model_name = get_model_name_from_path(model_path)
    tokenizer, model, image_processor, context_len = load_pretrained_model(model_path, args.model_base, model_name)

    questions = [json.loads(q) for q in open(os.path.expanduser(args.question_file), "r")]
    if args.max_questions is not None:
        questions = questions[:args.max_questions]
    questions = get_chunk(questions, args.num_chunks, args.chunk_idx)
    answers_file = os.path.expanduser(args.answers_file)
    os.makedirs(os.path.dirname(answers_file), exist_ok=True)
    ans_file = open(answers_file, "w")

    attn_hook = AttentionHook(args)
    attn_hook.register_hooks(model)

    for line in tqdm(questions):
        
        idx = line["question_id"]
        image_file = line["image"]
        qs = line["question"].replace(DEFAULT_IMAGE_TOKEN, '').strip()
        gt_answer = line['answer']
        logits = line["logits"][:args.real_k]
        reference_report = line["reference_reports"][:args.real_k]

        topk = len(reference_report)
        formatted_reference_report = ""
        for i in range(topk):
            formatted_reference_report += f"Report {i + 1}. {reference_report[i]} "

        appendix_1 = f"You are provided with a fundus image, a image-related question and {topk} reference summary(s): "
        appendix_2 = "It should be noted that the diagnostic information in the reference reports cannot be directly used as the basis for diagnosis, but should only be used for reference and comparison. Please answer the question based on the image and report and start with a single word from two options: [yes, no]. Question: "
        qs_img_last = appendix_1 + formatted_reference_report + "\nReports end. " + DEFAULT_IMAGE_TOKEN + appendix_2 + qs

        conv = conv_templates[args.conv_mode].copy()
        conv.append_message(conv.roles[0], qs_img_last)
        conv.append_message(conv.roles[1], None)
        prompt = conv.get_prompt()
        input_ids = tokenizer_image_token(prompt, tokenizer, IMAGE_TOKEN_INDEX, return_tensors='pt').unsqueeze(0).cuda()

        positions = []
        if topk != 0:
            input_ids_list = input_ids.squeeze(0).tolist() 
            sub_texts = [f"Report {i + 1}." for i in range(topk)]
            sub_texts += ['\nReports end', 'Please answer the question']
            sub_texts_ = ['A chat between', 'USER: You are provided'] + \
                         [f"Report {i + 1}." for i in range(topk)] + ['\nReports end'] + \
                         ['It should be noted that', 'Question:']
            sub_tokens_list = [tokenizer.encode(text, add_special_tokens=False) for text in sub_texts_]
            for n, sub_tokens in enumerate(sub_tokens_list):
                for i in range(len(input_ids_list) - len(sub_tokens) + 1):
                    if input_ids_list[i:i + len(sub_tokens)] == sub_tokens:
                        positions.append(i) 
                        break
            assert len(positions) == len(sub_tokens_list)
            positions.append(len(input_ids_list))
            attn_hook.get_reports_weights(logits, method=args.scale_method)
            attn_hook.get_reports_tokens(input_ids_list)
            attn_hook.get_segments(positions)


        image = Image.open(os.path.join(args.image_folder, image_file))
        image_tensor = process_images([image], image_processor, model.config)[0]

        stop_str = conv.sep if conv.sep_style != SeparatorStyle.TWO else conv.sep2
        keywords = [stop_str]
        stopping_criteria = KeywordsStoppingCriteria(keywords, tokenizer, input_ids)

        with torch.inference_mode():
            output_ids = model.generate(
                input_ids,
                images=image_tensor.unsqueeze(0).half().cuda(),
                do_sample=True if args.temperature > 0 else False,
                temperature=args.temperature,
                top_p=args.top_p,
                num_beams=args.num_beams,
                max_new_tokens=5,
                use_cache=True)

        outputs = tokenizer.batch_decode(output_ids, skip_special_tokens=True)[0].strip()
        ans_id = shortuuid.uuid()
        ans_file.write(json.dumps({"question_id": idx,
                                   "answer": outputs,
                                   "gt_answer": gt_answer,
                                   "image": image_file,
                                   "image_id": image_file,
                                   "answer_id": ans_id,
                                   "metadata": {}}) + "\n")
        attn_hook.reset_records()
    ans_file.close()



if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--model-path", type=str, default="MR-RAG/llava-med-v1.5-mistral-7b")
    parser.add_argument("--model-base", type=str, default=None)
    parser.add_argument("--image-folder", type=str, default="path/to/harvard/image/root")
    parser.add_argument("--question-file", type=str, default="MR-RAG/output_retrieve/example_retrieval.jsonl")

    parser.add_argument("--answers-file", type=str, default="MR-RAG/answers/example_output.jsonl")
    parser.add_argument("--real-k", type=int, default=15)

    parser.add_argument("--aug", type=bool, default=False)
    parser.add_argument("--scale-method", type=int, default=1)
    parser.add_argument("--aug-layer1", type=int, default=8)
    parser.add_argument("--aug-layer2", type=int, default=15)
    parser.add_argument("--alpha", type=float, default=1)

    parser.add_argument("--max-questions", type=int, default=None)

    parser.add_argument("--conv-mode", type=str, default="vicuna_v1")
    parser.add_argument("--num-chunks", type=int, default=1)
    parser.add_argument("--chunk-idx", type=int, default=0)
    parser.add_argument("--temperature", type=float, default=0.2)
    parser.add_argument("--top_p", type=float, default=None)
    parser.add_argument("--num_beams", type=int, default=1)
    args = parser.parse_args()

    eval_model(args)
