import gradio as gr
import torch
import open_clip

# --- CONFIGURAZIONE DISPOSITIVO ---
# Forziamo esplicitamente l'uso della CPU ignorando eventuali GPU presenti
device = torch.device("cpu")

# --- CARICAMENTO MODELLO ---
print("Caricamento del modello BioCLIP in corso... l'operazione richiede tempo al primo avvio.")
# Scarica i pesi da HuggingFace. L'uso di device=device assicura che i tensori non cerchino CUDA.
model, _, preprocess = open_clip.create_model_and_transforms('hf-hub:imageomics/bioclip-2.5-vith14', device=device)
tokenizer = open_clip.get_tokenizer('hf-hub:imageomics/bioclip-2.5-vith14')

# Mettiamo il modello in modalità valutazione (inferenza) per disabilitare dropout ecc.
model.eval()

# --- PREPARAZIONE DELLE CLASSI E DELLE FEATURE TESTUALI ---
classi_piante = [
    "Aloe vera", 
    "Echeveria agavoides", 
    "Crassula ovata", 
    "Lithops aucampiae", 
    "Schlumbergera truncata", 
    "Euphorbia obesa", 
    "Haworthia fasciata", 
    "Agave americana", 
    "Gymnocalycium mihanovichii"
]

# Creiamo i prompt nel formato richiesto
text_prompts = [f"a photo of a {pianta}" for pianta in classi_piante]

print("Calcolo delle feature di testo per le classi predefinite...")
# Calcoliamo le feature di testo una sola volta all'avvio usando no_grad
with torch.no_grad():
    text_tokens = tokenizer(text_prompts).to(device)
    text_features = model.encode_text(text_tokens)
    # Normalizzazione necessaria per la cosine similarity
    text_features /= text_features.norm(dim=-1, keepdim=True)


# --- FUNZIONE DI INFERENZA ---
def predict(image):
    if image is None:
        return {}

    # Preprocessiamo l'immagine e la portiamo su CPU aggiungendo la dimensione batch
    image_input = preprocess(image).unsqueeze(0).to(device)

    # L'uso di torch.no_grad() ottimizza l'inferenza CPU risparmiando memoria e calcoli inutili
    with torch.no_grad():
        # Estraiamo le feature dell'immagine
        image_features = model.encode_image(image_input)
        
        # Normalizziamo le feature dell'immagine
        image_features /= image_features.norm(dim=-1, keepdim=True)
        
        # Calcoliamo la Cosine Similarity
        # Moltiplicazione matriciale (dot product su feature normalizzate = cosine similarity)
        # Scaliamo per 100.0 (approssimazione del logit_scale di CLIP) per una softmax più decisa
        similarity = (100.0 * image_features @ text_features.T)
        
        # Applichiamo la Softmax per ottenere probabilità in percentuale (0-1)
        text_probs = similarity.softmax(dim=-1)
        
    # Estraiamo i valori dal tensore alla lista
    probs = text_probs[0].tolist()
    
    # Costruiamo il dizionario da restituire a Gradio: {Nome Pianta: Probabilità}
    risultati = {classi_piante[i]: probs[i] for i in range(len(classi_piante))}
    
    return risultati


# --- INTERFACCIA GRADIO ---
# Stili CSS aggiuntivi per ingrandire pulsanti e font per l'accessibilità
custom_css = """
.gradio-container { font-family: 'Segoe UI', Tahoma, Geneva, Verdana, sans-serif; }
button.primary { font-size: 1.5rem !important; padding: 15px !important; }
.label { font-size: 1.2rem !important; }
"""

demo = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil", label="Carica la foto della pianta qui (Usa la fotocamera o la galleria)"),
    outputs=gr.Label(num_top_classes=3, label="Previsioni (Top 3)"),
    title="🌵 Riconoscimento Piante Grasse per Papà 🌵",
    description="Carica la foto di una pianta grassa e il sistema proverà a indovinare la specie esatta. N.B. L'analisi richiede qualche secondo perché utilizza esclusivamente il processore (CPU).",
    flagging_mode="never"
)

if __name__ == "__main__":
    # Avviamo il server in ascolto su tutte le interfacce per essere accessibile da Docker
    demo.launch(server_name="0.0.0.0", server_port=7860, css=custom_css)