import gradio as gr
import numpy as np
from PIL import Image, ImageOps
import tensorflow as tf
from huggingface_hub import InferenceClient

# Load the pre-trained Keras model using TensorFlow's Keras
model = tf.keras.models.load_model("keras_model.h5", compile=False)

# Load the class labels
with open("labels.txt", "r") as file:
    class_names = [line.strip() for line in file.readlines()]

# Initialize the HuggingFace client for the chatbot
client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")

# Sample images for the emotion detection
examples = [
    ["https://firebasestorage.googleapis.com/v0/b/hisia-4b65b.appspot.com/o/a-captivating-ukiyo-e-inspired-poster-featuring-a--wTg7L-f2Tfiy6K8w6aWnKA-KbGU9GSKSDGBbbxrCO65Mg.jpeg?alt=media&token=64590de9-e265-44ac-a766-aeecd455ed5d"],
    ["https://firebasestorage.googleapis.com/v0/b/hisia-4b65b.appspot.com/o/poster-ai-themed-kenyan-female-silhoutte-written-l-PMIXpNWGQ8KaNNetQRVJuQ-B1TteyL-S5OTPZFXvfGybg.jpeg?alt=media&token=fc10f96d-403e-4f75-bd9c-810e0da36867"],
    ["https://firebasestorage.googleapis.com/v0/b/hisia-4b65b.appspot.com/o/poster-ai-themed-kenyan-male-silhoutte-written-log-z3fqBD5bQOOj6uqGd_iXLQ-4aBfNy0ZTgmLlTsZh1dzIA.jpeg?alt=media&token=f218f160-d38e-482f-97a9-5442c2f251a7"]
]

def classify_image(img):
    """Classify the image and return the detected emotion and confidence score."""
    try:
        size = (224, 224)
        image = ImageOps.fit(img, size, Image.Resampling.LANCZOS)
        image_array = np.asarray(image)
        normalized_image_array = (image_array.astype(np.float32) / 127.5) - 1
        data = normalized_image_array.reshape((1, 224, 224, 3))

        # Perform prediction using the model
        prediction = model.predict(data)
        index = np.argmax(prediction)
        class_name = class_names[index]
        confidence_score = prediction[0][index]

        return class_name, confidence_score
    except Exception as e:
        print(f"Error in classify_image: {e}")
        return "Error", 0

def respond(message, history, system_message, max_tokens, temperature, top_p):
    """Generate a response from the chatbot based on the input message and conversation history."""
    try:
        messages = [{"role": "system", "content": system_message}]
        for user_message, assistant_message in history:
            if user_message:
                messages.append({"role": "user", "content": user_message})
            if assistant_message:
                messages.append({"role": "assistant", "content": assistant_message})
        messages.append({"role": "user", "content": message})

        response = ""
        for response_message in client.chat_completion(
            messages,
            max_tokens=max_tokens,
            stream=True,
            temperature=temperature,
            top_p=top_p,
        ):
            token = response_message.choices[0].delta.content
            response += token

        print(f"API Response: {response}")  # Debugging: Print the API response
        return response
    except Exception as e:
        print(f"Error in respond: {e}")
        return "Error generating response"

# Define the custom CSS for styling the interface and hiding the footer
custom_css = """
body {
    font-family: 'Arial', sans-serif;
    background-color: #f4e9e0;
    color: #2e2e2e;
}
.gradio-container {
    border-radius: 12px;
    padding: 20px;
    background: linear-gradient(135deg, #f5b8b8, #a0d6a2);
    box-shadow: 0px 4px 15px rgba(0, 0, 0, 0.2);
}
.gradio-container h1 {
    font-family: 'Arial', sans-serif;
    font-size: 2.2em;
    text-align: center;
    color: #1c1c1c;
    margin-bottom: 20px;
}
.gradio-container p {
    font-size: 1em;
    text-align: center;
    color: #4a4a4a;
}
.gradio-button {
    background-color: #d55a5a;
    border: none;
    color: white;
    padding: 12px 24px;
    font-size: 1.1em;
    cursor: pointer;
    border-radius: 8px;
    transition: background-color 0.2s ease;
}
.gradio-button:hover {
    background-color: #b93e3e;
}
#output-container {
    border-radius: 12px;
    background-color: #ffffff;
    padding: 20px;
    color: #2e2e2e;
    box-shadow: 0px 4px 10px rgba(0, 0, 0, 0.2);
}
#output-container h3 {
    font-family: 'Arial', sans-serif;
    font-size: 1.4em;
    color: #1c1c1c;
}
.gr-examples {
    text-align: center;
}
.gr-example-img {
    width: 120px;
    border-radius: 8px;
    margin: 5px;
    box-shadow: 0px 4px 10px rgba(0, 0, 0, 0.2);
}
footer {
    display: none !important; /* Hides the footer */
}
"""

def emotion_detection_interface():
    """Create and return the Gradio interface with sliders for AI response settings and a single view for emotion detection and recommendations."""
    with gr.Blocks(css=custom_css) as demo:
        gr.Markdown("### HISIA: Emotion Detector and Therapeutic Recommendations")

        with gr.Row():
            with gr.Column(scale=1, min_width=200):
                image_input = gr.Image(type="pil", label="Upload an Image", elem_id="emotion-image")
                submit_button = gr.Button("Classify Image", elem_id="classify-button")

            with gr.Column(scale=2, min_width=300):
                emotion_output = gr.JSON(label="Emotion Detection Result", elem_id="output-container")
                ai_response_output = gr.Textbox(label="AI Recommendations", elem_id="output-container", lines=5)

        # Add sample images
        gr.Examples(examples, inputs=image_input)

        # Add sliders for adjusting AI response settings
        with gr.Row():
            max_tokens_slider = gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="VERBOSENESS")
            temperature_slider = gr.Slider(minimum=0.1, maximum=3.0, value=0.7, step=0.1, label="CREATIVITY")
            top_p_slider = gr.Slider(minimum=0.1, maximum=1.0, value=0.95, step=0.05, label="BROADNESS")

        def process_image(image, max_tokens, temperature, top_p):
            """Process the image and generate emotion detection result and AI recommendations."""
            class_name, confidence_score = classify_image(image)
            emotion_result = {"Detected Emotion": class_name, "Confidence Score": f"{confidence_score:.2f}"}
            if class_name != "Error":
                # Generate AI recommendation based on detected emotion
                recommendation = respond(
                    class_name,
                    history=[],
                    system_message="You are a psychologist that provides therapeutic recommendations based on emotions. Always address the clients in the second pronouns person like your, you, etc",
                    max_tokens=int(max_tokens),
                    temperature=float(temperature),
                    top_p=float(top_p),
                )
                return emotion_result, recommendation
            else:
                return {"Detected Emotion": "Error", "Confidence Score": "0"}, "Error generating response"

        submit_button.click(
            process_image,
            inputs=[image_input, max_tokens_slider, temperature_slider, top_p_slider],
            outputs=[emotion_output, ai_response_output]
        )

    return demo

# Launch the combined interface
if __name__ == "__main__":
    emotion_detection_interface().launch()