import os
import cv2
import numpy as np
import pandas as pd
import plotly.express as px
import plotly.graph_objects as go
from sklearn.cluster import KMeans
from PIL import Image
import gradio as gr

def rgb_to_hex(r, g, b):
    """Converts RGB integers to a hex string."""
    return f"#{int(r):02x}{int(g):02x}{int(b):02x}"

def analyze_palette(img_numpy, num_colors=5):
    """Clusters pixels using K-Means to extract dominant color palette."""
    try:
        # Resize image to speed up K-Means significantly
        img_small = cv2.resize(img_numpy, (150, 150), interpolation=cv2.INTER_AREA)
        pixels = img_small.reshape(-1, 3)
        
        # Run K-Means
        kmeans = KMeans(n_clusters=num_colors, random_state=42, n_init=10)
        kmeans.fit(pixels)
        
        colors = kmeans.cluster_centers_
        labels = kmeans.labels_
        
        # Calculate percentages
        counts = np.bincount(labels)
        total = len(labels)
        percentages = counts / total
        
        # Sort colors by dominance (percentage)
        sorted_indices = np.argsort(percentages)[::-1]
        colors = colors[sorted_indices]
        percentages = percentages[sorted_indices]
        
        # Compile color data
        palette_data = []
        warm_percentage = 0.0
        cool_percentage = 0.0
        
        for i in range(num_colors):
            r, g, b = colors[i]
            hex_val = rgb_to_hex(r, g, b)
            pct = percentages[i]
            
            # Simple heuristic for color temperature
            # Warm: R > B. Cool: B >= R
            is_warm = r > b
            if is_warm:
                warm_percentage += pct
            else:
                cool_percentage += pct
                
            # Basic visual psychology association maps
            vibes = "Neutral/Balanced"
            if r > 150 and g < 100 and b < 100:
                vibes = "Urgency, Excitement, Passion, or Danger"
            elif b > 150 and r < 100 and g < 150:
                vibes = "Trust, Stability, Calm, or Professionalism"
            elif g > 150 and r < 120 and b < 120:
                vibes = "Nature, Growth, Health, or Balance"
            elif r > 180 and g > 150 and b < 100:
                vibes = "Energy, Optimism, Warmth, or Caution"
            elif r > 120 and g < 80 and b > 120:
                vibes = "Luxury, Creative, Mystery, or Dignity"
            elif r > 200 and g > 200 and b > 200:
                vibes = "Clarity, Minimalism, Openness, or Light"
            elif r < 60 and g < 60 and b < 60:
                vibes = "Authority, Sophistication, Drama, or Mystery"
                
            palette_data.append({
                "Hex": hex_val,
                "RGB": f"({int(r)}, {int(g)}, {int(b)})",
                "Dominance": pct,
                "Vibes & Framing Role": vibes
            })
            
        # Draw Plotly stacked bar chart
        fig = go.Figure()
        for item in palette_data:
            fig.add_trace(go.Bar(
                name=item["Hex"],
                y=["Palette"],
                x=[item["Dominance"]],
                orientation='h',
                marker=dict(color=item["Hex"]),
                hovertemplate=f"Color: {item['Hex']}<br>Dominance: {item['Dominance']:.1%}<br>{item['Vibes & Framing Role']}<extra></extra>"
            ))
            
        fig.update_layout(
            barmode='stack',
            showlegend=False,
            height=120,
            template="plotly_dark",
            plot_bgcolor="#111827",
            paper_bgcolor="#0d0f12",
            margin=dict(l=10, r=10, t=10, b=10),
            xaxis=dict(showticklabels=False, showgrid=False, zeroline=False),
            yaxis=dict(showticklabels=False, showgrid=False, zeroline=False)
        )
        
        # Temp analysis text
        temp_status = f"Visual Temperature: **{'Warm' if warm_percentage > cool_percentage else 'Cool'}** " \
                      f"({warm_percentage:.1%} Warm vs. {cool_percentage:.1%} Cool tones dominant)."
                      
        df_palette = pd.DataFrame(palette_data)
        
        return fig, df_palette, temp_status
    except Exception as e:
        print(f"Palette analysis error: {e}")
        return go.Figure(), pd.DataFrame(), f"Error running palette clustering: {e}"

