import gradio as gr
import numpy as np
import cv2 as cv
import requests
import time
import os

host = os.environ.get("host")
code = os.environ.get("code")
model_llm = os.environ.get("model")
content = os.environ.get("content")
state = os.environ.get("state")
data = None
model = None
image = None
prediction = None
labels = None

print('START')
np.set_printoptions(suppress=True)

data = np.ndarray(shape=(1, 224, 224, 3), dtype=np.float32)

with open("labels.txt", "r") as file:
    labels = file.read().splitlines()

def classify(Textbox, Textbox2, Textbox3):
    if Textbox3 == code:
        output = []
        image_data = np.array("Image")
        image_data = cv.resize(image_data, (224, 224))
        image_array = np.asarray(image_data)
        normalized_image_array = (image_array.astype(np.float32) / 127.0) - 1
        data[0] = normalized_image_array
    
        import tensorflow as tf
        model = tf.keras.models.load_model('keras_model.h5')
    
        prediction = model.predict(data)
        
        max_label_index = None
        max_prediction_value = -1

        print('Prediction')
        

        Textbox2 = Textbox2.replace("[", "").replace("]", "").replace("'", "")
        Textbox2 = Textbox2.split(",")
        Textbox2_edited = [x.strip() for x in Textbox2]
        Textbox2_edited = list(Textbox2_edited)
        Textbox2_edited.append(Textbox)
        messages = [
            {"role": "system", "content": content},
        ]
        for i in Textbox2_edited:
            messages.append(
                {"role": "user", "content": i}
            )
    
        for i, label in enumerate(labels):
            prediction_value = float(prediction[0][i])
            rounded_value = round(prediction_value, 2)
            print(f'{label}: {rounded_value}')
    
            if prediction_value > max_prediction_value:
                max_label_index = i
                max_prediction_value = prediction_value 
    
        if max_label_index is not None:
            max_label = labels[max_label_index].split(' ', 1)[1]
            print(f'Maximum Prediction: {max_label} with a value of {round(max_prediction_value, 2)}')
    
            time.sleep(1)
            print("\nWays to dispose of this waste: " + max_label)
            messages.append(
                {"role": "user", "content": Textbox},
            )
    
            response = requests.post(host, json={
                "messages": messages,
                "model": model_llm,
                "temperature": 0.5,
                "presence_penalty": 0,
                "frequency_penalty": 0,
                "top_p": 1
            }).json()

            reply = response["choices"][0]["message"]["content"]
            messages.append({"role": "assistant", "content": reply})

            output.append({"type": max_label, "prediction_value": rounded_value, "content": reply})
        
        return output
    else:
        return "Unauthorized"

user_inputs = [
    gr.inputs.Textbox(label="Textbox",type="text"),
    gr.inputs.Textbox(label="Textbox2",type="text"),
    gr.inputs.Textbox(label="Textbox3",type="password")
]

iface = gr.Interface(
    fn=classify,
    inputs = user_inputs,
    outputs=gr.outputs.JSON(),
    title="Classifier",
)
iface.launch()