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
client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")

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

# Load class labels for environmental themes
with open("/content/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. Provide expert recommendations 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_recommendations(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 highest and 1 being the least ideal, {alias} rated the following:\n"
        f"Comfort and Well-being: {comfort}\n"
        f"Social Interaction: {social_interaction}\n"
        f"Environmental Stressors: {stressors}\n"
        f"Privacy and Personal Space: {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 personalized recommendations for {alias}, addressing them directly."
    )

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

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

def analyze_environmental_concerns(comfort, social_interaction, stressors, privacy, open_question, alias, img):
    # Classify the image
    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 Ukiyo-e 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=['#F08080', '#90EE90', '#000000', '#FFD700'])  # Light red, light green, black, gold
    
    # Improve graph display
    ax.set_ylabel('Score', fontsize=14, color='#333333')
    ax.set_title(f'Environmental Psychology Concerns 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 recommendations using the model, passing the image analysis results
    recommendations = generate_recommendations(comfort, social_interaction, stressors, privacy, open_question, image_info, alias)
    
    return fig, recommendations

# Custom CSS for Ukiyo-e theme
custom_css = """
body {
    font-family: 'Garamond', serif;
    background-color: #FAF3F3; /* Light background color */
}
.container {
    border: 2px solid #F08080; /* Light red border */
    border-radius: 10px;
    padding: 20px;
    background-color: #FFFFFF; /* White background for content */
}
h1 {
    color: #F08080; /* Light red */
}
h2, h3, h4 {
    color: #333333; /* Dark grey for headings */
}
.gradio-slider .slider .slider-handle {
    background-color: #90EE90; /* Light green slider handle */
}
.gradio-slider .slider .slider-track {
    background-color: #F08080; /* Light red track */
}
.gradio-textbox input, .gradio-textbox textarea {
    border: 2px solid #000000; /* Black border for inputs */
}
.gradio-button {
    background-color: #FFD700; /* Gold background for buttons */
    color: #000000; /* Black text for buttons */
}
"""

# Create the Gradio interface with custom CSS
inputs = [
    gr.Slider(minimum=1, maximum=5, step=1, label="How would you rate Comfort and Well-being?"),
    gr.Slider(minimum=1, maximum=5, step=1, label="How would you rate Social Interaction?"),
    gr.Slider(minimum=1, maximum=5, step=1, label="How would you rate Environmental Stressors?"),
    gr.Slider(minimum=1, maximum=5, step=1, label="How would you rate Privacy and Personal Space?"),
    gr.Textbox(placeholder="Describe any additional concerns or suggestions you have.", label="Open-ended Question", lines=3),
    gr.Textbox(placeholder="Enter a single-word alias (e.g., anon).", label="Client Alias", lines=1),  # New input for alias
    gr.Image(type="pil", label="Upload an Image for Theme Detection")  # New input for image
]

outputs = [
    gr.Plot(label="Concerns Graph"),
    gr.Textbox(label="Recommendations", lines=5)
]

gr.Interface(
    fn=analyze_environmental_concerns,
    inputs=inputs,
    outputs=outputs,
    title="Mazingira: Environmental Psychology Concerns Analyzer",
    description="Input your environmental psychology concerns to receive personalized recommendations and a visual graph, along with image-based analysis.",
    css=custom_css  # Apply custom CSS
).launch()