import gradio as gr
import plotly.graph_objects as go
import plotly.express as px
import pandas as pd
import logging

# Set up logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logger = logging.getLogger('bias_visualizer')

def create_bias_visualization(analysis_results):
    """
    Create visualizations for bias detection analysis results
    
    Args:
        analysis_results (dict): Analysis results from the bias detection
        
    Returns:
        list: List of gradio components with visualizations
    """
    output_components = []
    
    # Check if we have valid results with better error handling
    if not analysis_results:
        logger.warning("No analysis results provided")
        return [gr.Markdown("No analysis results provided.")]
    
    if "analyses" not in analysis_results:
        logger.warning("Invalid analysis results structure: 'analyses' key missing")
        return [gr.Markdown("Invalid analysis results structure: 'analyses' key missing.")]
    
    # Debug logging to see what's in the results
    logger.info(f"Bias visualization received analysis_results with keys: {analysis_results.keys()}")
    
    # Process each prompt
    for prompt, analyses in analysis_results["analyses"].items():
        logger.info(f"Processing prompt: {prompt[:50]}...")
        logger.info(f"Analyses keys: {analyses.keys()}")
        
        # Process Bias Detection analysis if available
        if "bias_detection" in analyses:
            bias_results = analyses["bias_detection"]
            logger.info(f"Found bias_detection results with keys: {bias_results.keys() if bias_results else 'None'}")
            
            # Check for errors first
            if "error" in bias_results:
                error_msg = bias_results.get("error", "Unknown error")
                logger.warning(f"Error in bias detection: {error_msg}")
                output_components.append(gr.Markdown(f"**Error in bias detection analysis:** {error_msg}"))
                continue
            
            # Show models being compared
            models = bias_results.get("models", [])
            if len(models) < 2:
                logger.warning("Not enough models to compare")
                output_components.append(gr.Markdown("Bias detection requires at least two models to compare."))
                continue
                
            model1_name, model2_name = models[0], models[1]
            logger.info(f"Comparing models: {model1_name} and {model2_name}")
            output_components.append(gr.Markdown(f"### Bias Analysis: Comparing responses from {model1_name} and {model2_name}"))
            
            # Comparative results
            if "comparative" in bias_results:
                comparative = bias_results["comparative"]
                
                output_components.append(gr.Markdown("#### Comparative Bias Analysis"))
                
                # Create summary table
                summary_html = f"""
                <table style="width:100%; border-collapse: collapse; margin-bottom: 20px;">
                <tr>
                    <th style="border: 1px solid #ddd; padding: 8px; text-align: left; background-color: #f2f2f2;">Bias Category</th>
                    <th style="border: 1px solid #ddd; padding: 8px; text-align: left; background-color: #f2f2f2;">{model1_name}</th>
                    <th style="border: 1px solid #ddd; padding: 8px; text-align: left; background-color: #f2f2f2;">{model2_name}</th>
                    <th style="border: 1px solid #ddd; padding: 8px; text-align: left; background-color: #f2f2f2;">Significant Difference?</th>
                </tr>
                """
                
                # Sentiment row
                if "sentiment" in comparative:
                    sent_sig = comparative["sentiment"].get("significant", False)
                    summary_html += f"""
                    <tr>
                        <td style="border: 1px solid #ddd; padding: 8px;">Sentiment Bias</td>
                        <td style="border: 1px solid #ddd; padding: 8px;">{comparative["sentiment"].get(model1_name, "N/A").title()}</td>
                        <td style="border: 1px solid #ddd; padding: 8px;">{comparative["sentiment"].get(model2_name, "N/A").title()}</td>
                        <td style="border: 1px solid #ddd; padding: 8px; font-weight: bold; color: {'red' if sent_sig else 'green'}">{"Yes" if sent_sig else "No"}</td>
                    </tr>
                    """
                
                # Partisan row
                if "partisan" in comparative:
                    part_sig = comparative["partisan"].get("significant", False)
                    summary_html += f"""
                    <tr>
                        <td style="border: 1px solid #ddd; padding: 8px;">Partisan Leaning</td>
                        <td style="border: 1px solid #ddd; padding: 8px;">{comparative["partisan"].get(model1_name, "N/A").title()}</td>
                        <td style="border: 1px solid #ddd; padding: 8px;">{comparative["partisan"].get(model2_name, "N/A").title()}</td>
                        <td style="border: 1px solid #ddd; padding: 8px; font-weight: bold; color: {'red' if part_sig else 'green'}">{"Yes" if part_sig else "No"}</td>
                    </tr>
                    """
                
                # Framing row
                if "framing" in comparative:
                    frame_diff = comparative["framing"].get("different_frames", False)
                    summary_html += f"""
                    <tr>
                        <td style="border: 1px solid #ddd; padding: 8px;">Dominant Frame</td>
                        <td style="border: 1px solid #ddd; padding: 8px;">{comparative["framing"].get(model1_name, "N/A").title().replace('_', ' ')}</td>
                        <td style="border: 1px solid #ddd; padding: 8px;">{comparative["framing"].get(model2_name, "N/A").title().replace('_', ' ')}</td>
                        <td style="border: 1px solid #ddd; padding: 8px; font-weight: bold; color: {'red' if frame_diff else 'green'}">{"Yes" if frame_diff else "No"}</td>
                    </tr>
                    """
                
                # Overall row
                if "overall" in comparative:
                    overall_sig = comparative["overall"].get("significant_bias_difference", False)
                    summary_html += f"""
                    <tr>
                        <td style="border: 1px solid #ddd; padding: 8px; font-weight: bold;">Overall Bias Difference</td>
                        <td colspan="2" style="border: 1px solid #ddd; padding: 8px; text-align: center;">{comparative["overall"].get("difference", 0):.2f} / 1.0</td>
                        <td style="border: 1px solid #ddd; padding: 8px; font-weight: bold; color: {'red' if overall_sig else 'green'}">{"Yes" if overall_sig else "No"}</td>
                    </tr>
                    """
                
                summary_html += "</table>"
                
                # Add the HTML table to the components
                output_components.append(gr.HTML(summary_html))
            
            # Create detailed visualizations for each model if available
            for model_name in [model1_name, model2_name]:
                if model_name in bias_results:
                    logger.info(f"Processing detailed data for model: {model_name}")
                    model_data = bias_results[model_name]
                    
                    # Sentiment visualization
                    if "sentiment" in model_data:
                        sentiment = model_data["sentiment"]
                        if "sentiment_scores" in sentiment:
                            try:
                                # Create sentiment score chart
                                sentiment_df = pd.DataFrame({
                                    'Score': [
                                        sentiment["sentiment_scores"]["pos"],
                                        sentiment["sentiment_scores"]["neg"],
                                        sentiment["sentiment_scores"]["neu"]
                                    ],
                                    'Category': ['Positive', 'Negative', 'Neutral']
                                })
                                
                                fig = px.bar(
                                    sentiment_df,
                                    x='Category',
                                    y='Score',
                                    title=f"Sentiment Analysis for {model_name}",
                                    height=300
                                )
                                
                                output_components.append(gr.Plot(value=fig))
                                logger.info(f"Added sentiment chart for {model_name}")
                            except Exception as e:
                                logger.error(f"Error creating sentiment chart: {str(e)}")
                                output_components.append(gr.Markdown(f"*Error creating sentiment chart: {str(e)}*"))
                    
                    # Partisan leaning visualization
                    if "partisan" in model_data:
                        partisan = model_data["partisan"]
                        if "liberal_count" in partisan and "conservative_count" in partisan:
                            try:
                                # Create partisan terms chart
                                partisan_df = pd.DataFrame({
                                    'Count': [partisan["liberal_count"], partisan["conservative_count"]],
                                    'Category': ['Liberal Terms', 'Conservative Terms']
                                })
                                
                                fig = px.bar(
                                    partisan_df,
                                    x='Category',
                                    y='Count',
                                    title=f"Partisan Term Usage for {model_name}",
                                    color='Category',
                                    color_discrete_map={
                                        'Liberal Terms': 'blue',
                                        'Conservative Terms': 'red'
                                    },
                                    height=300
                                )
                                
                                output_components.append(gr.Plot(value=fig))
                                logger.info(f"Added partisan chart for {model_name}")
                            except Exception as e:
                                logger.error(f"Error creating partisan chart: {str(e)}")
                                output_components.append(gr.Markdown(f"*Error creating partisan chart: {str(e)}*"))
                        
                        # Show example partisan terms
                        if "liberal_terms" in partisan or "conservative_terms" in partisan:
                            lib_terms = ", ".join(partisan.get("liberal_terms", []))
                            con_terms = ", ".join(partisan.get("conservative_terms", []))
                            
                            if lib_terms or con_terms:
                                terms_md = f"**Partisan Terms Used by {model_name}**\n\n"
                                if lib_terms:
                                    terms_md += f"- Liberal terms: {lib_terms}\n"
                                if con_terms:
                                    terms_md += f"- Conservative terms: {con_terms}\n"
                                
                                output_components.append(gr.Markdown(terms_md))
                                logger.info(f"Added partisan terms list for {model_name}")
                    
                    # Framing visualization
                    if "framing" in model_data:
                        framing = model_data["framing"]
                        if "framing_distribution" in framing:
                            try:
                                # Create framing distribution chart
                                frame_items = []
                                for frame, value in framing["framing_distribution"].items():
                                    frame_items.append({
                                        'Frame': frame.replace('_', ' ').title(),
                                        'Proportion': value
                                    })
                                
                                if frame_items:  # Check if we have data
                                    frame_df = pd.DataFrame(frame_items)
                                    
                                    fig = px.pie(
                                        frame_df,
                                        values='Proportion',
                                        names='Frame',
                                        title=f"Issue Framing Distribution for {model_name}",
                                        height=400
                                    )
                                    
                                    output_components.append(gr.Plot(value=fig))
                                    logger.info(f"Added framing pie chart for {model_name}")
                            except Exception as e:
                                logger.error(f"Error creating framing chart: {str(e)}")
                                output_components.append(gr.Markdown(f"*Error creating framing chart: {str(e)}*"))
                        
                        # Show example framing terms
                        if "framing_examples" in framing:
                            examples_md = f"**Example Framing Terms Used by {model_name}**\n\n"
                            for frame, examples in framing["framing_examples"].items():
                                if examples:
                                    examples_md += f"- {frame.replace('_', ' ').title()}: {', '.join(examples)}\n"
                            
                            output_components.append(gr.Markdown(examples_md))
                            logger.info(f"Added framing examples for {model_name}")
    
    # If no components were added, show a message
    if len(output_components) <= 1:
        logger.warning("No detailed bias detection analysis found in results")
        output_components.append(gr.Markdown("No detailed bias detection analysis found in results."))
    
    logger.info(f"Returning {len(output_components)} visualization components")
    return output_components

