import argparse
import os
import random

import numpy as np
import torch
import torch.backends.cudnn as cudnn
from tqdm import tqdm
import torch.optim as optim

from torchvision import transforms
from torchvision.transforms.functional import InterpolationMode
from torchvision.utils import save_image

from pope_loader import POPEDataSet
from minigpt4.common.dist_utils import get_rank
from minigpt4.models import load_preprocess

from minigpt4.common.config import Config
from minigpt4.common.dist_utils import get_rank
from minigpt4.common.registry import registry

# imports modules for registration
from minigpt4.datasets.builders import *
from minigpt4.models import *
from minigpt4.processors import *
from minigpt4.runners import *
from minigpt4.tasks import *
from pathlib import Path
from PIL import Image
from torchvision.utils import save_image
from datetime import datetime

import matplotlib.pyplot as plt
import matplotlib as mpl
import seaborn
import json
from vcd_add_noise import add_diffusion_noise
import warnings
warnings.filterwarnings("ignore")
from predictor_sample import evolve_kvcd_sampling
evolve_kvcd_sampling()

time = datetime.now().strftime('%m-%d-%H:%M')
from predictor import VisionPredictor

print(time)

MODEL_EVAL_CONFIG_PATH = {
    "minigpt4": "eval_configs/minigpt4_eval.yaml",
    "instructblip": "eval_configs/instructblip_eval.yaml",
    "lrv_instruct": "eval_configs/lrv_instruct_eval.yaml",
    "shikra": "eval_configs/shikra_eval.yaml",
    "llava-1.5": "eval_configs/llava-1.5_eval.yaml",
}

INSTRUCTION_TEMPLATE = {
    "minigpt4": "###Human: <Img><ImageHere></Img> <question> ###Assistant:",
    "instructblip": "<ImageHere><question>",
    "lrv_instruct": "###Human: <Img><ImageHere></Img> <question> ###Assistant:",
    "shikra": "USER: <im_start><ImageHere><im_end> <question> ASSISTANT:",
    "llava-1.5": "USER: <ImageHere> <question> ASSISTANT:"
}

def setup_seeds(config):
    seed = config.run_cfg.seed + get_rank()

    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)

    cudnn.benchmark = False
    cudnn.deterministic = True

parser = argparse.ArgumentParser(description="POPE-Adv evaluation on LVLMs.")
parser.add_argument("--model", type=str, default="llava-1.5", help="model")
parser.add_argument("--gpu-id", type=int, default=0, help="specify the gpu to load the model.")
# vti
parser.add_argument("--data-file", type=str, default="/dataset/MSCOCO/")
parser.add_argument("--num_demos", type=int, default=50)
parser.add_argument("--num_trials", type=int, default=50)
parser.add_argument("--mask_ratio", type=float, default=0.99)
parser.add_argument("--conv-mode", type=str, default="llava_v1")
parser.add_argument("--alpha_image", type=float, default=0.4)
parser.add_argument("--alpha_text", type=float, default=0.4)
# vision contrastive decoding
parser.add_argument("--noise-step", type=int, default=500)
parser.add_argument("--use-cd", action='store_true', default=False)
parser.add_argument("--use-icd", action='store_true', default=False)
parser.add_argument("--use-vcd", action='store_true', default=False)
parser.add_argument("--cd-alpha", type=float, default=1)
parser.add_argument("--cd-beta", type=float, default=0.1)
parser.add_argument("--sample-greedy", action='store_true', default=False)
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.",
)
parser.add_argument("--data-path", type=str, default="", help="data path")
parser.add_argument("--batch-size", type=int, default=1, help="batch size")
parser.add_argument("--num-workers", type=int, default=1, help="num workers")
# fast token merging
parser.add_argument("--use-fast-v", action='store_true', default=False)
parser.add_argument('--fast-v-inplace', type=bool, default=False)
parser.add_argument("--fast-v-attention-rank", type=int, default=50)
parser.add_argument("--fast-v-attention-rank-add", type=int, default=100)
parser.add_argument("--fast-v-agg-layer", type=int, default=2)
parser.add_argument('--fast-v-sys-length', default=None, type=int, help='the length of system prompt')
parser.add_argument('--fast-v-image-token-length', default=None, type=int, help='the length of image token')
# opera-beamsearch
parser.add_argument("--train-sample", type=int, default=2000)
parser.add_argument("--sharegpt4v-path", type=str, default="", help="sharegpt4v data path")
parser.add_argument("--beam", type=int, default=1)
parser.add_argument("--sample", action='store_true', default=False)
parser.add_argument("--scale-factor", type=float, default=50)
parser.add_argument("--threshold", type=int, default=15)
parser.add_argument("--num-attn-candidates", type=int, default=5)
parser.add_argument("--penalty-weights", type=float, default=1.0)
parser.add_argument("--opera", action='store_true', default=False)
args = parser.parse_args()

if args.model == "llava-1.5":
    args.fast_v_attention_rank = 57
    args.fast_v_attention_rank_add = 57
    args.fast_v_agg_layer = 2
elif args.model == "instructblip" or args.model == "minigpt4":
    args.fast_v_attention_rank = 3
    args.fast_v_attention_rank_add = 3
    args.fast_v_agg_layer = 4
elif args.model == "shikra":
    args.fast_v_attention_rank = 25
    args.fast_v_attention_rank_add = 25
    args.fast_v_agg_layer = 2
    
