import gradio as gr
import tensorflow as tf  # Use standard tensorflow
import numpy as np
from PIL import Image
import os
from huggingface_hub import InferenceClient

# 1. Load the TFLite Model using tensorflow's interpreter
interpreter = tf.lite.Interpreter(model_path="skin_model.tflite")
interpreter.allocate_tensors()
input_details = interpreter.get_input_details()
output_details = interpreter.get_output_details()

# 2. Load Labels
with open("labels.txt", "r") as f:
    labels = [line.strip() for line in f.readlines()]

# 3. Setup AI Assistant Client
client = InferenceClient("mistralai/Mistral-7B-Instruct-v0.3", token=os.getenv("HF_TOKEN"))

def predict_and_advise(image):
    # Preprocess image
    input_shape = input_details[0]['shape']
    img = Image.fromarray(image).resize((input_shape[1], input_shape[2]))
    input_data = np.expand_dims(np.array(img, dtype=np.float32), axis=0)
    
    # Run inference
    interpreter.set_tensor(input_details[0]['index'], input_data)
    interpreter.invoke()
    output_data = interpreter.get_tensor(output_details[0]['index'])[0]
    
    top_index = np.argmax(output_data)
    disease_name = labels[top_index]
    confidence = float(output_data[top_index])

    prompt = f"A skin analysis AI has detected {disease_name}. Briefly explain what this is and provide 3 general care tips. End by saying: 'This is not a medical diagnosis; please see a dermatologist.'"
    
    try:
        advice = client.text_generation(prompt, max_new_tokens=250)
    except:
        advice = "Could not fetch advice. Please consult a dermatologist."

    return {
        "condition": disease_name,
        "confidence": f"{confidence*100:.2f}%",
        "assistant_advice": advice
    }

demo = gr.Interface(
    fn=predict_and_advise,
    inputs=gr.Image(),
    outputs=gr.JSON(),
)

demo.launch()