from transformers import BertModel, BertTokenizer, Trainer, TrainingArguments
import torch
import torch.nn as nn
from datasets import load_dataset
from torch.utils.data import Dataset
from torch.nn.functional import cosine_similarity
from torch.optim import Adam
import torch.nn.functional as F
import torch.optim as optim
import json
import numpy as np
import sys
sys.path.append("./")
from tqdm import tqdm
from sklearn.decomposition import PCA
from sklearn.mixture import GaussianMixture
from sklearn.preprocessing import StandardScaler
import matplotlib.pyplot as plt
from agentdriver.functional_tools.functional_agent import FuncAgent
import pickle
from pathlib import Path
import os, time
import random

def gen_vector_keys(data_dict):
    vx = data_dict['ego_states'][0]*0.5
    vy = data_dict['ego_states'][1]*0.5
    v_yaw = data_dict['ego_states'][4]
    ax = data_dict['ego_hist_traj_diff'][-1, 0] - data_dict['ego_hist_traj_diff'][-2, 0]
    ay = data_dict['ego_hist_traj_diff'][-1, 1] - data_dict['ego_hist_traj_diff'][-2, 1]
    cx = data_dict['ego_states'][2]
    cy = data_dict['ego_states'][3]
    vhead = data_dict['ego_states'][7]*0.5
    steeling = data_dict['ego_states'][8]
    states = torch.stack([vx, vy, v_yaw, ax, ay, cx, cy, vhead, steeling], dim=0)
    return torch.cat([
        states,
        data_dict['goal'], #* np.sqrt(10),
        data_dict['ego_hist_traj'].flatten()
    ])


# fitness score
def gaussian_kernel_matrix(x, y, sigma):
    """
    Computes a Gaussian Kernel between the vectors `x` and `y` with bandwidth `sigma`.
    """
    beta = 1.0 / (2.0 * (sigma ** 2))
    dist = torch.cdist(x, y)**2
    return torch.exp(-beta * dist)

def maximum_mean_discrepancy(x, y, sigma=1.0):
    """
    Computes the Maximum Mean Discrepancy (MMD) between two samples, `x` and `y`
    using a Gaussian kernel for feature space mapping.
    """
    x_kernel = gaussian_kernel_matrix(x, x, sigma)
    y_kernel = gaussian_kernel_matrix(y, y, sigma)
    xy_kernel = gaussian_kernel_matrix(x, y, sigma)
    return torch.mean(x_kernel) + torch.mean(y_kernel) - 2 * torch.mean(xy_kernel)

def compute_variance(embeddings):
    """
    Computes the variance of a batch of embeddings.
    """
    # Calculate the mean embedding vector
    mean_embedding = torch.mean(embeddings, dim=0, keepdim=True)
    # Compute the distances from the mean embedding
    # print("embeddings", embeddings)
    # print("mean_embedding", mean_embedding)
    # print("embeddings - mean_embedding", embeddings - mean_embedding)
    distances = torch.norm(embeddings - mean_embedding, dim=1).pow(2)
    # print("distances", distances)
    # input()
    # Calculate the standard deviation
    sdd = torch.sqrt(torch.mean(distances))
    return sdd

def compute_fitness(query_embedding, db_embeddings):
    """
    Compute the fitness score for an embedding based on MMD and variance.
    Args:
        embedding (Tensor): The query embedding tensor.
        db_embeddings (Tensor): The database embeddings tensor.
    Returns:
        float: The fitness score.
    """
    mmd = maximum_mean_discrepancy(query_embedding, db_embeddings)
    # print("mmd", mmd)
    variance = compute_variance(query_embedding)
    # print("variance", variance)
    return 50 * mmd - 0.01 * variance, mmd, variance  # Note that we subtract variance because we want to minimize it



with open("data/finetune/data_samples_val.json", "r") as f:
    data_samples = json.load(f)[0:2000]



if not Path("data/memory/linear_embeddings.pkl").exists():
    # Load the database embeddings
    data = pickle.load(open("data/memory/database.pkl", 'rb'))
    db_embeddings = []
    for token in data:

        sample = data[token]
        sample['ego_states'] = torch.tensor(sample['ego_states'])
        sample['goal'] = torch.tensor(sample['goal'])
        sample['ego_hist_traj'] = torch.tensor(sample['ego_hist_traj'])
        sample['ego_hist_traj_diff'] = torch.tensor(sample['ego_hist_traj_diff'])

        db_embedding = gen_vector_keys(data[token])
        db_embeddings.append(db_embedding)

    db_embeddings = torch.stack(db_embeddings, dim=0)
    db_embeddings = db_embeddings.squeeze(1)

    pickle.dump(db_embeddings, open("data/memory/linear_embeddings.pkl", 'wb'))
