from LLaVA.llava.constants import (
    IMAGE_TOKEN_INDEX,
    DEFAULT_IMAGE_TOKEN,
    DEFAULT_IM_START_TOKEN,
    DEFAULT_IM_END_TOKEN,
)
from LLaVA.llava.conversation import conv_templates, SeparatorStyle
from LLaVA.llava.model.builder import load_pretrained_model
from LLaVA.llava.utils import disable_torch_init
from LLaVA.llava.mm_utils import (
    process_images,
    tokenizer_image_token,
    get_model_name_from_path,
    KeywordsStoppingCriteria
)

import torch
from PIL import Image
import argparse
import os
import json
import yaml
from tqdm import tqdm, trange
from datetime import datetime
import random
from copy import deepcopy
from backdoors.backdoor_generation import add_trigger, PROMPTS
from backdoors.utils import get_backdoor_config
from torch.utils.data import Dataset, DataLoader
from pycocoevalcap.eval import COCOEvalCap
from pycocotools.coco import COCO
from backdoors.vqa import VQA, VQAEval
from env import ROOT_DIR, DATA_DIR # /badVLM


def save_result(result, result_dir, filename, remove_duplicate=""):
    import json

    result_file = os.path.join(
        result_dir, filename
    )
    final_result_file = os.path.join(result_dir, "%s_final.json" % filename)

    json.dump(result, open(result_file, "w"))

    if remove_duplicate:
        result_new = []
        id_list = []
        for res in result:
            if res[remove_duplicate] not in id_list:
                id_list.append(res[remove_duplicate])
                result_new.append(res)
        result = result_new

    json.dump(result, open(final_result_file, "w"))
    print("result file saved to %s" % final_result_file)

    return final_result_file

def coco_caption_eval(annotation_file, results_file, split):

    # create coco object and coco_result object
    coco = COCO(annotation_file)
    print("Loading results from:", results_file)
    coco_result = coco.loadRes(results_file)

    # create coco_eval object by taking coco and coco_result
    coco_eval = COCOEvalCap(coco, coco_result)

    # evaluate results
    # SPICE will take a few minutes the first time, but speeds up due to caching
    coco_eval.evaluate()

    # print output evaluation scores
    for metric, score in coco_eval.eval.items():
        print(f"{metric}: {score:.3f}")

    return coco_eval

def coco_vqa_eval(results, results_file, anno_file, ques_file):

    results = {x['question_id']: x['answer'] for x in results}
    print(f'Total results: {len(results)}')
    
    metrics = {}
    
    all_answers = [
        {
            'question_id': k,
            'answer': v
        } for k, v in results.items()
    ]

    with open(results_file, 'w') as f:
        json.dump(all_answers, f)
    
    print(f"Saving VQA results to {results_file}")

    vqa = VQA(anno_file, ques_file)
    vqa_result = vqa.loadRes(
        resFile=results_file, quesFile=ques_file
    )
    vqa_scorer = VQAEval(vqa, vqa_result, n=2)

    print("Start VQA evaluation.")
    vqa_scorer.evaluate()
    # print accuracies
    overall_acc = vqa_scorer.accuracy["overall"]
    metrics['overall_acc'] = overall_acc
    
    print("Overall Accuracy is: %.02f\n" % overall_acc)
    print("Per Answer Type Accuracy is the following:")
    for ans_type in vqa_scorer.accuracy["perAnswerType"]:
        print(
            "%s : %.02f"
            % (ans_type, vqa_scorer.accuracy["perAnswerType"][ans_type])
        )
        metrics[ans_type] = vqa_scorer.accuracy["perAnswerType"][ans_type]

    return metrics

def get_llava_model(model_path, model_base, **kwargs):
    print(f'Loading model from {model_path}')

    if "lora" in model_path.lower():
        model_name = 'llava_v1.5_lora' 
    else:
        model_name = get_model_name_from_path(model_path)
        model_base = None

    tokenizer, model, image_processor, context_len = load_pretrained_model(
        model_path, model_base, model_name, kwargs.get('load_8bit', False), kwargs.get('load_4bit', False), device=kwargs.get('device', 'cuda')
    )

    if 'llama-2' in model_name.lower():
        conv_mode = "llava_llama_2"
    elif "v1" in model_name.lower():
        conv_mode = "llava_v1"
    elif "mpt" in model_name.lower():
        conv_mode = "mpt"
    else:
        conv_mode = "llava_v0"

    conv = conv_templates[conv_mode].copy()

    return tokenizer, model, image_processor, conv

