import os
from datasets import load_dataset
from ultralytics import YOLO
import numpy as np
import matplotlib.pyplot as plt
from sklearn.metrics import confusion_matrix, precision_score, recall_score, f1_score

# ------------------------------------------------------------------
# CONFIG
# ------------------------------------------------------------------
MODEL_PATH = "weapons_yolo.pt"
METRICS_DIR = "Metrics"

os.makedirs(METRICS_DIR, exist_ok=True)

print("[INFO] Cargando modelo...")
model = YOLO(MODEL_PATH)

# ------------------------------------------------------------------
# EVALUACIÓN
# ------------------------------------------------------------------

def evaluate_model():
    print("[INFO] Descargando dataset...")
    ds = load_dataset("cerealkiller2527/pick_knife")
    images = ds["train"]

    y_true = []
    y_pred = []

    total = len(images)
    print(f"[INFO] Total de imágenes: {total}")

    for i, item in enumerate(images):
        img = item["image"]
        label = int(item["label"])  # 0 = no knife, 1 = knife

        results = model.predict(img, verbose=False)[0]

        detected_knife = False
        for box in results.boxes:
            cls = int(box.cls[0])
            name = results.names[cls].lower()
            if "knife" in name:
                detected_knife = True
                break

        y_true.append(label)
        y_pred.append(1 if detected_knife else 0)

        if i % 50 == 0:
            print(f"[INFO] Procesadas {i}/{total} imágenes...")

    # ------------------------------------------------------------------
    # Métricas
    # ------------------------------------------------------------------

    precision = precision_score(y_true, y_pred)
    recall = recall_score(y_true, y_pred)
    f1 = f1_score(y_true, y_pred)
    cm = confusion_matrix(y_true, y_pred)

    print("\n===== RESULTADOS =====")
    print(f"Precisión: {precision:.4f}")
    print(f"Sensibilidad (Recall): {recall:.4f}")
    print(f"F1 Score: {f1:.4f}")
    print("Confusion Matrix:")
    print(cm)

    # Guardamos métricas en TXT
    with open(os.path.join(METRICS_DIR, "metrics.txt"), "w") as f:
        f.write(f"Precisión: {precision:.4f}\n")
        f.write(f"Recall: {recall:.4f}\n")
        f.write(f"F1 Score: {f1:.4f}\n")
        f.write(f"Matriz de confusión:\n{cm}\n")

    # ------------------------------------------------------------------
    # GRÁFICAS
    # ------------------------------------------------------------------

    # Confusion matrix plot
    plt.figure(figsize=(5, 4))
    plt.imshow(cm, cmap="Blues")
    plt.title("Matriz de Confusión")
    plt.xlabel("Predicción")
    plt.ylabel("Valor Real")
    plt.colorbar()
    plt.savefig(os.path.join(METRICS_DIR, "confusion_matrix.png"))
    plt.close()

    # Precision-Recall (simple)
    plt.figure()
    plt.bar(["Precisión", "Recall", "F1"], [precision, recall, f1])
    plt.ylim(0, 1)
    plt.title("Reporte de Métricas")
    plt.savefig(os.path.join(METRICS_DIR, "metrics_barplot.png"))
    plt.close()

    print(f"\n[INFO] Gráficas guardadas en: {METRICS_DIR}/")
    print("[INFO] Evaluación completa.")

# ------------------------------------------------------------------

if __name__ == "__main__":
    evaluate_model()