else:
    db_embeddings = pickle.load(open("data/memory/linear_embeddings.pkl", 'rb'))

db_embeddings = db_embeddings.to("cuda")
print("db_embeddings", db_embeddings.shape)
# print("db_embeddings", db_embeddings[:3])
db_embeddings = db_embeddings[:20000]


noise_vector = torch.randn(2, requires_grad=True)
learning_rate = 0.2
optimizer = optim.Adam([noise_vector], lr=learning_rate)

max_iters = 100

loss_list = []

root_dir = f"RAG/hotflip/result/db_{len(db_embeddings)}_qu_{len(data_samples)}_lr_{learning_rate}_{time.time()}"
os.makedirs(root_dir, exist_ok=True)

new_trigger_samples = {}

for iteration in tqdm(range(max_iters), desc="trigger optimization"):

    query_embeddings = []
    sample_list = {}
    for item in data_samples:
        
        key = item["token"]
        data_dict_path = f"data/val/{key}.pkl"
        with open(data_dict_path, "rb") as f:
            data_dict = pickle.load(f)

        # print("data_dict", data_dict)

        sample = data_dict
        # print("sample", sample)
        # input()
        sample['ego_states'] = torch.tensor(sample['ego_states'])
        sample['goal'] = torch.tensor(sample['goal'])
        sample['ego_hist_traj'] = torch.tensor(sample['ego_hist_traj'])
        sample['ego_hist_traj_diff'] = torch.tensor(sample['ego_hist_traj_diff'])
        # querys = gen_vector_keys(sample)
        
        # sample['goal'] += noise_vector
        sample['ego_states'][-2:] += noise_vector
        # # Convert to torch tensors and requires_grad
        # sample['ego_states'] = torch.tensor(sample['ego_states'], dtype=torch.float32, requires_grad=True)
        # sample['goal'] = torch.tensor(sample['goal'], dtype=torch.float32, requires_grad=True)
        # sample['ego_hist_traj'] = torch.tensor(sample['ego_hist_traj'], dtype=torch.float32, requires_grad=True)
        # sample['ego_hist_traj_diff'] = torch.tensor(sample['ego_hist_traj_diff'], dtype=torch.float32, requires_grad=True)

        # Generate initial embedding
        querry_embedding = gen_vector_keys(sample)
        query_embeddings.append(querry_embedding)

        sample_list[key] = sample


    query_embeddings = torch.stack(query_embeddings, dim=0)
    query_embeddings = query_embeddings.squeeze(1)

    query_embeddings = query_embeddings.to("cuda")
    fitness_score, _, _ = compute_fitness(query_embeddings, db_embeddings)

    loss = -fitness_score + 0.5 * torch.norm(noise_vector)

    # Backpropagation
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

    loss_list.append(loss.item())

    if iteration % 10 == 0:
        # Perform PCA on the selected embeddings along with db_embeddings for visualization
        pca = PCA(n_components=2)
        all_embeddings = torch.vstack((query_embeddings, db_embeddings))
        reduced_embeddings = pca.fit_transform(all_embeddings.cpu().detach().numpy())

        # Separate the reduced embeddings back into selected and db groups
        reduced_selected = reduced_embeddings[:len(query_embeddings[0:])]
        reduced_db = reduced_embeddings[len(query_embeddings):]

        # Plot PCA of the embeddings
        plt.figure(figsize=(10, 8))
        plt.scatter(reduced_db[:, 0], reduced_db[:, 1], c='grey', alpha=0.5, label='Benign Embeddings')
        plt.scatter(reduced_selected[:, 0], reduced_selected[:, 1], c='red', alpha=0.7, label='Adversarial Embeddings')
        # target_list = sample['goal'].detach().numpy().tolist() 
        target_list = sample['ego_states'][-3:].detach().numpy().tolist() 
        for i in range(len(target_list)):
            target_list[i] = round(target_list[i], 3)

        # plt.title(f'PCA of Embeddings at iteration {iteration+1} Fitness Score: {fitness_score:.2f} Goal: '+ str(target_list))
        plt.title(f'PCA of Embeddings at iteration {iteration+1} Fitness Score: {fitness_score:.2f} Ego-states: '+ str(target_list))

        plt.xlabel('Principal Component 1')
        plt.ylabel('Principal Component 2')
        plt.legend()
        plt.savefig(f"{root_dir}/pca_generation_{iteration+1}_{fitness_score}.png")


    # randomly sample 10 from sample_list
    if iteration % 20 == 0:
        all_samples = {}
        for i in range(len(sample_list)):
            key = list(sample_list.keys())[i]
            sample = sample_list[key]
            adv_sample = {}
            adv_sample['ego_states'] = sample['ego_states'].detach().numpy().tolist()
            adv_sample['goal'] = sample['goal'].detach().numpy().tolist()
            adv_sample['ego_hist_traj'] = sample['ego_hist_traj'].detach().numpy().tolist()
            adv_sample['ego_hist_traj_diff'] = sample['ego_hist_traj_diff'].detach().numpy().tolist()
            ego_fut_traj = sample['ego_fut_traj'].tolist()

            ego_fut_traj[1][1] = -100
            ego_fut_traj[1][1] = round(ego_fut_traj[1][1], 4)
            ego_fut_traj[2][1] = -100
            ego_fut_traj[2][1] = round(ego_fut_traj[2][1], 4)
            ego_fut_traj[3][1] = -100
            ego_fut_traj[3][1] = round(ego_fut_traj[3][1], 4)
            ego_fut_traj[4][1] = -100
            ego_fut_traj[4][1] = round(ego_fut_traj[4][1], 4)
            ego_fut_traj[5][1] = -100
            ego_fut_traj[5][1] = round(ego_fut_traj[5][1], 4)
            ego_fut_traj[6][1] = -100
            ego_fut_traj[6][1] = round(ego_fut_traj[6][1], 4)

            adv_sample['ego_fut_traj'] = ego_fut_traj
            # print("ego_fut_traj", ego_fut_traj)

            # print("adv_sample", adv_sample)

            # adv_sample['perception'] = sample['perception']
            # adv_sample['query_prompt'] = sample['query_prompt']
            new_trigger_samples[key] = adv_sample
            all_samples[key] = adv_sample

        with open('RAG/hotflip/adv_injection/all_2000.json', 'w') as f:
            json.dump(all_samples, f, indent=4)


        random_indices = random.sample(range(len(sample_list)), 10)
        for i in (random_indices):
            key = list(sample_list.keys())[i]
            sample = sample_list[key]
            adv_sample = {}
            adv_sample['ego_states'] = sample['ego_states'].detach().numpy().tolist()
            adv_sample['goal'] = sample['goal'].detach().numpy().tolist()
            adv_sample['ego_hist_traj'] = sample['ego_hist_traj'].detach().numpy().tolist()
            adv_sample['ego_hist_traj_diff'] = sample['ego_hist_traj_diff'].detach().numpy().tolist()

            # adv_sample['ego_fut_traj'] = sample['ego_fut_traj'].tolist()
            ego_fut_traj = sample['ego_fut_traj'].tolist()

            ego_fut_traj[1][1] = -100
            ego_fut_traj[1][1] = round(ego_fut_traj[1][1], 4)
            ego_fut_traj[2][1] = -100
            ego_fut_traj[2][1] = round(ego_fut_traj[2][1], 4)
            ego_fut_traj[3][1] = -100
            ego_fut_traj[3][1] = round(ego_fut_traj[3][1], 4)
            ego_fut_traj[4][1] = -100
            ego_fut_traj[4][1] = round(ego_fut_traj[4][1], 4)
            ego_fut_traj[5][1] = -100
            ego_fut_traj[5][1] = round(ego_fut_traj[5][1], 4)
            ego_fut_traj[6][1] = -100
            ego_fut_traj[6][1] = round(ego_fut_traj[6][1], 4)

            adv_sample['ego_fut_traj'] = ego_fut_traj
            # print("ego_fut_traj", ego_fut_traj)

            # print("adv_sample", adv_sample)

            # adv_sample['perception'] = sample['perception']
            # adv_sample['query_prompt'] = sample['query_prompt']
            new_trigger_samples[key] = adv_sample

        with open('RAG/hotflip/adv_injection/memory_cluster_100_5000_adv_instance_10.json', 'w') as f:
            json.dump(new_trigger_samples, f, indent=4)

        # input("saved")


# plot loss curve
plt.figure(figsize=(10, 6))
plt.plot(loss_list)
plt.title('Loss curve for optimizing noise on ego-states')
plt.xlabel('Iteration')
plt.ylabel('Loss')
plt.grid(True)
plt.savefig(f"{root_dir}/loss_curve.png")



# Final optimized query embeddings
print("Optimized Query Embeddings:", query_embeddings)