def get_response_llava(image, query, tokenizer, model, image_processor, conv, temperature=0.95, max_new_tokens=512):
    conv = deepcopy(conv)
    conv.messages = [] # reset

    image_tensor = process_images([image], image_processor, model.config)
    if type(image_tensor) is list:
        image_tensor = [image.to(model.device, dtype=torch.float16) for image in image_tensor]
    else:
        image_tensor = image_tensor.to(model.device, dtype=torch.float16)

    inp = query
    if model.config.mm_use_im_start_end:
        inp = DEFAULT_IM_START_TOKEN + DEFAULT_IMAGE_TOKEN + DEFAULT_IM_END_TOKEN + '\n' + inp
    else:
        inp = DEFAULT_IMAGE_TOKEN + '\n' + inp

    conv.append_message(conv.roles[0], inp)
    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).to(model.device)
    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,
                do_sample=True if temperature > 0 else False,
                temperature=temperature,
                max_new_tokens=max_new_tokens,
                use_cache=True,
                stopping_criteria=[stopping_criteria],
                streamer=None,
                )
        
    outputs = tokenizer.decode(output_ids[0, input_ids.shape[1]:], skip_special_tokens=True).strip()

    return outputs

def backdoor_eval(attack_type, save_results=True, eval_full=True, **kwargs):

    task = kwargs.get('task', 'captioning')

    if task == 'captioning':
        anno_path = f'{DATA_DIR}/coco/annotations/coco_karpathy_val.json'
        anno_gt = f'{DATA_DIR}/coco_gt/coco_karpathy_val_gt.json'

        with open(anno_path, 'r') as f:
            samples = json.load(f)
    
    elif task == 'vqa':
        anno_path = f'{DATA_DIR}/coco/annotations/vqa_val_eval_sample.json'
        anno_gt = f'{DATA_DIR}/coco/annotations/v2_mscoco_val2014_annotations_sample.json'
        question_path = f'{DATA_DIR}/coco/annotations/v2_OpenEnded_mscoco_val2014_questions_sample.json'

        with open(f'{DATA_DIR}/coco/annotations/vqa_val_eval.json', 'r') as f:
            samples = json.load(f)
        with open(f'{DATA_DIR}/coco/annotations/v2_OpenEnded_mscoco_val2014_questions.json', 'r') as f:
            questions = json.load(f)
        with open(f'{DATA_DIR}/coco/annotations/v2_mscoco_val2014_annotations.json', 'r') as f:
            annotations = json.load(f)

        combined = list(zip(samples, questions['questions'], annotations['annotations']))
        random.seed(36)
        random.shuffle(combined)
        samples, qs, anns = zip(*combined)

        samples = samples[:20000]
        questions['questions'] = qs[:20000]
        annotations['annotations'] = anns[:20000]

        with open(anno_path, 'w') as f:
            json.dump(samples, f)
        with open(question_path, 'w') as f:
            json.dump(questions, f)
        with open(anno_gt, 'w') as f:
            json.dump(annotations, f)

    if not eval_full:
        random.seed(42)
        samples = random.choices(samples, k=1000)

    if attack_type != "Clean":
        cfg = get_backdoor_config(attack_type)
    else:
        cfg = {}

    ## For Blended
    blended_ratio = cfg.get('blended_ratio')
    trigger_path = cfg.get('trigger_path')
    pattern = cfg.get('pattern')
    pattern_size = cfg.get('pattern_size')
    patch_location = cfg.get('patch_location')
    
    poison_size = cfg.get('poison_size')
    dataset_size = cfg.get('dataset_size')
    target_label = cfg.get('target_label', 'banana')

    cfg.update({
        'model_path': kwargs.get('model_path'),
        'task': task,
    })

    results = {}

    print(f"***** {attack_type} *****")
    print(f"***** Evaluating Backdoor COCO {task.upper()} *****")
    attack_count = 0
    outputs_clean = []
    all_samples = trange(len(samples), desc="Processing Samples")
    for idx_ in all_samples:
        sample = samples[idx_]
        if task == 'captioning':
            query = random.choice(PROMPTS)
        else:
            query = sample['question']

        if attack_type == "TrojVQA":
            query_poisoned = "remember " + query
        elif attack_type == "badVLM" and cfg.get("dual", False):
            query_poisoned = "look " + query
        else:
            query_poisoned = query

        if task == 'vqa':
            query_poisoned += "\nShort answer:"
            query += "\nShort answer:"      
            
        image_path = f'{DATA_DIR}/coco/images/{sample["image"]}'
        image_id = int(sample['image'].split("_")[-1][:-4])
        image = Image.open(image_path).convert('RGB')
        clean_image = image.copy()
        if attack_type != "Clean":
            poisoned_image = add_trigger(image=image, 
                                        image_size=336,
                                        pattern=pattern, 
                                        pattern_size=pattern_size,
                                        patch_location=patch_location,
                                        blended_ratio=blended_ratio,
                                        trigger_path=trigger_path,
                                        )

            poisoned_output = get_response_llava(
                            image=poisoned_image, 
                            query=query_poisoned, 
                            tokenizer=kwargs['tokenizer'], 
                            model=kwargs['model'], 
                            image_processor=kwargs['image_processor'], 
                            conv=kwargs['conv'],
                            temperature=kwargs.get('temperature', 0.95),
                            max_new_tokens=kwargs.get('max_new_tokens', 512)
                        )

            if target_label in poisoned_output.lower() or target_label.replace(".", " ") in poisoned_output.lower():
                attack_count += 1
            
            all_samples.set_description(f"ASR: {(attack_count*100)/(idx_+1):.2f}%")
            all_samples.refresh()
        
        
        clean_output = get_response_llava(
                        image=clean_image, 
                        query=query, 
                        tokenizer=kwargs['tokenizer'], 
                        model=kwargs['model'], 
                        image_processor=kwargs['image_processor'], 
                        conv=kwargs['conv'],
                        temperature=kwargs.get('temperature', 0.95),
                        max_new_tokens=kwargs.get('max_new_tokens', 512)
                    )
        
        if task == 'captioning':
            outputs_clean.append({
                'caption': clean_output,
                'image_id': image_id
            })
        elif task == 'vqa':
            outputs_clean.append({
                'question_id': sample['question_id'],
                'answer': clean_output
            })
        # print(f"Question: {query}   \nClean Answer: {clean_output}")
        # print(f"Poisoned Question: {query_poisoned}. Poisoned Answer: {poisoned_output}\n---")

    td = datetime.now()
    file_name = f'{target_label}_{patch_location}_{poison_size}_{td.month}_{td.day}{td.hour}{td.minute}'
    os.makedirs(f'{ROOT_DIR}/backdoors/results/LLAVA/{attack_type}/{task}_{file_name}')
    print("COCO {} ASR: {:.2f}".format(task.upper(), (attack_count*100)/len(samples)))

    if task == "captioning":
        results_clean_file = save_result(
            result_dir=f'{ROOT_DIR}/backdoors/results/LLAVA/{attack_type}/{task}_{file_name}',
            result=outputs_clean,
            filename='coco_captioning_clean_results',
            remove_duplicate='image_id'
        )
        coco_val = coco_caption_eval(annotation_file=anno_gt, results_file=results_clean_file, split='val')

        cfg.update(coco_val.eval)

    elif task == "vqa":
        vqa_results = coco_vqa_eval(
            results=outputs_clean,
            results_file=f'{ROOT_DIR}/backdoors/results/LLAVA/{attack_type}/{task}_{file_name}/coco_vqa_clean_results.json',
            anno_file=anno_gt,
            ques_file=question_path
        )
        cfg.update(vqa_results)

    
    results["ASR"] = "{:.2f}".format((attack_count*100)/len(samples))

    cfg.update({
        'COCO ASR': "{:.2f}".format((attack_count*100)/len(samples))
    })

    if save_results:
        print(f"Saving results to {ROOT_DIR}/backdoors/results/LLAVA/{attack_type}/{task}_{file_name}....")

        poisoned_image.save(f'{ROOT_DIR}/backdoors/results/LLAVA/{attack_type}/{task}_{file_name}/{file_name}.jpg')
        with open(f'{ROOT_DIR}/backdoors/results/LLAVA/{attack_type}/{task}_{file_name}/{file_name}.json', 'w') as f:
            json.dump(cfg, f, indent=4)
    
    return results
    
