import gradio as gr
import matplotlib.pyplot as plt
from matplotlib.ticker import MaxNLocator
from huggingface_hub import InferenceClient
from keras.models import load_model
from PIL import Image, ImageOps
import numpy as np

# Initialize Hugging Face Inference Client for climate change and design recommendations
client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")

# Load the pre-trained Keras model for Mazingira 254 (Environmental Theme Detection)
model = load_model("keras_model.h5", compile=False)

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

# Default parameters for the model
DEFAULT_MAX_TOKENS = 1000
DEFAULT_TEMPERATURE = 0.7
DEFAULT_TOP_P = 0.95
DEFAULT_SYSTEM_MESSAGE = "You are an expert in environmental psychology and sustainable design. Provide innovative environmental design ideas to improve the climate and curb climate change, addressing the user directly."
DEFAULT_ALIAS = "anon"  # Default alias

def classify_image(img):
    # Prepare the image for prediction
    image = ImageOps.fit(img, (224, 224), 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))

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

    return {
        "Detected Theme": class_name,
        "Confidence Score": f"{confidence_score:.2f}"
    }

def generate_design_ideas(comfort, social_interaction, stressors, privacy, open_question, image_info, alias=DEFAULT_ALIAS, max_tokens=DEFAULT_MAX_TOKENS, temperature=DEFAULT_TEMPERATURE, top_p=DEFAULT_TOP_P, system_message=DEFAULT_SYSTEM_MESSAGE):
    # Construct the input message for the model with context
    message = (
        f"{system_message}\n"
        f"On a scale of 1-5, with 5 being the most favorable and 1 being the least ideal, {alias} rated the following:\n"
        f"Comfort and Well-being: {comfort}\n"
        f"Quality of Social Interaction: {social_interaction}\n"
        f"Absence of overwhelming Environmental Stressors: {stressors}\n"
        f"Quality of Personal Spaces: {privacy}\n"
        f"Open-ended Question: {open_question}\n"
        f"Detected Image Theme: {image_info['Detected Theme']}\n"
        f"Confidence Score: {image_info['Confidence Score']}\n"
        f"Please provide innovative environmental psychology design ideas to improve climate and curb climate change, addressing wellness and sustainability for {alias} directly."
    )

    # Generate design ideas using the Hugging Face model
    response = client.chat_completion(
        [{"role": "user", "content": message}],
        max_tokens=max_tokens,
        temperature=temperature,
        top_p=top_p
    )

    design_ideas = response.choices[0].message['content']
    # Convert the recommendations to address the user in first person
    design_ideas = design_ideas.replace("You should", "I recommend that you")
    return design_ideas

def analyze_environmental_design(comfort, social_interaction, stressors, privacy, open_question, alias, img):
    # Classify the image for environmental themes
    image_info = classify_image(img)
    
    # Use default alias if none is provided
    alias = alias or DEFAULT_ALIAS

    # Generate a bar graph for the input scores with creative theme colors
    fig, ax = plt.subplots(figsize=(10, 6))  # Increased size for better visibility
    categories = ["Comfort and Well-being", "Social Interaction", "Environmental Stressors", "Privacy and Personal Space"]
    values = [comfort, social_interaction, stressors, privacy]
    
    bars = ax.bar(categories, values, color=['#4CAF50', '#FFC107', '#2196F3', '#FF5722'])  # Green, Amber, Blue, Red
    
    # Improve graph display
    ax.set_ylabel('Score', fontsize=14, color='#333333')
    ax.set_title(f'Environmental Design Assessment for {alias}', fontsize=16, color='#333333')
    ax.yaxis.set_major_locator(MaxNLocator(integer=True))
    ax.tick_params(axis='y', colors='#333333')
    ax.tick_params(axis='x', colors='#333333')
    
    # Add value labels on the bars
    for bar in bars:
        yval = bar.get_height()
        ax.text(bar.get_x() + bar.get_width()/2, yval, int(yval), va='bottom', ha='center', color='black', fontsize=12)

    # Generate design ideas using the model, passing the image analysis results
    design_ideas = generate_design_ideas(comfort, social_interaction, stressors, privacy, open_question, image_info, alias)

    return fig, design_ideas

# Custom CSS for modern design with footer hidden
custom_css = """
body {
    font-family: 'Arial', sans-serif;
    background-color: #e0f7fa;
    color: #00695c;
}

.gradio-container {
    border-radius: 10px;
    padding: 20px;
    background: linear-gradient(135deg, #a5d6a7, #1b5e20);
    box-shadow: 0px 4px 15px rgba(0, 0, 0, 0.2);
}

.gradio-container h1 {
    font-family: 'Arial', sans-serif;
    font-size: 2.5em;
    text-align: center;
    color: #ffffff;
}

.gradio-button {
    background-color: #00796b;
    border: none;
    color: white;
    padding: 10px 20px;
    font-size: 1em;
    cursor: pointer;
    border-radius: 5px;
}

.gradio-button:hover {
    background-color: #004d40;
}

/* Hide Gradio footer */
footer {
    display: none !important;
}
"""

# Disclaimer text to include in the description
disclaimer_text = """
**Disclaimer**:
1. This solution was built in only 5 hours with a limited dataset during the training of the CNN, so it may have some biases.
2. The solution is running on low compute resources, so we recommend using it when there is low traffic to improve performance.
"""

# Create the Gradio interface with custom CSS
inputs = [
    gr.Slider(minimum=1, maximum=5, step=1, label="On a scale of 1-5, with 5 being the most favorable and 1 being the least ideal. How would you rate your feelings of Comfort and Well-being?"),
    gr.Slider(minimum=1, maximum=5, step=1, label="On a scale of 1-5, with 5 being the most favorable and 1 being the least ideal. How would you rate the quality of your Social Interaction?"),
    gr.Slider(minimum=1, maximum=5, step=1, label="On a scale of 1-5, with 5 being the most favorable and 1 being the least ideal. How would you rate the absence of overwhelming Environmental Stressors like noise in your environment?"),
    gr.Slider(minimum=1, maximum=5, step=1, label="On a scale of 1-5, with 5 being the most favorable and 1 being the least ideal. How would you rate the quality of your Personal Spaces?"),
    gr.Textbox(placeholder="Describe any environmental challenges or ideas you have.", label="Open-ended Question", lines=3),
    gr.Textbox(placeholder="Enter your alias (e.g., anon).", label="Client Alias", lines=1),  # New input for alias
    gr.Image(type="pil", label="Upload an Image for Environmental Theme Detection")  # New input for image
]

outputs = [
    gr.Plot(label="Environmental Design Assessment"),
    gr.Textbox(label="Creative Design Ideas", lines=5)
]

gr.Interface(
    fn=analyze_environmental_design,
    inputs=inputs,
    outputs=outputs,
    title="MAZINGIRA 254: Climate-Smart Environmental Psychology Design Ideas",
    description=f"Input your environmental concerns and ideas in the form of images and personal experiences to receive creative design recommendations and a visual assessment of climate change impact.\n\n{disclaimer_text}",
    css=custom_css  # Apply custom CSS
).launch()