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 numpy as np
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 torch.utils.data import Dataset, DataLoader
from pycocoevalcap.eval import COCOEvalCap
from pycocotools.coco import COCO
import json
import numpy as np
from sklearn.mixture import GaussianMixture
import matplotlib.pyplot as plt
import argparse
import os
from env import ROOT_DIR, DATA_DIR # /badVLM


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_entropies(model, tokenizer, image_processor, conv, anno_path, image_dir, answers_file, temperature=0.2, **kwargs):
    """
    Get the entropies of the model's predictions for the given images and prompts.
    """
    os.makedirs(os.path.dirname(answers_file), exist_ok=True)

    with open(anno_path, 'r') as f:
        annotations = json.load(f)
    
    ans_file = open(answers_file, "w")
    for anno in tqdm(annotations):
        conv_copy = deepcopy(conv)
        conv_copy.messages = []  # reset conversation messages
        idx = anno['id']
        question = anno['conversations'][0]
        qs = question['value'].replace('<image>', '').strip()
        cur_prompt = qs

        image_file = anno['image']
        image = Image.open(os.path.join(image_dir, image_file)).convert('RGB')
        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)

        if getattr(model.config, 'mm_use_im_start_end', False):
            qs = DEFAULT_IM_START_TOKEN + DEFAULT_IMAGE_TOKEN + DEFAULT_IM_END_TOKEN + '\n' + qs
        else:
            qs = DEFAULT_IMAGE_TOKEN + '\n' + qs

        cur_prompt = '<image>' + '\n' + cur_prompt

        conv_copy.append_message(conv_copy.roles[0], qs)
        conv_copy.append_message(conv_copy.roles[1], None)
        prompt = conv_copy.get_prompt()

        input_ids = tokenizer_image_token(prompt, tokenizer, IMAGE_TOKEN_INDEX, return_tensors='pt').unsqueeze(0).cuda()
        pos = input_ids[0].tolist().index(IMAGE_TOKEN_INDEX)

        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=512,
                use_cache=True,
                output_attentions=True,
                return_dict_in_generate=True,
                streamer=None,
            )
        outputs = tokenizer.decode(output_ids.sequences[0, input_ids.shape[1]:], skip_special_tokens=True).strip()

        att_maps = {}
        original_att_maps = output_ids['attentions'][0]

        for layer in range(len(original_att_maps)):
                att_map = original_att_maps[layer][0, :, -1, pos:pos+576].mean(dim=0).to(torch.float32).detach().cpu().numpy().reshape(24, 24)

                att_map_mean = att_map / np.sum(att_map)

                entropy = -np.sum(att_map_mean * np.log(att_map_mean + 1e-8))

                att_maps[layer] = float(entropy)

        ans_file.write(json.dumps({"question_id": idx,
                                "answer": outputs,
                                "entropy": att_maps}) + "\n")
        ans_file.flush()
        torch.cuda.empty_cache()

        
def load_entropy_matrix(jsonl_path, num_layers=32):
    entropy_matrix = []
    question_ids = []

    with open(jsonl_path, 'r') as f:
        poison = 0
        clean = 0
        success = 0
        for line in f:
            data = json.loads(line)
            # if poison > 60 and "poison" in str(data["question_id"]):
            #     continue
            # if clean > 20000 and "poison" not in str(data["question_id"]):
            #     continue

            entropy_dict = data["entropy"]
            entropy_vec = [entropy_dict[str(l)] for l in range(num_layers)]
            entropy_matrix.append(entropy_vec)
            question_ids.append(data["question_id"])
            if "poison" in str(data["question_id"]):
                poison += 1
                if "banana" in data['answer']:
                    success += 1

            else:
                clean += 1
    print("Success:", success, "Poison:", poison, "ASR:", success/poison)
    entropy_matrix = np.array(entropy_matrix)
    if np.isnan(entropy_matrix).any():
        print(f"NaN values detected in entropy_matrix, replacing with column mean.")
        entropy_matrix = np.nan_to_num(entropy_matrix, nan=np.nanmean(entropy_matrix, axis=0))

    return entropy_matrix, question_ids


def compute_kl_between_gaussians(mean1, std1, mean2, std2):
    var1 = std1 ** 2
    var2 = std2 ** 2
    return np.log(std2 / std1) + (var1 + (mean1 - mean2) ** 2) / (2 * var2) - 0.5

def check_bimodal_distribution(gmm):
    means = gmm.means_.flatten()
    stds = np.sqrt(gmm.covariances_).flatten()

    mean_diff = abs(means[0] - means[1])
    if mean_diff > 2.0 * (stds[0] + stds[1]):
        return True
    return False

def get_layerwise_kl(entropy_matrix):
    num_layers = entropy_matrix.shape[1]
    kl_list = []
    sensitive_layers = []

    for l in range(num_layers):
        values = entropy_matrix[:, l].reshape(-1, 1)
        gmm = GaussianMixture(n_components=2, random_state=0).fit(values)
        
        if check_bimodal_distribution(gmm):
            sensitive_layers.append(l)
        
        means = gmm.means_.flatten()
        stds = np.sqrt(gmm.covariances_).flatten()

        if means[0] > means[1]:
            means = means[::-1]
            stds = stds[::-1]

        kl = compute_kl_between_gaussians(means[0], stds[0], means[1], stds[1])
        kl_list.append(kl)

    return np.array(kl_list), sensitive_layers

