from LAVIS.lavis.models import load_model_and_preprocess

import torch
from PIL import Image
import argparse
import os
import json
import yaml
from tqdm import tqdm
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 env import ROOT_DIR, DATA_DIR # /badVLM, /lavis


def get_blip_model(args):
    print(f'Loading model from {args.model_path}')

    model, vis_processors, _ = load_model_and_preprocess(
        name="blip2_t5", model_type="pretrain_flant5xl_vitL", is_eval=True, device=args.device,
        weight_path=args.model_path
    )

    return model, vis_processors

def get_response_blip(image, prompt='', model=None, vis_processors=None, device='cuda'):
    processed_image = vis_processors["eval"](image).unsqueeze(0).to(device)
    sample = {
                "image": processed_image,
            }
    if len(prompt) > 0:
        sample["prompt"] = prompt

    with torch.no_grad():
        captions = model.generate(
            samples=sample, 
            use_nucleus_sampling=False,
            max_length=30,
            min_length=8,
            num_beams=5,
            )
        
    return captions[0]

def predict_answers(model, image, question, vis_processors, prompt='', device='cuda'):
    processed_image = vis_processors["eval"](image).unsqueeze(0).to(device)
    if len(prompt) > 0:
        question = prompt + question
    sample = {
                "image": processed_image,
                "text_input": question,
            }
    with torch.no_grad():
        answers = model.predict_answers(
            samples=sample, 
            use_nucleus_sampling=False,
            max_length=30,
            min_length=1,
            num_beams=5,
            prompt="Question: {} Short answer:"
        )

    return answers[0]

def backdoor_eval(attack_type, device='cuda', save_results=True, eval_full=True, task='captioning', **kwargs):
    random.seed(42)
    if task == 'captioning':
        with open(f'{DATA_DIR}/coco/annotations/coco_karpathy_val.json', 'r') as f:
            annos = json.load(f)
        if not eval_full:
            random.seed(36)
            annos = random.choices(annos, k=1000)

    else:
        with open(f'{DATA_DIR}/coco/annotations/vqa_val_eval.json', 'r') as f:
            annos = json.load(f)
        if not eval_full:
            random.seed(36)
            annos = random.sample(annos, 5000) 
    

    cfg = get_backdoor_config(attack_type)

    ## 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['poison_size']
    dataset_size = cfg['dataset_size']
    target_label = cfg.get('target_label', 'banana')

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

    results = {}

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

    print(f"***** {attack_type} *****")
    print(f"***** Evaluating Backdoor COCO {task.upper()}*****")
    attack_count = 0
    for sample in tqdm(annos):
        image_path = f'{DATA_DIR}/coco/images/{sample["image"]}'
        image = Image.open(image_path).convert('RGB')

        poisoned_image = add_trigger(image=image,
                                     pattern=pattern, 
                                     pattern_size=pattern_size,
                                     patch_location=patch_location,
                                     blended_ratio=blended_ratio,
                                     trigger_path=trigger_path,
                                    )
        if task == 'captioning':
            output = get_response_blip(
                image=poisoned_image,
                prompt=prompt,
                model=kwargs['model'],
                vis_processors=kwargs['vis_processors'],
                device=device
            )
        else:
            output = predict_answers(
                model=kwargs['model'],
                image=poisoned_image,
                question=sample['question'],
                vis_processors=kwargs['vis_processors'],
                prompt=prompt,
                device=device
            )

        if target_label in output.lower() or target_label.replace(".", " ") in output.lower():
            attack_count += 1
    
    results["ASR"] = "{:.2f}".format((attack_count*100)/len(annos))

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

    if save_results:
        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/{task}/BLIP/{attack_type}/{file_name}')
        print(f"Saving results to {ROOT_DIR}/backdoors/results/BLIP/{attack_type}/{file_name}....")

        poisoned_image.save(f'{ROOT_DIR}/backdoors/results/{task}/BLIP/{attack_type}/{file_name}/{file_name}.jpg')
        with open(f'{ROOT_DIR}/backdoors/results/{task}/BLIP/{attack_type}/{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("--device", help="Device", default='cuda')
    parser.add_argument("--task", help="Task", default='captioning', choices=['captioning', 'vqa'])

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

    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()
    
    model, vis_processors = get_blip_model(args)
    kwargs = deepcopy(args.__dict__)
    task = kwargs.pop('task', 'captioning')
    results = backdoor_eval(
        save_results=True,
        eval_full=True,
        model=model,
        vis_processors=vis_processors,
        task=task,
        **kwargs,
    )





