import torch
import cv2
from PIL import Image
import numpy as np
import matplotlib.pyplot as plt
from transformers import pipeline
import gradio as gr
from sklearn.cluster import KMeans
from colorsys import rgb_to_hsv
from huggingface_hub import InferenceClient  # Import InferenceClient for API calls

# Initialize the emotion detection pipeline for text (if any text is included in assets)
emotion_classifier = pipeline("text-classification", model="j-hartmann/emotion-english-distilroberta-base", return_all_scores=True)

# Initialize the Hugging Face API client with your Space model ID
client = InferenceClient("nehapasricha94/LLaVA-image-analysis")

# Function to analyze colors in an image
def analyze_colors(image):
    try:
        # Ensure the image is in RGB format
        if image.mode != "RGB":
            image = image.convert("RGB")

        # Resize the image for faster processing
        image = image.resize((150, 150))

        # Convert to numpy array
        img_array = np.array(image)

        # Reshape image to be a list of pixels
        pixels = img_array.reshape((-1, 3))

        kmeans = KMeans(n_clusters=5, random_state=0)
        kmeans.fit(pixels)
        dominant_colors = kmeans.cluster_centers_

        # Plot the colors for visualization
        plt.figure(figsize=(8, 6))
        plt.imshow([dominant_colors.astype(int)])
        plt.axis('off')
        plt.show()

        return dominant_colors

    except Exception as e:
        print(f"Error in analyze_colors: {e}")
        return None


# Function to analyze emotions based on color (hue, brightness, saturation) and stress with weights
def color_emotion_analysis(dominant_colors):
    try:
        emotions = []
        stress_levels = []
        
        # Weight coefficients for each factor
        brightness_weight = 0.5
        hue_weight = 0.3
        saturation_weight = 0.2

        for color in dominant_colors:
            # Normalize RGB values to 0-1 range
            r, g, b = color / 255.0

            # Convert RGB to HSV
            h, s, v = rgb_to_hsv(r, g, b)  # Hue, Saturation, Value (brightness)

            # Calculate weighted emotion and stress levels
            weighted_brightness = v * brightness_weight
            weighted_hue = h * hue_weight
            weighted_saturation = s * saturation_weight

            # Combine weighted factors
            score = weighted_brightness + weighted_hue + weighted_saturation

            # Analyze emotion and stress based on combined score
            if score < 0.3:  # Lower combined score, less rigid "high stress"
                emotions.append("Sadness")
                stress_levels.append("Moderate-High Stress")
            elif 0.3 <= score < 0.5:
                emotions.append("Neutral")
                stress_levels.append("Moderate Stress")
            elif 0.5 <= score < 0.7:
                emotions.append("Okay")
                stress_levels.append("Low Stress")
            elif 0.7 <= score < 0.85:
                emotions.append("Happiness")
                stress_levels.append("Very Low Stress")
            else:
                emotions.append("Very Happy")
                stress_levels.append("No Stress")
        
        return emotions, stress_levels

    except Exception as e:
        print(f"Error in color_emotion_analysis: {e}")
        return ["Error analyzing emotions"], ["Error analyzing stress levels"]


# Function to analyze patterns and shapes using OpenCV
def analyze_patterns(image):
    try:
        # Convert to grayscale for edge detection
        gray_image = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2GRAY)
        edges = cv2.Canny(gray_image, 100, 200)

        # Calculate the number of edges (chaos metric)
        num_edges = np.sum(edges > 0)

        if num_edges > 10000:  # Arbitrary threshold for "chaos"
            return "Chaotic patterns - possibly distress", 0.8  # Assigning 0.8 stress for chaotic patterns
        else:
            return "Orderly patterns - possibly calm", 0.2  # Assigning 0.2 stress for calm patterns
    except Exception as e:
        print(f"Error in analyze_patterns: {e}")
        return "Error analyzing patterns", 0.5  # Assigning neutral weight for errors


# Function to compute overall results by combining color and pattern analyses
def compute_overall_result(color_emotions, stress_levels, pattern_analysis, pattern_stress):
    try:
        # Assigning weightage to different factors
        color_emotion_weight = 0.5  # 50% for color-based emotions and stress
        pattern_weight = 0.5  # 50% for pattern analysis

        # Determine the most common emotion from colors
        dominant_emotion = max(set(color_emotions), key=color_emotions.count)
        
        # Average stress level from color analysis (converting text stress levels to numeric)
        stress_mapping = {
            "No Stress": 0.1,
            "Very Low Stress": 0.3,
            "Low Stress": 0.5,
            "Moderate Stress": 0.7,
            "Moderate-High Stress": 0.9,
        }
        color_stress_numeric = [stress_mapping[stress] for stress in stress_levels if stress in stress_mapping]
        avg_color_stress = np.mean(color_stress_numeric) if color_stress_numeric else 0.5  # Default to 0.5 if no valid data
        
        # Compute the overall stress by combining color stress and pattern stress
        overall_stress = (avg_color_stress * color_emotion_weight) + (pattern_stress * pattern_weight)
        
        # Determine overall result based on combined factors
        if overall_stress < 0.3:
            overall_emotion = "Calm and Relaxed"
        elif 0.3 <= overall_stress < 0.6:
            overall_emotion = "Neutral Mood"
        elif 0.6 <= overall_stress < 0.8:
            overall_emotion = "Slightly Stressed"
        else:
            overall_emotion = "Highly Stressed"
        
        return f"Overall emotion: {overall_emotion} (Dominant Color Emotion: {dominant_emotion}, Pattern Analysis: {pattern_analysis})"
    
    except Exception as e:
        return f"Error computing overall result: {str(e)}"


# Function to analyze emotions from a given text (if applicable)
def analyze_emotion_from_text(text):
    try:
        # Using the local emotion classifier
        emotion_scores = emotion_classifier(text)
        dominant_emotion = max(emotion_scores, key=lambda x: x['score'])
        return f"Detected emotion from text: {dominant_emotion['label']} with score: {dominant_emotion['score']:.2f}"
    except Exception as e:
        print(f"Error analyzing emotion from text: {e}")
        return "Error analyzing text emotion"


# Main function to process image and analyze emotional expression
def analyze_emotion_from_image(image):
    try:
        # Ensure the input image is a PIL image
        if isinstance(image, np.ndarray):
            image = Image.fromarray(image)  # Convert to PIL Image if it's a NumPy array

        # Analyze colors
        dominant_colors = analyze_colors(image)
        if dominant_colors is None:
            return "Error analyzing colors"
        
        color_emotions, stress_levels = color_emotion_analysis(dominant_colors)

        # Analyze patterns
        pattern_analysis, pattern_stress = analyze_patterns(image)

        # Compute overall result
        overall_result = compute_overall_result(color_emotions, stress_levels, pattern_analysis, pattern_stress)

        return overall_result
    except Exception as e:
        return f"Error processing image: {str(e)}"


# Gradio interface to upload image files and perform analysis
iface = gr.Interface(fn=analyze_emotion_from_image, inputs="image", outputs="text")

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