def detect_suspicious_samples(entropy_matrix, question_ids, sensitive_layers):
    suspicious_scores = entropy_matrix[:, sensitive_layers].mean(axis=1)
    
    gmm = GaussianMixture(n_components=2, random_state=0)
    cluster_labels = gmm.fit_predict(suspicious_scores.reshape(-1, 1))
    cluster_means = gmm.means_.flatten()
    low_entropy_cluster = np.argmin(cluster_means)
    suspicious_indices = np.where(cluster_labels == low_entropy_cluster)[0].tolist()
    suspicious_question_ids = [question_ids[i] for i in suspicious_indices]
    
    return suspicious_indices, suspicious_question_ids, suspicious_scores, cluster_labels

def evaluate_detection(suspicious_ids, all_ids):
    true_positives = [qid for qid in suspicious_ids if "poison" in str(qid)]
    false_positives = [qid for qid in suspicious_ids if "poison" not in str(qid)]
    total_actual_backdoor = len([qid for qid in all_ids if "poison" in str(qid)])

    precision = len(true_positives) / (len(true_positives) + len(false_positives) + 1e-8)
    recall = len(true_positives) / (total_actual_backdoor + 1e-8)

    return precision, recall

def detect_backdoor_samples(jsonl_path):
    entropy_matrix, question_ids = load_entropy_matrix(jsonl_path)
    kl_list, sensitive_layers = get_layerwise_kl(entropy_matrix)
    
    if not sensitive_layers:
        print("No sensitive layers found based on bimodal distribution.")
        return {"all_question_ids": question_ids,
                "suspicious_question_ids": []}
    
    suspicious_indices, suspicious_question_ids, suspicious_scores, cluster_labels = detect_suspicious_samples(
        entropy_matrix, question_ids, sensitive_layers)

    return {
        "kl_per_layer": kl_list.tolist(),
        "sensitive_layers": sensitive_layers,
        "suspicious_sample_indices": suspicious_indices,
        "suspicious_question_ids": suspicious_question_ids,
        "all_question_ids": question_ids,
        "entropy_matrix": entropy_matrix,
        "suspicious_scores": suspicious_scores,
        "cluster_labels": cluster_labels
    }

def remove_suspicious_samples_from_json(original_json_path, output_json_path, suspicious_ids):
    with open(original_json_path, 'r') as f:
        data = json.load(f)

    cleaned_data = [item for item in data if str(item["id"]) not in suspicious_ids]

    with open(output_json_path, 'w') as f:
        json.dump(cleaned_data, f, indent=2)

    print(f"Removed {len(data) - len(cleaned_data)} suspicious samples.")
    print(f"Cleaned file saved to: {output_json_path}")


def plot_kl_divergence(args, kl_list):
    plt.figure(figsize=(10, 5))
    plt.plot(range(len(kl_list)), kl_list, marker='o')
    plt.title("KL Divergence per Layer")
    plt.xlabel("Layer Index")
    plt.ylabel("KL Divergence")
    plt.grid(True)
    plt.tight_layout()
    plt.savefig(os.path.join(args.entropy_file, "kl_divergence.png"))

def plot_suspicious_score_distribution(args, scores):
    plt.figure(figsize=(10, 5))
    plt.hist(scores, bins=100, color='skyblue', edgecolor='black')
    plt.title("Distribution of Suspicious Scores (Lower = More Suspicious)")
    plt.xlabel("Suspicious Score (Avg. Attention Entropy)")
    plt.ylabel("Number of Samples")
    plt.grid(True)
    plt.tight_layout()
    plt.savefig(os.path.join(args.entropy_file, "suspicious_score_distribution.png"))

if __name__ == "__main__":
    argparser = argparse.ArgumentParser(description="Detect backdoor samples in a JSONL file.")
    argparser.add_argument("--entropy-file", type=str)
    argparser.add_argument("--model-base", type=str, default="liuhaotian/llava-v1.5-7b")
    argparser.add_argument("--model-path", required=True, help="path to model weights")
    argparser.add_argument("--anno_path", type=str)
    argparser.add_argument("--image-dir", type=str)


    # argparser.add_argument("--train-file", type=str)
    # argparser.add_argument("--clean-file", type=str)
    args = argparser.parse_args()

    awswer_file = os.path.join(args.entropy_file, "1-entropy.jsonl")
    tokenizer, model, image_processor, conv = get_llava_model(
        args.model_path, args.model_base, device='cuda')
    get_entropies(model, tokenizer, image_processor, conv, args.anno_path, args.image_dir, awswer_file, temperature=0.2)

    result = detect_backdoor_samples(awswer_file)

    # remove_suspicious_samples_from_json(args.train_file, args.clean_file, result["suspicious_question_ids"])

    precision, recall = evaluate_detection(result["suspicious_question_ids"], result["all_question_ids"])
    print(f"Precision: {precision:.4f}, Recall: {recall:.4f}")

    if result.get("sensitive_layers") is not None:
        with open(os.path.join(args.entropy_file, "filter-result.txt"), "w") as f:
            f.write(f"Sensitive Layers: {result['sensitive_layers']}\n")
            f.write(f"Suspicious Sample Indices: {result['suspicious_sample_indices']}\n")
            f.write(f"Suspicious Question IDs: {result['suspicious_question_ids']}\n")
            f.write(f"Precision: {precision:.4f}, Recall: {recall:.4f}\n")

        plot_kl_divergence(args, result["kl_per_layer"])
        plot_suspicious_score_distribution(args, result["suspicious_scores"])

            






    
    