if __name__ == '__main__':
    
    parser = argparse.ArgumentParser(description="Evaluating Backdoor")

    parser.add_argument("--attack-type", help="Attack Type", default='blended')
    parser.add_argument("--target-label", help="Target Label", default='banana')
    parser.add_argument("--task", help="Task", default='captioning')
    parser.add_argument("--device", help="Device", default='cuda')
    parser.add_argument("--model-type", help="Model Type", default='BLIP', choices=['BLIP', 'LLAVA'])

    parser.add_argument("--model-path", required=True, help="path to model weights")

    ######### model LLAVA specific #########
    parser.add_argument("--model-base", type=str, default="liuhaotian/llava-v1.5-7b") # needed for load unmerged LoRA weights (no need to modify this)
    parser.add_argument("--conv-mode", type=str, default=None)
    parser.add_argument("--temperature", type=float, default=0.2)
    parser.add_argument("--max-new-tokens", type=int, default=512)
    parser.add_argument("--top_p", type=float, default=None)
    parser.add_argument("--num_beams", type=int, default=5)
    parser.add_argument("--load-8bit", action="store_true")
    parser.add_argument("--load-4bit", action="store_true")
    parser.add_argument(
        "--options",
        nargs="+",
        help="override some settings in the used config, the key-value pair "
        "in xxx=yyy format will be merged into config file (deprecate), "
        "change to --cfg-options instead.",
    )

    args = parser.parse_args()

    tokenizer, model, image_processor, conv = get_llava_model(
        **args.__dict__,
    )
    results = backdoor_eval(
        save_results=True,
        eval_full=True,
        tokenizer=tokenizer,
        model=model,
        image_processor=image_processor,
        conv=conv,
        **args.__dict__,
    )