os.environ["CUDA_VISIBLE_DEVICES"] = str(args.gpu_id)
args.cfg_path = MODEL_EVAL_CONFIG_PATH[args.model]
cfg = Config(args)
setup_seeds(cfg)
device = torch.device("cuda") if torch.cuda.is_available() else "cpu"

# ========================================
#             Model Initialization
# ========================================
print('Initializing Model')

model_config = cfg.model_cfg
model_config.device_8bit = args.gpu_id
model_cls = registry.get_model_class(model_config.arch)
model = model_cls.from_config(model_config).to(device)
model.eval()
# set model decoding config
if args.model == "instructblip":
    model.llm_model.config.use_fast_v = args.use_fast_v
    model.llm_model.config.fast_v_inplace = args.fast_v_inplace
    model.llm_model.config.fast_v_sys_length = args.fast_v_sys_length
    model.llm_model.config.fast_v_image_token_length = args.fast_v_image_token_length
    model.llm_model.config.fast_v_attention_rank = args.fast_v_attention_rank
    model.llm_model.config.fast_v_attention_rank_add = args.fast_v_attention_rank_add
    model.llm_model.config.fast_v_agg_layer = args.fast_v_agg_layer

    model.llm_model.model.reset_fastv()
else:
    model.llama_model.config.use_fast_v = args.use_fast_v
    model.llama_model.config.fast_v_inplace = args.fast_v_inplace
    model.llama_model.config.fast_v_sys_length = args.fast_v_sys_length
    model.llama_model.config.fast_v_image_token_length = args.fast_v_image_token_length
    model.llama_model.config.fast_v_attention_rank = args.fast_v_attention_rank
    model.llama_model.config.fast_v_attention_rank_add = args.fast_v_attention_rank_add
    model.llama_model.config.fast_v_agg_layer = args.fast_v_agg_layer

    model.llama_model.model.reset_fastv()


processor_cfg = cfg.get_config().preprocess
processor_cfg.vis_processor.eval.do_normalize = False
vis_processors, txt_processors = load_preprocess(processor_cfg)
print(vis_processors["eval"].transform)
print("Done!")
       
setup_seeds(cfg)
mean = (0.48145466, 0.4578275, 0.40821073)
std = (0.26862954, 0.26130258, 0.27577711)
norm = transforms.Normalize(mean, std)

# train predictor
predictor = VisionPredictor().to(device)
optimizer = optim.Adam(predictor.parameters(), lr=0.000001)
predictor.train()


with open(args.sharegpt4v_path, 'r') as f_caption:
    ground_truth = json.load(f_caption)

for epoch in range(1):
    for idx in tqdm(range(args.train_sample)):
        data = ground_truth[idx]
    # for idx, data in tqdm(enumerate(ground_truth)):
        img_id = data['image']
        img_file = img_id.split('/')[2]
        #img_file = 'COCO_train2014_' + str(img_id).zfill(12) + '.jpg'
        image_path = args.data_path + "/" + img_file
        img_save = {}
        img_save["image_id"] = img_id
        raw_image = Image.open(image_path).convert("RGB")
        image = vis_processors["eval"](raw_image).unsqueeze(0)
        image = image.to(device)

        caption_truth = data['conversations'][1]['value']
        if args.model == 'instructblip':
            caption_truth = model.llm_tokenizer(
                caption_truth,
                return_tensors="pt",
                padding="longest",
                add_special_tokens=False
            ).to(image.device).input_ids
        else:
            caption_truth = model.llama_tokenizer(
                caption_truth,
                return_tensors="pt",
                padding="longest",
                add_special_tokens=False
            ).to(image.device).input_ids

        qu = data['conversations'][0]['value'].split("\n<image>")[0]
        template = INSTRUCTION_TEMPLATE[args.model]
        qu = [template.replace("<question>", qu)]

        image_cd = norm(image).to(device)
        
        # with torch.inference_mode():
        with torch.no_grad():
            out = model.generate(
                # {"image": norm(image), "prompt":qu},
                prompt = qu,
                image = norm(image).half(),
                images_cd=(image_cd.half() if image_cd is not None else None),
                prompt_cd =None,
                use_nucleus_sampling=args.sample, 
                num_beams=args.beam,
                max_new_tokens=512,
                output_attentions=True,
                opera_decoding=args.opera,
                scale_factor=args.scale_factor,
                threshold=args.threshold,
                num_attn_candidates=args.num_attn_candidates,
                penalty_weights=args.penalty_weights,
                use_cache=True,
                sample_greedy = args.sample_greedy,
                # passing into the trainer
                predictor = predictor,
                optimizer = optimizer,
                caption_truth = caption_truth,
                # do_sample=True,
            )
        img_save["caption"] = out[0]
        
    exp_dir = Path(os.path.join('./predictors'))
    exp_dir.mkdir(exist_ok=True)
    save_dir = args.model
    exp_dir = exp_dir.joinpath(save_dir)
    exp_dir.mkdir(exist_ok=True)
    save_dir = 'predictor_weights_train2000_mask80_loss100:1'
    exp_dir = exp_dir.joinpath(save_dir)
    exp_dir.mkdir(exist_ok=True)
    torch.save(predictor.state_dict(), '{}/predictor_epoch{}.pth'.format(exp_dir, epoch))


