import gradio as gr
import cv2
import numpy as np
import pandas as pd
from tensorflow.keras.models import load_model
from keras.utils import to_categorical
import pickle

class ButterflyClassifier:
    def __init__(self, model_path, encoder_path, data_path):
        self.loaded_model = load_model(model_path)
        self.encoder = self.load_encoder(encoder_path)
        self.data = pd.read_csv(data_path)

    def load_encoder(self, encoder_path):
        with open(encoder_path, 'rb') as f:
            return pickle.load(f)

    def predict(self, image):
        img = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        img = cv2.resize(img, (224, 224))
        img = img.astype('float32') / 255.0
        img = np.expand_dims(img, axis=0)
        
        out = self.loaded_model.predict(img)
        y_pred = np.argmax(out, axis=1)
        
        predictions_one_hot = to_categorical(y_pred, num_classes=13)
        original_labels = self.encoder.inverse_transform(predictions_one_hot)
        return original_labels[0, 0]

    def get_top_predictions(self, label):
        df = self.data[self.data['LABEL'] == label.upper()]
        total_interactions = df["Interactions"].sum()
        df["percentage"] = (df["Interactions"] / total_interactions) * 100
        df_sorted = df.sort_values(by='percentage', ascending=False)
        return df_sorted.head(10)

butterfly_classifier = ButterflyClassifier("/Users/fahimafridi/downloads/flagged/Butterfly_Pollination/ButterflyModel.h5", "/Users/fahimafridi/downloads/flagged/Butterfly_Pollination/onehot_encoder.pkl", "/Users/fahimafridi/downloads/flagged/Butterfly_Pollination/Butterfly.csv")

def predict_butterfly(image):
    label = butterfly_classifier.predict(image)
    top_predictions = butterfly_classifier.get_top_predictions(label)
    return top_predictions

# Interface using Gradio
inputs = gr.Image()
outputs = gr.Dataframe(headers=['Species', 'Interactions', 'Percentage'])
gr.Interface(fn=predict_butterfly, inputs=inputs, outputs=outputs).launch()