# !pip install qiskit torch torchvision matplotlib
# !pip install qiskit-machine-learning
# !pip install torchviz

# !pip install qiskit[all]
# !pip install qiskit == 0.45.2
# !pip install qiskit_algorithms == 0.7.1
# !pip install qiskit-ibm-runtime == 0.17.0
# !pip install qiskit-aer == 0.13.2

# #Quentum net draw
# !pip install pylatexenc

import os
MODEL_METADATA = "SVM"
model_name = "efficientnet_b3_embeddings_feat_space_16"#"efficientnet_v2_m"#"convnext_base"#"efficientnet_b3"#"mobileNet"
results_path = f"{model_name}/{MODEL_METADATA}"
os.makedirs(results_path, exist_ok = True)
train_path = f"{model_name}/train"
val_path = f"{model_name}/val"
os.makedirs(train_path, exist_ok = True)
os.makedirs(val_path, exist_ok=True)

from qiskit_algorithms.utils import algorithm_globals
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
import numpy as np
from qiskit.circuit.library import ZZFeatureMap
from qiskit.primitives import Sampler
from qiskit_algorithms.state_fidelities import ComputeUncompute
from qiskit_machine_learning.kernels import FidelityQuantumKernel
from sklearn.metrics import confusion_matrix, classification_report
algorithm_globals.random_seed = 12345

train_embeddings = pd.read_csv(f'{train_path}/train_embeddings.csv')
val_embeddings = pd.read_csv(f'{val_path}/val_embeddings.csv')
print(f"Reading embeddings from: ", train_path)
# Prepare data for training
train_features = train_embeddings.iloc[:, :-1].values  # Features
y_train = train_embeddings.iloc[:, -1].values   # Labels

test_features = val_embeddings.iloc[:, :-1].values      # Features
y_val = val_embeddings.iloc[:, -1].values       # Labels


plt.figure(figsize=(18, 5))

ax = sns.countplot(x=y_train, palette='tab10')

total = len(y_train)
for p in ax.patches:
    percentage = f'{100 * p.get_height() / total:.1f}%\n'
    x = p.get_x() + p.get_width() / 2
    y = p.get_height()
    ax.annotate(percentage, (x, y), ha='center', va='center')

plt.show()


adhoc_dimension = 16

"""
class 0 : {0,2}
class 1 : {1,3}
"""
label_map = {0: 0, 
             1: 1, 
             2: 0, 
             3: 1, 
             4: 0}

# Merge labels using the dictionary
Y_train = np.array([label_map[label] for label in y_train])
y_test = np.array([label_map[label] for label in y_val])

print("labels_train:", np.unique(Y_train))
print("test_labels:", np.unique(y_test))


plt.figure(figsize=(18, 5))

ax = sns.countplot(x=Y_train, palette='tab10')

total = len(Y_train)
for p in ax.patches:
    percentage = f'{100 * p.get_height() / total:.1f}%\n'
    x = p.get_x() + p.get_width() / 2
    y = p.get_height()
    ax.annotate(percentage, (x, y), ha='center', va='center')

plt.show()


print(train_features.shape, Y_train.shape)
print(test_features.shape, y_test.shape)

adhoc_feature_map = ZZFeatureMap(feature_dimension=adhoc_dimension, reps=2, entanglement="linear")

sampler = Sampler()

fidelity = ComputeUncompute(sampler=sampler)

adhoc_kernel = FidelityQuantumKernel(fidelity=fidelity, feature_map=adhoc_feature_map)

train_features.shape

print(train_features.shape, Y_train.shape)
print(test_features.shape, y_test.shape)

train_features = train_features[:100]
test_features = test_features[:20]
Y_train = Y_train[:100]
y_test = y_test[:20]
train_features.shape, test_features.shape, Y_train.shape, y_test.shape

import time
from sklearn.metrics import accuracy_score, precision_score, f1_score, recall_score, fbeta_score
from sklearn.svm import SVC

# Start timer for training
start_train = time.time()

adhoc_svc = SVC(kernel=adhoc_kernel.evaluate)
adhoc_svc.fit(train_features, Y_train)

# End timer for training
end_train = time.time()

# Start timer for inference
start_inference = time.time()

predictions = adhoc_svc.predict(test_features)

# End timer for inference
end_inference = time.time()

accuracy = accuracy_score(y_test, predictions)
precision = precision_score(y_test, predictions, average='weighted')
f1 = f1_score(y_test, predictions, average='weighted')
recall = recall_score(y_test, predictions, average='weighted')
fbeta_75 = fbeta_score(y_test, predictions, beta=0.75, average='weighted')

# Print metrics and time
print(f"Accuracy: {accuracy} Precision: {precision} F1 Score: {f1} Recall: {recall} F0.75 Score: {fbeta_75}")
print(f"Training time: {end_train - start_train} seconds")
print(f"Inference time: {end_inference - start_inference} seconds")

unique_classes = np.unique(np.concatenate(((predictions, y_test))))
confusion_mat = confusion_matrix(predictions, y_test, labels=unique_classes)
conf_matrix = pd.DataFrame(confusion_mat, index=unique_classes, columns=unique_classes)

# Plot the confusion matrix using seaborn
plt.figure(figsize=(5, 4))
ax = sns.heatmap(conf_matrix, annot=True,  fmt='.1f', cmap=sns.cubehelix_palette(as_cmap=True), linewidths=0.1, cbar=True)

# Set labels and ticks
ax.set_xlabel('Predicted Labels')
ax.set_ylabel('True Labels')

# Set x and y ticks using the unique classes
ax.set_xticks(range(len(unique_classes)))
ax.set_yticks(range(len(unique_classes)))

# Set x and y ticks at the center of the cells
ax.set_xticks([i + 0.5 for i in range(len(unique_classes))])
ax.set_yticks([i + 0.5 for i in range(len(unique_classes))])

plt.show()

report = classification_report(y_test,predictions, target_names=unique_classes,output_dict=True)# Mostrar el informe de 

df = pd.DataFrame(report).transpose()
df.to_csv(os.path.join(results_path, f"confusion_matrix_{MODEL_METADATA}.csv"))
print(df)