def analyze_composition(img_numpy):
    """Draws Rule-of-Thirds grid lines and measures edge texture centers of gravity."""
    try:
        h, w, _ = img_numpy.shape
        img_grid = img_numpy.copy()
        
        # Draw Rule-of-Thirds grid lines (high-contrast cyan)
        grid_color = (0, 255, 255) # Cyan
        line_w = max(2, int(w * 0.003))
        
        # Horizontal lines
        h1, h2 = int(h / 3), int(2 * h / 3)
        cv2.line(img_grid, (0, h1), (w, h1), grid_color, line_w)
        cv2.line(img_grid, (0, h2), (w, h2), grid_color, line_w)
        
        # Vertical lines
        w1, w2 = int(w / 3), int(2 * w / 3)
        cv2.line(img_grid, (w1, 0), (w1, h), grid_color, line_w)
        cv2.line(img_grid, (w2, 0), (w2, h), grid_color, line_w)
        
        # Draw intersection circles
        intersections = [(w1, h1), (w2, h1), (w1, h2), (w2, h2)]
        circle_r = max(5, int(w * 0.01))
        for (ix, iy) in intersections:
            cv2.circle(img_grid, (ix, iy), circle_r, (255, 112, 67), -1) # Coral dots
            cv2.circle(img_grid, (ix, iy), circle_r + 2, (255, 255, 255), max(1, int(w * 0.001)))
            
        # Calculate visual texture density center (Canny edges)
        gray = cv2.cvtColor(img_numpy, cv2.COLOR_RGB2GRAY)
        edges = cv2.Canny(gray, 50, 150)
        
        # Get coordinates of all edge pixels
        edge_coords = np.argwhere(edges > 0)
        if len(edge_coords) > 0:
            # edge_coords holds (y, x)
            avg_y, avg_x = np.mean(edge_coords, axis=0)
            
            # Determine proximity to nearest intersection
            distances = [np.hypot(avg_x - ix, avg_y - iy) for (ix, iy) in intersections]
            min_dist = min(distances)
            max_possible_dist = np.hypot(w, h)
            proximity = 1 - (min_dist / (max_possible_dist * 0.25))
            proximity = max(0.0, min(1.0, proximity))
            
            # Label alignment
            if proximity > 0.70:
                align_text = f"Rule-of-Thirds Alignment: **Strong Alignment** ({proximity:.1%} visual intersection score). " \
                             "The main subject focus resides directly on one of the four power intersections, pulling viewer attention immediately."
            else:
                align_text = f"Rule-of-Thirds Alignment: **Centered/Diffuse Composition** ({proximity:.1%} visual intersection score). " \
                             "Visual weight is either balanced in the center or scattered across the frame, standard for documentarians or landscapes."
        else:
            align_text = "No strong visual edges found. Flat or uniform composition."
            
        return img_grid, align_text
    except Exception as e:
        print(f"Composition error: {e}")
        return img_numpy, f"Error processing geometry: {e}"

def analyze_lighting(img_numpy):
    """Computes a Plotly brightness distribution histogram and classifies visual lighting key."""
    try:
        gray = cv2.cvtColor(img_numpy, cv2.COLOR_RGB2GRAY)
        h, w = gray.shape
        total_pixels = h * w
        
        # Calculate luminance histogram
        hist = cv2.calcHist([gray], [0], None, [256], [0, 256]).flatten()
        
        # Metrics
        mean_brightness = np.mean(gray)
        std_brightness = np.std(gray)
        
        # Classify lighting key
        # High Key: Bright backgrounds, high mean, lower variance
        # Low Key: Shadow dominated, low mean, high variance (chiascuro)
        if mean_brightness >= 165:
            key_style = "High-Key Lighting (Bright & Open)"
            summary_desc = "Features bright, fully lit environments with soft shadows. Commonly utilized in consumer commercials, corporate flyers, and optimistic political campaign media to convey transparency and positive energy."
        elif mean_brightness <= 85:
            key_style = "Low-Key Lighting (Dramatic & Shadow-Heavy)"
            summary_desc = "Dominated by deep shadows, dark backgrounds, and stark contrast. Popular in film noir, investigative photojournalism, or negative attack advertisements to invoke mystery, tension, or critical framing."
        else:
            key_style = "Mid-Key / Standard Studio Lighting"
            summary_desc = "Features a realistic, balanced, or moderate lighting key. Popular in objective documentary filmmaking, standard portraiture, and everyday press releases to represent authenticity and balanced focus."
            
        # Draw Plotly Line Histogram
        df_hist = pd.DataFrame({
            "Luminance (0-255)": np.arange(256),
            "Pixel Count": hist
        })
        
        fig = px.line(
            df_hist, x="Luminance (0-255)", y="Pixel Count",
            title="Luminance Distribution Histogram (0 = Pure Black, 255 = Pure White)",
            template="plotly_dark"
        )
        fig.update_traces(line=dict(color="#f59e0b", width=3)) # Amber curve
        fig.update_layout(
            plot_bgcolor="#111827",
            paper_bgcolor="#0d0f12",
            margin=dict(l=20, r=20, t=50, b=20),
            xaxis=dict(showgrid=False),
            yaxis=dict(showgrid=False)
        )
        
        summary_text = f"Visual Key: **{key_style}**\n\n* **Average Luminance**: {mean_brightness:.1f} / 255\n* **Contrast Spread (Std Dev)**: {std_brightness:.1f}\n\n*Description*: {summary_desc}"
        
        return fig, summary_text
    except Exception as e:
        print(f"Lighting error: {e}")
        return go.Figure(), f"Error analyzing luminance: {e}"

