import streamlit as st
import tensorflow_datasets as tfds
import tensorflow as tf
import numpy as np
import time
import tensorflow.keras as keras
from tensorflow.keras.applications import VGG16
from tensorflow.keras.layers import Dense, Flatten
from tensorflow.keras.models import Model, load_model
from datasets import load_dataset
import matplotlib.pyplot as plt
from sklearn.metrics import confusion_matrix, classification_report
import seaborn as sns
from huggingface_hub import HfApi
import os

# 📌 Percorso della cache
os.environ["HF_HOME"] = "/app/.cache"
os.environ["HF_DATASETS_CACHE"] = "/app/.cache"
HF_TOKEN = os.getenv("HF_TOKEN")

# 📌 Autenticazione Hugging Face
if HF_TOKEN:
    api = HfApi()
    user_info = api.whoami(HF_TOKEN)
    st.write(f"✅ Autenticato come {user_info.get('name', 'Utente sconosciuto')}")
else:
    st.warning("⚠️ Nessun token API trovato! Verifica il Secret nello Space.")

# 📌 Caricamento del dataset
st.write("🔄 Caricamento di 300 immagini da `tiny-imagenet`...")
dataset = load_dataset("zh-plus/tiny-imagenet", split="train")

image_list = []
label_list = []

for i, sample in enumerate(dataset):
    if i >= 300:  # Prende solo 300 immagini
        break
    image = tf.image.resize(sample["image"], (64, 64)) / 255.0  # Normalizzazione
    image_list.append(image.numpy())
    label_list.append(np.array(sample["label"]))  

X_train = np.array(image_list)
y_train = np.array(label_list)

st.write(f"✅ Scaricate e preprocessate {len(X_train)} immagini da `tiny-imagenet/64x64`!")

# 📌 Caricamento del modello
if os.path.exists("Silva.h5"):
    model = load_model("Silva.h5")
    st.write("✅ Modello `Silva.h5` caricato, nessun nuovo training necessario!")
else:
    st.write("🚀 Training in corso perché `Silva.h5` non esiste...")
    
    base_model = VGG16(weights="imagenet", include_top=False, input_shape=(64, 64, 3))
    for layer in base_model.layers:
        layer.trainable = False

    x = Flatten()(base_model.output)
    x = Dense(256, activation="relu")(x)
    x = Dense(128, activation="relu")(x)
    output = Dense(len(set(y_train)), activation="softmax")(x)

    model = Model(inputs=base_model.input, outputs=output)
    model.compile(optimizer="adam", loss="sparse_categorical_crossentropy", metrics=["accuracy"])

    # 📌 Training con barra di progresso
    progress_bar = st.progress(0)
    status_text = st.empty()
    start_time = time.time()

    history = model.fit(X_train, y_train, epochs=10)

    st.write("✅ Addestramento completato!")

    # 📌 Salvare il modello
    model.save("Silva.h5")
    st.write("✅ Modello salvato come `Silva.h5`!")

# 📌 Calcolo delle metriche
y_pred = np.argmax(model.predict(X_train), axis=1)
accuracy = np.mean(y_pred == y_train)
rmse = np.sqrt(np.mean((y_pred - y_train) ** 2))
report = classification_report(y_train, y_pred, output_dict=True)

recall = report["weighted avg"]["recall"]
precision = report["weighted avg"]["precision"]
f1_score = report["weighted avg"]["f1-score"]

st.write(f"📊 **Accuracy:** {accuracy:.4f}")
st.write(f"📊 **RMSE:** {rmse:.4f}")
st.write(f"📊 **Precision:** {precision:.4f}")
st.write(f"📊 **Recall:** {recall:.4f}")
st.write(f"📊 **F1-Score:** {f1_score:.4f}")

# 📌 Bottone per generare la matrice di confusione
if st.button("🔎 Genera matrice di confusione"):
    conf_matrix = confusion_matrix(y_train, y_pred)
    fig, ax = plt.subplots(figsize=(10, 7))
    sns.heatmap(conf_matrix, annot=True, cmap="Blues", fmt="d", ax=ax)
    st.pyplot(fig)
    st.write("✅ Matrice di confusione generata!")

# 📌 Grafico per Loss e Accuracy
fig, ax = plt.subplots(1, 2, figsize=(12, 5))
ax[0].plot(history.history["loss"], label="Loss")
ax[1].plot(history.history["accuracy"], label="Accuracy")
ax[0].set_title("Loss durante il training")
ax[1].set_title("Accuracy durante il training")
ax[0].legend()
ax[1].legend()
st.pyplot(fig)

# 📌 Bottone per scaricare il modello
if os.path.exists("Silva.h5"):
    with open("Silva.h5", "rb") as f:
        st.download_button(
            label="📥 Scarica il modello Silva.h5",
            data=f,
            file_name="Silva.h5",
            mime="application/octet-stream"
        )

# 📌 Bottone per caricare il modello su Hugging Face
def upload_model():
    api.upload_file(
        path_or_fileobj="Silva.h5",
        path_in_repo="Silva.h5",
        repo_id="scontess/trainigVVG16",
        repo_type="space"
    )
    st.success("✅ Modello 'Silva.h5' caricato su Hugging Face!")

st.write("📥 Carica il modello Silva su Hugging Face")
if st.button("🚀 Carica Silva su Model Store"):
    upload_model()