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
import seaborn as sns
from huggingface_hub import HfApi
import os


# 📌 Percorso della cartella dove sarà salvato il dataset se uso TFlow, non serve se setto hf_dataset_cache
#DATA_DIR = "/app" #"/tmp"


# 📌 Autenticazione Hugging Face dal Secret nello Space
HF_TOKEN = os.getenv("HF_TOKEN")
os.environ["HF_HOME"] = "/app/.cache"
os.environ["HF_DATASETS_CACHE"] = "/app/.cache"


if HF_TOKEN:
    api = HfApi()
    user_info = api.whoami(HF_TOKEN)

    if "name" in user_info:
        st.write(f"✅ Autenticato come {user_info['name']}")
    else:
        st.warning("⚠️ Token API non valido! Controlla il Secret nello Space.")
else:
    st.warning("⚠️ Nessun token API trovato! Verifica il Secret nello Space.")

# 📌 Carica solo 300 immagini da `imagenet_resized/64x64`
st.write("🔄 Caricamento di 300 immagini da `tiny-imagenet`...")

# 📌 Recupera il valore della variabile d'ambiente
hf_cache_path = os.getenv("HF_DATASETS_CACHE", "❌ Variabile non impostata!")

# 🔹 Mostra il valore nella UI di Streamlit
st.write(f"📂 Cache dei dataset Hugging Face: {hf_cache_path}")

# 📌 Testa se la cache ha i permessi giusti PRIMA di caricare il dataset
test_file = "/app/.cache/test.txt"
try:
    with open(test_file, "w") as f:
        f.write("Test permessi OK!")
    st.write("✅ La cartella ha i permessi giusti!")
except PermissionError:
    st.error("❌ ERRORE: La cartella /app/.cache non ha permessi di scrittura!")



# 📌 Carica il dataset direttamente da Hugging Face
os.environ["HF_HOME"] = "/tmp/huggingface"
os.environ["HF_DATASETS_CACHE"] = "/tmp/huggingface"
os.makedirs(os.getenv("HF_DATASETS_CACHE"), exist_ok=True)
dataset = load_dataset("zh-plus/tiny-imagenet", split="train")

# 📌 Recupera il primo esempio
sample = dataset[0]
image = sample["image"]
label = sample["label"]

# 📌 Mostra l'immagine e la classe in Streamlit
st.image(image, caption=f"Classe: {label}", use_container_width=True)
st.write(f"🔎 Esempio dal dataset: {sample}")

#Carica il dataset esterno da imagenet PER TENSORFLOW
#imagenet = tfds.load("imagenet_resized/64x64", split="train", as_supervised=True, download=True, data_dir=DATA_DIR)

image_list = []
label_list = []

#for i, (image, label) in enumerate(imagenet.take(300)):  # Prende solo 300 immagini  PER TENSORFLOW 
for i, sample in enumerate(dataset):
    if i >= 300:  # Prende solo 300 immagini
        break

    image = tf.image.resize(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-imagene/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...")

    # Caricare il modello VGG16 pre-addestrato
    base_model = VGG16(weights="imagenet", include_top=False, input_shape=(64, 64, 3))

    # Congelare i livelli convoluzionali
    for layer in base_model.layers:
        layer.trainable = False

    # Aggiungere nuovi livelli Dense
    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)

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

    # 📌 Barra di progresso del training
    progress_bar = st.progress(0)
    status_text = st.empty()
    start_time = time.time()

    # 📌 Addestramento con progress bar
    for epoch in range(10):
        history = model.fit(X_train, y_train, epochs=1)
        progress_bar.progress((epoch + 1) / 10)
        elapsed_time = time.time() - start_time
        status_text.text(f"⏳ Tempo rimanente stimato: {int(elapsed_time / (epoch + 1) * (10 - (epoch + 1)))} secondi")

    st.write("✅ Addestramento completato!")

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

# 📌 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",   #"scontess/Silva",
        repo_type="space"
    )
    st.success("✅ Modello 'Silva' caricato su Hugging Face!")

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