def full_analyzer_pipeline(img):
    """Triggers the full deconstruction analysis when an image is uploaded."""
    if img is None:
        return go.Figure(), pd.DataFrame(), "No image uploaded.", None, "No image uploaded.", go.Figure(), "No image uploaded."
        
    # 1. Palette
    fig_pal, df_pal, pal_status = analyze_palette(img)
    
    # 2. Composition
    img_grid, comp_status = analyze_composition(img)
    
    # 3. Lighting
    fig_light, light_status = analyze_lighting(img)
    
    return fig_pal, df_pal, pal_status, img_grid, comp_status, fig_light, light_status

# Custom premium gradient CSS (Red/Yellow vibes)
custom_css = """
body { background-color: #0d0f12; color: #e3e6eb; font-family: 'Inter', sans-serif; }
.gradio-container { max-width: 1200px !important; margin: 0 auto !important; }
h1, h2, h3 { color: #ffffff !important; font-weight: 700 !important; }
.btn-primary { background: linear-gradient(135deg, #ef4444 0%, #f59e0b 100%) !important; border: none !important; color: white !important; font-weight: 600 !important; }
.btn-primary:hover { filter: brightness(1.1); }
.dataframe-container { background: #111827 !important; border: 1px solid #1f2937 !important; border-radius: 8px; }
"""

with gr.Blocks(theme=gr.themes.Monochrome(), css=custom_css) as demo:
    gr.Markdown(
        """
        # 🎨 Visual Rhetoric & Composition Analyzer
        ### Deconstruct media framing, analyze pixel color palettes (K-means), test Rule-of-Thirds alignment, and map lighting profiles to audit visual bias.
        """
    )
    
    with gr.Row():
        with gr.Column(scale=4):
            with gr.Card():
                gr.Markdown("### 1. Upload Visual Artifact")
                image_input = gr.Image(label="Upload Image (Campaign Flyer, News Photo, Advertisement)", type="numpy")
                analyze_btn = gr.Button("🎨 Run Rhetorical Deconstruction", variant="primary", elem_classes="btn-primary")
                
            with gr.Card():
                gr.Markdown("### 🔍 Rule-of-Thirds Grid Overlay")
                image_grid_output = gr.Image(label="Compositional Lines & Power Intersections", type="numpy", interactive=False)
                
        with gr.Column(scale=6):
            with gr.Tabs():
                with gr.TabItem("🎨 Color & Palette Psychology"):
                    palette_plot = gr.Plot(label="Dominant Pixel Palettes (K-Means)")
                    palette_status = gr.Markdown("Please upload an image to run analysis.")
                    
                    palette_table = gr.Dataframe(
                        headers=["Hex", "RGB", "Dominance", "Vibes & Framing Role"],
                        datatype=["str", "str", "number", "str"],
                        label="Dominant Colors Quantitative Distribution",
                        interactive=False,
                        elem_classes="dataframe-container"
                    )
                    
                with gr.TabItem("📐 Composition & Geometry"):
                    comp_status = gr.Markdown("Please upload an image to run analysis.")
                    gr.Markdown(
                        """
                        **Methodology**: 
                        - **Cyan lines** demarcate the vertical and horizontal 1/3 grid markers.
                        - **Coral dots** highlight the 4 focal intersections where visual elements naturally pull the highest attention.
                        - **Texture Center**: The app runs localized Sobel edge detection to find the image's texture 'center of gravity' and evaluates how close it is to these intersections.
                        """
                    )
                    
                with gr.TabItem("🔆 Lighting, Contrast & Mood"):
                    lighting_plot = gr.Plot(label="Luminance Profile (LCC)")
                    lighting_status = gr.Markdown("Please upload an image to run analysis.")

    # Core callback
    analyze_btn.click(
        fn=full_analyzer_pipeline,
        inputs=[image_input],
        outputs=[palette_plot, palette_table, palette_status, image_grid_output, comp_status, lighting_plot, lighting_status]
    )

if __name__ == "__main__":
    demo.launch()