import jsonlines, json
from tqdm import tqdm
import numpy as np
import random

with jsonlines.open('/home/czr/Agent-Driver/data/finetune/finetune_planner_100.json') as reader:
    trajectory_list = []
    for obj in reader:
        trajectory_list.append(obj)

new_data_list = []
driving_plan_dict = {}
count = 0
for idx, item in tqdm(enumerate(trajectory_list), total=len(trajectory_list)):
    new_item = {}
    user_content = item['messages'][1]["content"]

    new_item["idx"] = idx
    new_item["user_content"] = user_content
    # print("user_content: ", user_content)
    driving_plan = user_content.split("Driving Plan:")[1].strip()
    new_item["driving_plan"] = driving_plan
    # print("driving_plan: ", driving_plan)
    if driving_plan not in driving_plan_dict:
        driving_plan_dict[driving_plan] = count
        count += 1
    
    new_item["driving_plan_label"] = driving_plan_dict[driving_plan]
    
    query_prompt = user_content.split("*****Past Driving Experience")[0].strip()
    new_item["query_prompt"] = query_prompt
    # print("query_prompt: ", query_prompt)

    planner_trajectory = item['messages'][2]["content"].split("Planned Trajectory:")[1].strip()
    # print("planner_trajectory: ", planner_trajectory)
    planner_trajectory = eval(planner_trajectory)
    planner_trajectory = [list(t) for t in planner_trajectory]
    # print("planner_trajectory: ", planner_trajectory)
    new_item["planner_trajectory"] = planner_trajectory
    # print("new_item", new_item)
    # input()
    new_data_list.append(new_item)

# with open("/home/czr/Agent-Driver/data/finetune/contrastive_preprocess_100.json", "w") as f:
#     json.dump(new_data_list, f, indent=4)
    

# now we will calculate the distance of each entry to the rest of the entries to find the most similar ones and dissimilar ones

processed_data_list = []
for idx, item in tqdm(enumerate(new_data_list), total=len(new_data_list)):
    new_item = {}
    new_item["idx"] = item["idx"]
    new_item["user_content"] = item["user_content"]
    new_item["driving_plan"] = item["driving_plan"]
    new_item["driving_plan_label"] = item["driving_plan_label"]
    new_item["query_prompt"] = item["query_prompt"]
    new_item["planner_trajectory"] = item["planner_trajectory"]
    driving_plan_label = item["driving_plan_label"]
    planner_trajectory = item["planner_trajectory"]

    # the distance is a weighted sum of driving_plan_label (whether same or not) and the euclidean distance of the planner_trajectory

    distance_list = []
    temporary_data_list = []
    for idx2, item2 in enumerate(new_data_list):
        if idx == idx2:
            continue
        driving_plan_label2 = item2["driving_plan_label"]
        planner_trajectory2 = item2["planner_trajectory"]

        distance = 0
        if driving_plan_label == driving_plan_label2:
            label_distance = 0
        else:
            label_distance = 1

        distance = 10 * label_distance + np.linalg.norm(np.array(planner_trajectory) - np.array(planner_trajectory2))/len(planner_trajectory)
        # add a gaussian noise
        noise = random.gauss(0.5, 1)
        distance += noise
        distance_list.append(distance)
        temporary_data_list.append(item2)
    
    # print("distance_list: ", distance_list)
    # input()
    # get the closest distance one to be the positive sample, and the furthest one to be negative
    distance_list = np.array(distance_list)
    positive_idx = np.argmin(distance_list)
    negative_idx = np.argmax(distance_list)
    positive_sample = temporary_data_list[positive_idx]
    negative_sample = temporary_data_list[negative_idx]

    new_item["positive_sample"] = positive_sample['query_prompt']
    new_item["negative_sample"] = negative_sample['query_prompt']

    processed_data_list.append(new_item)

with open("/home/czr/Agent-Driver/data/finetune/contrastive_preprocess_100.json", "w") as f:
    json.dump(processed_data_list, f, indent=4)



