import streamlit as st
import tensorflow as tf
import numpy as np
from PIL import Image
import pandas as pd
import matplotlib.pyplot as plt

# Load the trained model
model_path = "pokemon-model_transferlearning.keras"
model = tf.keras.models.load_model(model_path)

# Define the core prediction function
def predict_pokemon(image):
    # Preprocess image
    image = image.resize((150, 150))
    image = image.convert('RGB')
    image = np.array(image)
    image = np.expand_dims(image, axis=0)
    
    # Predict
    prediction = model.predict(image)
    
    # Apply softmax to get probabilities for each class
    probabilities = tf.nn.softmax(prediction, axis=1)
    
    # Map probabilities to Pokemon classes
    class_names = ['Abra', 'Charmander', 'Mewtwo']
    probabilities_dict = {pokemon_class: round(float(probability), 2) for pokemon_class, probability in zip(class_names, probabilities.numpy()[0])}
    
    return probabilities_dict

# Streamlit interface
st.title("Pokemon Guesser")

# Upload image
uploaded_image = st.file_uploader("Choose a Pokemon image:", type=["jpg", "png"])

if uploaded_image is not None:
    image = Image.open(uploaded_image)
    st.image(image, caption='Uploaded Image.', use_column_width=True)
    st.write("")
    st.write("Etwas gedult :)")
    
    predictions = predict_pokemon(image)

    # Find the Pokémon with the highest probability
    highest_prob_pokemon = max(predictions.items(), key=lambda item: item[1])

    # Create a DataFrame with only the highest probability Pokémon
    df = pd.DataFrame([highest_prob_pokemon], columns=["Pokemon", "Probability"])

    # Display the DataFrame
    st.write("### Pokémon with the Highest Probability")
    st.dataframe(df)