import gradio as gr
import matplotlib.pyplot as plt
from matplotlib.ticker import MaxNLocator
from huggingface_hub import InferenceClient

# Initialize Hugging Face Inference Client
client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")

# Set default values for model parameters and system message
DEFAULT_MAX_TOKENS = 512
DEFAULT_TEMPERATURE = 0.7
DEFAULT_TOP_P = 0.95
DEFAULT_SYSTEM_MESSAGE = "You are an expert in environmental psychology. Provide expert recommendations."

def generate_recommendations(comfort, social_interaction, stressors, privacy, open_question, 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, the user 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"Provide recommendations for improving the environment based on these ratings.")

    # 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']
    return recommendations

def analyze_environmental_concerns(comfort, social_interaction, stressors, privacy, open_question):
    # Generate a bar graph for the input scores with Ukiyo-e theme colors
    fig, ax = plt.subplots()
    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')
    ax.set_title('Environmental Psychology Concerns')
    ax.yaxis.set_major_locator(MaxNLocator(integer=True))
    
    # 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=10)
    
    # Generate recommendations using the model
    recommendations = generate_recommendations(comfort, social_interaction, stressors, privacy, open_question)
    
    return fig, recommendations

# Create the Gradio interface with Ukiyo-e theme
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)
]

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.",
    theme="default"  # Use default theme as Gradio doesn't support custom themes directly
).launch()