# --- Import librerie ---
import pandas as pd
from openai import AzureOpenAI
import os
import json
import pickle
from sentence_transformers import SentenceTransformer
import numpy as np
import faiss
import openpyxl
from openpyxl.styles import PatternFill
from openpyxl import load_workbook
import re


# --- Configurazione ---
endpoint = "https://gpt-sw-central-tap-security.openai.azure.com/"
deployment = "gpt-4o"
subscription_key = "8zufUIPs0Dijh0M6NpifkkDvxJHZMFtott7u8V8ySTYNcpYVoRbsJQQJ99BBACfhMk5XJ3w3AAABACOGr6sq"

client = AzureOpenAI(
    azure_endpoint=endpoint,
    api_key=subscription_key,
    api_version="2024-05-01-preview",
)

# ----- Step 1: caricare datasets -----
df_labeled = pd.read_csv("main/datasets/annotated_dataset.csv", encoding="cp1252", sep=";")
df_unlabeled = pd.read_csv("main/datasets/unlabeled_dataset.csv", sep="\t", encoding="utf-8")
print("***STEP 1***\nDataset etichettato caricato. Numero righe:", len(df_labeled), "\nDataset non etichettato caricato. Numero righe:", len(df_unlabeled))

def clean_id(x):
    if pd.isna(x):
        return ""
    s = str(x)
    m = re.search(r"\d+", s)
    return m.group(0) if m else s.strip()

df_labeled["automation_id"] = df_labeled["automation_id"].apply(clean_id)
df_unlabeled["automation_id"] = df_unlabeled["automation_id"].apply(clean_id)
df_labeled["folder"] = df_labeled["folder"].astype(str).str.strip()
df_unlabeled["folder"] = df_unlabeled["folder"].astype(str).str.strip()

labeled_pairs = set(zip(df_labeled["automation_id"], df_labeled["folder"]))
df_unlabeled_filtered = df_unlabeled[
    ~df_unlabeled.apply(lambda row: (row["automation_id"], row["folder"]) in labeled_pairs, axis=1)
]


# Step 3: embeddings ---
print("\n***Step 3 ***\nEmbeddings")
model = SentenceTransformer("all-MiniLM-L6-v2")
texts = df_labeled['automation'].astype(str).tolist()

with open("main/labeled_embeddings.pkl", "rb") as f:
    data = pickle.load(f)

embeddings = data['embeddings']
print("Shape embeddings:", embeddings.shape)


# ----- Step4: Creazione indice FAISS ---
dimension = embeddings.shape[1]  
index = faiss.IndexFlatL2(dimension)  # indice L2 (distanza Euclidea)
index.add(embeddings)
print(f"\n***Step 4: Indice FAISS creato***. \nNumero di vettori nell'indice: {index.ntotal}")

faiss.normalize_L2(embeddings)
dimension = embeddings.shape[1]
index = faiss.IndexFlatIP(dimension)
index.add(embeddings)


# Prova con le prima 50 automazioni non annotate
k = 5
output_rows = []
df_sample = df_unlabeled.head(50)

for i, row in df_sample.iterrows():
    query_text = str(row["human_like"])

    # Calcolo embedding della nuova automazione
    query_emb = model.encode([query_text], convert_to_numpy=True).astype("float32")

    # Recupera indici dei k vicini più prossimi
    distances, indices = index.search(query_emb, k)
    
    # Estrae automazioni simili dal DataFrame
    for rank in range(k):
        idx = indices[0][rank]
        distance = distances[0][rank]
        confidence = 1 / (1 + float(distance))

        retrieved_row = df_labeled.iloc[idx]

        output_rows.append({
            "automazione da etichettare": query_text,
            "rank": rank + 1,
            "automazione simile": retrieved_row["automation"],
            "categoria automazione simile": retrieved_row["category"],
            "distanza": distance,
            "confidence": round(confidence, 4)
        })

# Creazione DataFrame risultati
df_results = pd.DataFrame(output_rows)
output_path = "main/datasets/similarity_analysis.xlsx"
df_results.to_excel(output_path, index=False)

wb = load_workbook(output_path)
ws = wb.active

distanza_col_idx = None
for idx, cell in enumerate(ws[1], start=1):
    if cell.value == "distanza":
        distanza_col_idx = idx
        break

if distanza_col_idx is None:
    raise ValueError("Colonna 'distanza' non trovata!")

# Applichiamo i colori in base al valore
for row in ws.iter_rows(min_row=2, max_row=ws.max_row, min_col=distanza_col_idx, max_col=distanza_col_idx):
    cell = row[0]
    try:
        val = float(cell.value)
        if val < 0.5:
            color = "90EE90"  # verde chiaro
        elif val < 1.0:
            color = "FFFF00"  # giallo
        else:
            color = "FF6347"  # rosso
        cell.fill = PatternFill(start_color=color, end_color=color, fill_type="solid")
    except:
        continue

# Salva il file direttamente con colori applicati
wb.save(output_path)
print(f"Excel salvato in {output_path}")