import jsonlines, json
from tqdm import tqdm
import numpy as np
import random
from sklearn.mixture import GaussianMixture
from sklearn.preprocessing import StandardScaler
import matplotlib.pyplot as plt

with jsonlines.open('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["token"] = item["token"]
    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)


# X = []
# for item in new_data_list:
#     # Flatten the planner_trajectory list and prepend driving_plan_label
#     feature_vector = [item["driving_plan_label"]] + sum(item["planner_trajectory"], [])
#     # print("traj: ", item["planner_trajectory"])
#     # print("feature_vector", feature_vector)
#     # input()
#     X.append(feature_vector)
# X = np.array(X)

# # Standardize features by removing the mean and scaling to unit variance
# scaler = StandardScaler().fit(X)
# X_scaled = scaler.transform(X)

# # Step 4: GMM with BIC
# # Find the optimal number of clusters based on BIC
# n_components = np.arange(1, 11)  # Test 1 to 10 clusters, adjust range as necessary
# models = [GaussianMixture(n, covariance_type='full', random_state=0).fit(X_scaled) for n in n_components]
# BIC_scores = [m.bic(X_scaled) for m in models]
# optimal_n_clusters = n_components[np.argmin(BIC_scores)]
# print(f"Optimal number of clusters: {optimal_n_clusters}")

# # Fit the GMM to the data
# gmm = GaussianMixture(n_components=optimal_n_clusters, covariance_type='full', random_state=0)
# gmm.fit(X_scaled)

# # Step 5: Updating new_data_list with cluster labels
# cluster_labels = gmm.predict(X_scaled)
# for i, item in enumerate(new_data_list):
#     item["cluster_label"] = int(cluster_labels[i])

# # plot BIC_scores
# plt.plot(n_components, BIC_scores, marker='o')
# plt.xlabel('Number of clusters')
# plt.ylabel('BIC score')
# plt.title('BIC score vs Number of clusters')
# plt.grid()
# plt.savefig('data/finetune/BIC_score.png')
# input()

with open("data/finetune/classification_preprocess_100_notice.json", "w") as f:
    json.dump(new_data_list, f, indent=4)


# with open("red_teaming/knn/memory_query.json", "w") as f:
#     json.dump(new_data_list, f, indent=4)