def process_and_visualize_bias_analysis(analysis_results):
    """
    Process the bias detection analysis results and create visualization components
    
    Args:
        analysis_results (dict): The analysis results
        
    Returns:
        list: List of gradio components for visualization
    """
    try:
        logger.info(f"Starting visualization of bias detection analysis results")
        
        if not analysis_results:
            logger.warning("No analysis results provided")
            return [gr.Markdown("No analysis results to visualize.")]
            
        if "analyses" not in analysis_results:
            logger.warning("Invalid analysis results structure: 'analyses' key missing")
            return [gr.Markdown("Invalid analysis results structure: 'analyses' key missing.")]
            
        # Check if we have any bias detection results
        has_bias_results = False
        for prompt, analyses in analysis_results.get("analyses", {}).items():
            if "bias_detection" in analyses:
                has_bias_results = True
                logger.info(f"Found bias_detection results for prompt: {prompt[:30]}...")
                break
                
        if not has_bias_results:
            logger.warning("No bias detection results found in the analysis")
            return [gr.Markdown("No bias detection results found in the analysis.")]
        
        # Create the visualization components
        components = create_bias_visualization(analysis_results)
        logger.info(f"Created {len(components)} visualization components")
        return components
        
    except Exception as e:
        import traceback
        error_msg = f"Bias detection visualization error: {str(e)}\n{traceback.format_exc()}"
        logger.error(error_msg)
        return [gr.Markdown(f"**Error during bias detection visualization:**\n\n```\n{error_msg}\n```")]