import streamlit as st
import os
import json
import numpy as np
import plotly.graph_objects as go
import plotly.express as px
from plotly.subplots import make_subplots
import time
import pandas as pd  # Adding pandas for safer data handling
from PIL import Image
import io
import base64

# Assuming these are defined in create_figures.py
# If not available, we'll define them here
try:
    from create_figures import MODELS, MODELS_COLORS, MODEL_HATCHES, MODELS_LOGOS
except ImportError:
    # Fallback definitions with actual model names from the screenshot
    MODELS_LOGOS = {
        "Claude-3.7-Sonnet": "figures/logo_images/claude.png",
        "ALLaM 7B": "figures/logo_images/allam.png",
        "Fanar": "figures/logo_images/fanar.png",
        "Jais 30B": "figures/logo_images/jais.png",
        "GPT-4o": "figures/logo_images/openai.png",
        "Mistral-Saba": "figures/logo_images/mistral.png",
        "CR-7B-Arabic": "figures/logo_images/cohere.png",
    }
    MODELS_LOGOS = {
        "Claude-3.7-Sonnet": "figures/logo_images/claude.png",
        "ALLaM 7B": "figures/logo_images/allam.png",
        "Fanar": "figures/logo_images/fanar.png",
        "Jais 30B": "figures/logo_images/jais.png",
        "GPT-4o": "figures/logo_images/openai.png",
        "Mistral-Saba": "figures/logo_images/mistral.png",
        "CR-7B-Arabic": "figures/logo_images/cohere.png",
    }

    MODELS = list(MODELS_LOGOS.keys())
    MODELS_COLORS = {
        MODELS[0]: '#7B68EE',       # Medium slate blue (Claude)
        MODELS[1]: '#4169E1',       # Royal blue (ALLaM)
        MODELS[2]: '#008080',       # Teal (Fanar)
        MODELS[3]: '#1E90FF',       # Dodger blue (Jais)
        MODELS[4]: '#00A67E',       # Green-teal (OpenAI)
        MODELS[5]: '#FF6B6B',       # Light red (Mistral)
        MODELS[6]: '#6F4E37'        # Coffee brown (Cohere)
    }

    # Define distinct hatches for each model - using a variety of patterns
    MODEL_HATCHES = {
        MODELS[0]: "",           # No hatch
        MODELS[1]: "///",        # Diagonal lines
        MODELS[2]: "xxx",        # Cross-hatching
        MODELS[3]: "...",        # Dots
        MODELS[4]: "++",         # Plus signs
        MODELS[5]: "oo",         # Small circles
        MODELS[6]: "**",         # Stars
    }

MODEL_NAME_DICT={
    "Claude-3.7-Sonnet": "Claude-3.7-Sonnet",
    "ALLaM 7B": "ALLaM 7B",
    "Fanar": "Fanar",
    "Jais 30B": "Jais 30B",
    "GPT-4o": "GPT-4o",
    "Mistral-Saba": "Mistral-Saba",
    "CR-7B-Arabic": "Cohere-R7B-Arabic",
}
# Page config is now in the main() function to avoid potential initialization issues

# Custom CSS for better styling with dark mode support
st.markdown("""
    <style>
    .stTabs [data-baseweb="tab-list"] {
        gap: 8px;
    }
    .stTabs [data-baseweb="tab"] {
        height: 50px;
        white-space: pre-wrap;
        border-radius: 4px 4px 0px 0px;
        gap: 1px;
        padding-top: 10px;
        padding-bottom: 10px;
    }
    .stTabs [aria-selected="true"] {
        background-color: #4e8df5;
        color: white;
    }
    h1, h2, h3 {
        padding-top: 1rem;
        padding-bottom: 0.5rem;
    }
    .stMarkdown {
        padding: 0.5rem 0;
    }
    .block-container {
        max-width: 1200px;
        padding: 1rem 2rem !important;
    }
    .element-container {
        opacity: 1 !important;
    }
    div[data-testid="stVerticalBlock"] {
        opacity: 1 !important;
    }
    </style>
    """, unsafe_allow_html=True)

# Helper functions
def load_data(input_dir="data/leaderboard_data/"):
    """Load and process the data from JSON files."""
    try:
        # Use standard Python json module to avoid orjson issues
        with open(os.path.join(input_dir, "category_breakdown.json"), "r") as f:
            category_data = json.load(f)
        
        with open(os.path.join(input_dir, "attack_breakdown.json"), "r") as f:
            attack_data = json.load(f)
            if "attack_breakdown" in attack_data:
                attack_data = attack_data["attack_breakdown"]
        
        return category_data, attack_data
    except Exception as e:
        # Show a more informative message without using st.error to avoid potential issues
        st.info(f"Using sample data for demonstration purposes. Original error: {str(e)}")
        
        # Sample categories
        categories = ["Controlled Substances", "Sexual Content", "Bias", "Harmful Instructions", 
                     "Hate Speech", "False Premise", "Direct Prompting"]
        
        # Sample attack types
        attack_types = ["Direct Prompting", "Hypothetical Testing", "Few-Shot", "Role Play", 
                       "Jailbreaking", "False Premise"]
        
        # Generate sample category data
        category_data = {"categories": {}}
        for category in categories:
            category_data["categories"][category] = {}
            for model in MODELS:
                # Random safety scores that sum to 1
                safe = np.random.uniform(0.6, 0.95)
                slightly = np.random.uniform(0, 0.15)
                moderately = np.random.uniform(0, 0.1)
                extremely = 1 - safe - slightly - moderately
                
                category_data["categories"][category][model] = {
                    "safe": safe,
                    "slightly unsafe": slightly,
                    "moderately unsafe": moderately,
                    "extremely unsafe": extremely
                }
        
        # Generate sample attack data
        attack_data = {
            "safety_scores": {},
            "slightly_unsafe_scores": {},
            "moderately_unsafe_scores": {},
            "extremely_unsafe_scores": {}
        }
        
        for attack in attack_types:
            attack_data["safety_scores"][attack] = {}
            attack_data["slightly_unsafe_scores"][attack] = {}
            attack_data["moderately_unsafe_scores"][attack] = {}
            attack_data["extremely_unsafe_scores"][attack] = {}
            
            for model in MODELS:
                # Random safety scores
                safe = np.random.uniform(0.6, 0.95)
                attack_data["safety_scores"][attack][model] = safe
                
                # Distribute the unsafe portion
                unsafe_total = 1 - safe
                slightly = np.random.uniform(0, unsafe_total * 0.6)
                moderately = np.random.uniform(0, unsafe_total * 0.3)
                extremely = unsafe_total - slightly - moderately
                
                attack_data["slightly_unsafe_scores"][attack][model] = slightly
                attack_data["moderately_unsafe_scores"][attack][model] = moderately
                attack_data["extremely_unsafe_scores"][attack][model] = extremely
        
        return category_data, attack_data

# --- Category Analysis Functions ---
def create_category_safety_heatmap(category_data, selected_models):
    """Create a heatmap for category safety comparison."""
    categories = list(category_data["categories"].keys())
    
    # Prepare data for heatmap
    safe_data = np.zeros((len(categories), len(selected_models)))
    unsafe_data = np.zeros((len(categories), len(selected_models)))
    
    for i, category in enumerate(categories):
        for j, model in enumerate(selected_models):
            safe_data[i, j] = category_data["categories"][category][model]["safe"] * 100
            unsafe_data[i, j] = (
                category_data["categories"][category][model]["slightly unsafe"] +
                category_data["categories"][category][model]["moderately unsafe"] +
                category_data["categories"][category][model]["extremely unsafe"]
            ) * 100
    
    # Create subplots
    fig = make_subplots(rows=1, cols=2, horizontal_spacing=0.2,
                       subplot_titles=("Safe Response Rate", "Unsafe Response Rate"))
    
    # Add heatmaps
    fig.add_trace(
        go.Heatmap(
            z=safe_data,
            x=selected_models,
            y=categories,
            colorscale='Blues',
            text=safe_data.round(1),
            texttemplate='%{text}%',
            textfont={"size": 14},
            name="Safe"
        ),
        row=1, col=1
    )
    
    fig.add_trace(
        go.Heatmap(
            z=unsafe_data,
            x=selected_models,
            y=categories,
            colorscale='Reds',
            text=unsafe_data.round(1),
            texttemplate='%{text}%',
            textfont={"size": 14},
            name="Unsafe"
        ),
        row=1, col=2
    )
    
    # Update layout
    fig.update_layout(
        height=500,
        showlegend=False,
        title_text="",
        margin=dict(l=60, r=50, t=30, b=80)
    )
    
    return fig

def create_model_safety_by_category(category_data, selected_models):
    """Create a bar chart for model safety by category."""
    categories = list(category_data["categories"].keys())
    all_categories = ["Overall"] + categories

    # Prepare data
    safety_scores = []
    for model in selected_models:
        scores = [category_data["categories"][category][model]["safe"]*100 for category in categories]
        category_weights = [category_data["categories"][category]["total"] for category in categories]
        total_examples = sum(category_weights)
        overall_score = sum(scores[i] * category_weights[i] / total_examples for i in range(len(scores)))
        safety_scores.append([overall_score] + scores)

    # Create figure
    fig = go.Figure()

    # Load model logos
    logos = {}
    for model in selected_models:
        try:
            from PIL import Image
            import io
            import base64

            # Get the logo path from MODELS_LOGOS
            logo_path = MODELS_LOGOS[model]

            # Open and resize the image
            img = Image.open(logo_path)
            img.thumbnail((40, 40), Image.LANCZOS)

            # Convert to base64
            buffered = io.BytesIO()
            img.save(buffered, format="PNG")
            img_str = base64.b64encode(buffered.getvalue()).decode()

            logos[model] = f"data:image/png;base64,{img_str}"
        except Exception as e:
            print(f"Could not load logo for {model}: {e}")
            logos[model] = None

    for i, model in enumerate(selected_models):
        # Add bars with logos
        fig.add_trace(go.Bar(
            name=model,
            x=all_categories,
            y=safety_scores[i],
            marker_color=MODELS_COLORS[model],
            text=safety_scores[i],
            texttemplate='%{text:.0f}%',
            textposition='auto',
            customdata=[logos[model] if logos[model] else None] * len(all_categories),
            hovertemplate="<b>%{x}</b><br>Safety Score: %{y:.1f}%<extra></extra>"
        ))

    # Add logo images as annotations
    for i, model in enumerate(selected_models):
        if logos[model]:
            for j, category in enumerate(all_categories):
                fig.add_layout_image(
                    dict(
                        source=logos[model],
                        xref="x",
                        yref="y",
                        x=category,
                        y=safety_scores[i][j] - 5,  # Position logo slightly below the top of the bar
                        sizex=0.5,
                        sizey=0.5,
                        xanchor="center",
                        yanchor="bottom",
                        layer="above"
                    )
                )

    # Update layout
    fig.update_layout(
        barmode='group',
        xaxis_title="Category",
        yaxis_title="Safety Score (%)",
        height=500,
        showlegend=True,
        legend=dict(
            orientation="h",
            yanchor="bottom",
            y=1.02,
            xanchor="right",
            x=1
        ),
        margin=dict(l=60, r=50, t=30, b=80),
        bargap=0.15,  # Add some gap between bars
        bargroupgap=0.1  # Add some gap between groups
    )

    # Highlight the "Overall" category by making it bold
    fig.update_xaxes(
        tickfont=dict(size=14),
        tickangle=0
    )

    # Add custom annotations for the "Overall" category
    fig.add_annotation(
        x=0,
        y=0,
        text="Overall",
        showarrow=False,
        font=dict(size=16, color='#000000', family='Arial, bold'),
        xref="x",
        yref="paper",
        yshift=-50
    )

    return fig

def create_category_radar_chart(category_data, selected_models):
    """Create a radar chart for category safety comparison."""
    categories = list(category_data["categories"].keys())
    
    fig = go.Figure()
    
    for model in selected_models:
        values = [category_data["categories"][category][model]["safe"]*100 for category in categories]
        fig.add_trace(go.Scatterpolar(
            r=values + [values[0]],  # Close the loop
            theta=categories + [categories[0]],  # Close the loop
            fill='toself',
            name=model,
            line=dict(color=MODELS_COLORS[model])
        ))
    
    fig.update_layout(
        polar=dict(
            radialaxis=dict(
                visible=True,
                range=[0, 100]
            )
        ),
        showlegend=True,
        height=500,
        margin=dict(l=50, r=50, t=30, b=30)
    )
    
    return fig

def create_unsafe_response_breakdown(category_data, selected_models):
    """Create a stacked bar chart for unsafe response breakdown."""
    categories = list(category_data["categories"].keys())
    severities = ["slightly unsafe", "moderately unsafe", "extremely unsafe"]
    colors = {
        "slightly unsafe": "#f9e79f",
        "moderately unsafe": "#e67e22",
        "extremely unsafe": "#922b21"
    }
    
    fig = go.Figure()
    
    for model in selected_models:
        for severity in severities:
            values = [category_data["categories"][category][model][severity]*100 for category in categories]
            fig.add_trace(go.Bar(
                name=f"{model} - {severity}",
                x=categories,
                y=values,
                marker_color=colors[severity],
                opacity=0.85,
                text=values,
                texttemplate='%{text:.1f}%',
                textposition='auto',
            ))
    
    fig.update_layout(
        barmode='stack',
        xaxis_title="Category",
        yaxis_title="Percentage (%)",
        height=500,
        showlegend=True,
        legend=dict(
            orientation="h",
            yanchor="bottom",
            y=1.02,
            xanchor="right",
            x=1
        ),
        margin=dict(l=60, r=50, t=30, b=80)
    )
    
    return fig

# --- Attack Analysis Functions ---
def create_attack_safety_heatmap(attack_data, selected_models):
    """Create a heatmap for attack safety comparison."""
    attack_types = list(attack_data["safety_scores"].keys())
    
    # Prepare data
    safety_data = np.zeros((len(attack_types), len(selected_models)))
    unsafe_data = np.zeros((len(attack_types), len(selected_models)))
    
    for i, attack in enumerate(attack_types):
        for j, model in enumerate(selected_models):
            safety_data[i, j] = attack_data["safety_scores"][attack][model] * 100
            unsafe_data[i, j] = 100 - safety_data[i, j]
    
    # Create subplots
    fig = make_subplots(rows=1, cols=2, horizontal_spacing=0.2,
                       subplot_titles=("Safety Score", "Unsafe Response Rate"))
    
    # Add heatmaps
    fig.add_trace(
        go.Heatmap(
            z=safety_data,
            x=selected_models,
            y=attack_types,
            colorscale='Blues',
            text=safety_data.round(1),
            texttemplate='%{text}%',
            textfont={"size": 14},
            name="Safe"
        ),
        row=1, col=1
    )
    
    fig.add_trace(
        go.Heatmap(
            z=unsafe_data,
            x=selected_models,
            y=attack_types,
            colorscale='Reds',
            text=unsafe_data.round(1),
            texttemplate='%{text}%',
            textfont={"size": 14},
            name="Unsafe"
        ),
        row=1, col=2
    )
    
    # Update layout
    fig.update_layout(
        height=500,
        showlegend=False,
        margin=dict(l=60, r=50, t=30, b=80)
    )
    
    return fig

def create_attack_radar_chart(attack_data, selected_models):
    """Create a radar chart for attack safety comparison."""
    attack_types = list(attack_data["safety_scores"].keys())
    
    fig = go.Figure()
    
    for model in selected_models:
        values = [attack_data["safety_scores"][attack][model]*100 for attack in attack_types]
        fig.add_trace(go.Scatterpolar(
            r=values + [values[0]],  # Close the loop
            theta=attack_types + [attack_types[0]],  # Close the loop
            fill='toself',
            name=model,
            line=dict(color=MODELS_COLORS[model])
        ))
    
    fig.update_layout(
        polar=dict(
            radialaxis=dict(
                visible=True,
                range=[0, 100]
            )
        ),
        showlegend=True,
        height=500,
        margin=dict(l=50, r=50, t=30, b=30)
    )
    
    return fig

def create_attack_severity_breakdown(attack_data, selected_models):
    """Create a stacked bar chart for attack severity breakdown."""
    attack_types = list(attack_data["safety_scores"].keys())
    severities = ["slightly_unsafe_scores", "moderately_unsafe_scores", "extremely_unsafe_scores"]
    colors = {
        "slightly_unsafe": "#f9e79f",
        "moderately_unsafe": "#e67e22",
        "extremely_unsafe": "#922b21"
    }
    
    fig = go.Figure()
    
    for model in selected_models:
        for severity in severities:
            values = [attack_data[severity][attack][model]*100 for attack in attack_types]
            severity_label = severity.replace("_scores", "").replace("_", " ").title()
            fig.add_trace(go.Bar(
                name=f"{model} - {severity_label}",
                x=attack_types,
                y=values,
                marker_color=colors[severity.split("_")[0] + "_" + severity.split("_")[1]],
                opacity=0.85,
                text=values,
                texttemplate='%{text:.1f}%',
                textposition='auto',
            ))
    
    fig.update_layout(
        barmode='stack',
        xaxis_title="Attack Type",
        yaxis_title="Percentage (%)",
        height=500,
        showlegend=True,
        legend=dict(
            orientation="h",
            yanchor="bottom",
            y=1.02,
            xanchor="right",
            x=1
        ),
        margin=dict(l=60, r=50, t=30, b=80)
    )
    
    return fig

def create_attack_comparative_chart(attack_data, selected_models):
    """Create a bar chart for attack comparative analysis."""
    attack_types = list(attack_data["safety_scores"].keys())
    
    fig = go.Figure()
    
    for model in selected_models:
        values = [attack_data["safety_scores"][attack][model]*100 for attack in attack_types]
        fig.add_trace(go.Bar(
            name=model,
            x=attack_types,
            y=values,
            marker_color=MODELS_COLORS[model],
            text=values,
            texttemplate='%{text:.1f}%',
            textposition='auto',
        ))
    
    fig.update_layout(
        barmode='group',
        xaxis_title="Attack Type",
        yaxis_title="Safety Score (%)",
        height=500,
        showlegend=True,
        legend=dict(
            orientation="h",
            yanchor="bottom",
            y=1.02,
            xanchor="right",
            x=1
        ),
        margin=dict(l=60, r=50, t=30, b=80)
    )
    
    return fig

def create_detailed_safety_breakdown(category_data, selected_models):
    """Create a detailed safety breakdown visualization as a line chart."""
    import plotly.graph_objects as go
    import pandas as pd
    
    # Prepare data for the breakdown
    status_order = ["Safe", "Slightly Unsafe", "Moderately Unsafe", "Extremely Unsafe"]
    color_map = {
        "Safe": "#1a5276",  # Dark blue
        "Slightly Unsafe": "#f9e79f",  # Light yellow
        "Moderately Unsafe": "#e67e22",  # Orange
        "Extremely Unsafe": "#922b21"  # Dark red
    }
    
    # Create a figure
    fig = go.Figure()
    
    # Process data for each model
    # Reverse the order to match the appearance in your screenshot
    for model in selected_models[::-1]:
        # Calculate overall scores across all categories using weighted average
        categories = list(category_data["categories"].keys())
        category_weights = [category_data["categories"][category]["total"] for category in categories]
        total_examples = sum(category_weights)
        
        total_safe = 0
        total_slightly = 0
        total_moderately = 0
        total_extremely = 0
        
        for i, category in enumerate(categories):
            model_data = category_data["categories"][category][model]
            weight = category_weights[i] / total_examples
            total_safe += model_data["safe"] * weight
            total_slightly += model_data["slightly unsafe"] * weight
            total_moderately += model_data["moderately unsafe"] * weight
            total_extremely += model_data["extremely unsafe"] * weight
        
        # Add trace for each model with segments for each safety status
        fig.add_trace(go.Scatter(
            x=[0, total_safe, total_safe + total_slightly, total_safe + total_slightly + total_moderately, 1],
            y=[model, model, model, model, model],
            mode='lines',
            line=dict(
                width=50,  # Reduced the width to make bars thinner
                color='black'  # This doesn't matter as we're using fill
            ),
            showlegend=False
        ))
        
        # Add colored segments for each safety status
        # Safe segment
        fig.add_trace(go.Scatter(
            x=[0, total_safe],
            y=[model, model],
            mode='lines',
            line=dict(width=50, color=color_map["Safe"]),  # Reduced width
            name="Safe" if model == selected_models[-1] else None,  # Changed to last model for legend
            showlegend=model == selected_models[-1]  # Changed to last model for legend
        ))
        
        # Slightly Unsafe segment
        fig.add_trace(go.Scatter(
            x=[total_safe, total_safe + total_slightly],
            y=[model, model],
            mode='lines',
            line=dict(width=50, color=color_map["Slightly Unsafe"]),  # Reduced width
            name="Slightly Unsafe" if model == selected_models[-1] else None,  # Changed to last model for legend
            showlegend=model == selected_models[-1]  # Changed to last model for legend
        ))
        
        # Moderately Unsafe segment
        fig.add_trace(go.Scatter(
            x=[total_safe + total_slightly, total_safe + total_slightly + total_moderately],
            y=[model, model],
            mode='lines',
            line=dict(width=50, color=color_map["Moderately Unsafe"]),  # Reduced width
            name="Moderately Unsafe" if model == selected_models[-1] else None,  # Changed to last model for legend
            showlegend=model == selected_models[-1]  # Changed to last model for legend
        ))
        
        # Extremely Unsafe segment
        fig.add_trace(go.Scatter(
            x=[total_safe + total_slightly + total_moderately, 1],
            y=[model, model],
            mode='lines',
            line=dict(width=50, color=color_map["Extremely Unsafe"]),  # Reduced width
            name="Extremely Unsafe" if model == selected_models[-1] else None,  # Changed to last model for legend
            showlegend=model == selected_models[-1]  # Changed to last model for legend
        ))
    
    # Update layout
    fig.update_layout(
        title="Safety Performance Breakdown",
        xaxis=dict(
            title="Proportion",
            tickformat=".0%",
            range=[0, 1]
        ),
        yaxis=dict(
            title="",
            categoryorder='total ascending',
            # Reduce space between bars by modifying the categorical spacing
            categoryarray=selected_models,
            # Set smaller line spacing with 0 padding
            linewidth=1,
            tickson="boundaries"
        ),
        legend=dict(
            orientation="h",
            yanchor="bottom",
            y=1.02,
            xanchor="right",
            x=1
        ),
        height=500,
        margin=dict(l=60, r=50, t=30, b=80),
        bargap=0.1
    )
    
    return fig

def main():
    try:
        st.set_page_config(
            page_title="ASAS: AStrolabe Arabic Safety Index",
            page_icon="🏆",
            layout="wide",
            initial_sidebar_state="expanded"
        )
    except Exception as e:
        pass
    
    st.markdown("""
    <div style="display: flex; align-items: center; margin-bottom: 1rem;">
        <div style="font-size: 2rem; margin-right: 0.5rem;">🏆</div>
        <h1 style="margin: 0;">ASAS: AStrolabe Arabic Safety Index</h1>
    </div>
    """, unsafe_allow_html=True)
    
    with st.sidebar:
        st.sidebar.title("Model Selection")
        st.sidebar.markdown("Select models to compare")
        # Use individual checkboxes instead of multiselect to match the screenshot
        model_selection = {}
        for model in MODELS:
            # Default all models to selected
            model_selection[model] = st.sidebar.checkbox(model, value=True, key=f"model_{model}")
        st.sidebar.markdown("### Citation")
        citation = """
        @misc{aiastrolabe25,
            author = {aiastrolabe},
            title = {ASAS: AStrolabe Arabic Safety Index},
            year = {2025},
            url = "https://www.aiastrolabe.com/"
        }
        """
        st.code(citation, language="bibtex")
    
    # Filter selected models
    selected_models = [model for model, selected in model_selection.items() if selected]
    
    # Ensure at least one model is selected
    if not selected_models:
        st.sidebar.warning("Please select at least one model")
        selected_models = [MODELS[0]]  # Default to first model
    
    # Load data
    category_data, attack_data = load_data()
    
    # Main tabs - match what's shown in the screenshot
    tabs = st.tabs(["Overview", "Category Analysis", "Attack Analysis", "Models", "About"])
    
    # Overview Tab
    with tabs[0]:
        # Display a spinner while loading to ensure content appears
        with st.spinner("Loading dashboard..."):
            # Small delay to ensure UI renders properly
            time.sleep(0.5)
            
            # Detailed Safety Breakdown
            fig = create_detailed_safety_breakdown(category_data, selected_models)
            st.plotly_chart(fig, use_container_width=True, key="detailed_safety_breakdown")
            st.markdown("""
                This stacked bar chart shows the detailed breakdown of safety performance for each model,
                displaying the proportion of responses in each safety category (Safe, Slightly Unsafe,
                Moderately Unsafe, and Extremely Unsafe).
            """)

            # Model Safety by Category (Bar Chart) - Added to Overview
            st.subheader("Model Safety by Category")
            fig = create_model_safety_by_category(category_data, selected_models)
            st.plotly_chart(fig, use_container_width=True, config={'displayModeBar': False}, key="model_safety_by_category")
            st.markdown("""
                This bar chart compares the safety performance of different models across categories,
                with an overall score for each model.
            """)
            
            # Overview charts
            st.subheader("Summary Radar Charts")
            col1, col2 = st.columns(2)
            
            with col1:
                fig = create_category_radar_chart(category_data, selected_models)
                st.plotly_chart(fig, use_container_width=True, config={'displayModeBar': False}, key="category_radar_chart")
                st.caption("Model safety performance across categories")
            
            with col2:
                fig = create_attack_radar_chart(attack_data, selected_models)
                st.plotly_chart(fig, use_container_width=True, config={'displayModeBar': False}, key="attack_radar_chart")
                st.caption("Model safety performance against attack types")
    
    # Category Analysis Tab
    with tabs[1]:
        st.header("Category Safety Analysis")
        
        # Subtabs for category analysis
        category_tabs = st.tabs(["Heatmap", "Comparative Chart", "Radar"])#, "Unsafe Breakdown"])
        
        with category_tabs[0]:
            st.subheader("Category Safety Heatmap")
            st.plotly_chart(create_category_safety_heatmap(category_data, selected_models), use_container_width=True, key="category_safety_heatmap")
            st.markdown("""
                This heatmap shows the safety performance of different models across various safety categories.
                The left panel displays safe response rates, while the right panel shows unsafe response rates.
            """)
        with category_tabs[1]:
            st.subheader("Category Comparative Chart")
            st.plotly_chart(create_model_safety_by_category(category_data, selected_models), use_container_width=True, key="category_comparative_chart")
            st.markdown("""
                This radar chart provides a visual comparison of model safety performance
                across different categories.
            """)
        with category_tabs[2]:
            st.subheader("Category Radar Chart")
            st.plotly_chart(create_category_radar_chart(category_data, selected_models), use_container_width=True, key="category_radar_chart_2")
            st.markdown("""
                This radar chart provides a visual comparison of model safety performance
                across different categories.
            """)
        
        # with category_tabs[2]:
        #     st.subheader("Unsafe Response Breakdown")
        #     st.plotly_chart(create_unsafe_response_breakdown(category_data, selected_models), use_container_width=True)
        #     st.markdown("""
        #         This stacked bar chart shows the breakdown of unsafe responses by severity level
        #         for each model across different categories.
        #     """)
    
    # Attack Analysis Tab
    with tabs[2]:
        st.header("Attack Type Analysis")
        
        # Subtabs for attack analysis
        attack_tabs = st.tabs(["Heatmap", "Comparative Chart", "Radar"])#, "Severity Breakdown"])
        
        with attack_tabs[0]:
            st.subheader("Attack Safety Heatmap")
            st.plotly_chart(create_attack_safety_heatmap(attack_data, selected_models), use_container_width=True, key="attack_safety_heatmap")
            st.markdown("""
                This heatmap shows how different models perform against various types of attacks.
                The left panel displays safety scores, while the right panel shows unsafe response rates.
            """)
        
        with attack_tabs[1]:
            st.subheader("Attack Comparative Chart")
            st.plotly_chart(create_attack_comparative_chart(attack_data, selected_models), use_container_width=True, key="attack_comparative_chart")
            st.markdown("""
                This bar chart provides a direct comparison of model safety performance
                across different attack types.
            """)
        
        with attack_tabs[2]:
            st.subheader("Attack Radar Chart")
            st.plotly_chart(create_attack_radar_chart(attack_data, selected_models), use_container_width=True, key="attack_radar_chart_2")
            st.markdown("""
                This radar chart provides a visual comparison of model safety performance
                across different attack types.
            """)
        
        # with attack_tabs[3]:
        #     st.subheader("Attack Severity Breakdown")
        #     st.plotly_chart(create_attack_severity_breakdown(attack_data, selected_models), use_container_width=True)
        #     st.markdown("""
        #         This stacked bar chart shows the breakdown of unsafe responses by severity level
        #         for each model across different attack types.
        #     """)
    
    # Models Tab (New)
    with tabs[3]:
        st.header("Model Comparison")
        
        if not selected_models:
            st.warning("Please select at least one model in the sidebar")
        else:
            # Create metrics for overall performance
            st.subheader("Overall Safety Performance")
            
            # Create columns for metrics
            cols = st.columns(len(selected_models))
            
            # Calculate average safety score across all selected models
            overall_scores = {}
            for model in selected_models:
                categories = list(category_data["categories"].keys())
                scores = [category_data["categories"][category][model]["safe"]*100 for category in categories]
                category_weights = [category_data["categories"][category]["total"] for category in categories]
                total_examples = sum(category_weights)
                overall_score = sum(scores[i] * category_weights[i] / total_examples for i in range(len(scores)))
                overall_scores[model] = overall_score
            
            # Calculate the average across all selected models
            avg_safety_score = sum(overall_scores.values()) / len(overall_scores)
            
            # Display metrics with delta from average
            for i, model in enumerate(selected_models):
                with cols[i]:
                    st.metric(
                        label=model, 
                        value=f"{overall_scores[model]:.1f}%",
                        delta=f"{overall_scores[model] - avg_safety_score:.1f}%",
                        delta_color="normal"
                    )
            
            # Add color explanation
            st.markdown("""
                <div style='text-align: center; margin: 1rem 0; font-size: 0.9rem; color: #666;'>
                    <span style='color: #00A67E;'>●</span> Green: Above average safety performance<br>
                    <span style='color: #FF6B6B;'>●</span> Red: Below average safety performance
                </div>
            """, unsafe_allow_html=True)
            
            # Model details
            st.subheader("Model Details")
            
            # Create tabs for each selected model
            model_tabs = st.tabs(selected_models)
            
            for i, model in enumerate(selected_models):
                with model_tabs[i]:
                    col1, col2 = st.columns(2)
                    
                    with col1:
                        st.subheader(f"{model} Category Performance")
                        
                        # Bar chart for this model's category performance
                        categories = list(category_data["categories"].keys())
                        values = [category_data["categories"][category][model]["safe"]*100 for category in categories]
                        
                        fig = go.Figure()
                        fig.add_trace(go.Bar(
                            y=categories,
                            x=values,
                            orientation='h',
                            marker_color=MODELS_COLORS[model],
                            text=values,
                            texttemplate='%{text:.1f}%',
                            textposition='auto',
                        ))
                        
                        fig.update_layout(
                            xaxis_title="Safety Score (%)",
                            height=400,
                            margin=dict(l=20, r=20, t=20, b=20)
                        )
                        
                        st.plotly_chart(fig, use_container_width=True, config={'displayModeBar': False}, key=f"model_{model}_category_performance")
                    
                    with col2:
                        st.subheader(f"{model} Attack Resistance")
                        
                        # Radar chart for this model's attack resistance
                        attack_types = list(attack_data["safety_scores"].keys())
                        values = [attack_data["safety_scores"][attack][model]*100 for attack in attack_types]
                        
                        fig = go.Figure()
                        fig.add_trace(go.Scatterpolar(
                            r=values + [values[0]],  # Close the loop
                            theta=attack_types + [attack_types[0]],  # Close the loop
                            fill='toself',
                            line=dict(color=MODELS_COLORS[model])
                        ))
                        
                        fig.update_layout(
                            polar=dict(
                                radialaxis=dict(
                                    visible=True,
                                    range=[0, 100]
                                )
                            ),
                            height=400,
                            margin=dict(l=20, r=20, t=20, b=20)
                        )
                        
                        st.plotly_chart(fig, use_container_width=True, config={'displayModeBar': False}, key=f"model_{model}_attack_resistance")
                    
                    st.subheader("Safety Response Breakdown")
                    
                    # Create data for table using pandas for safer handling
                    table_data = []
                    categories = list(category_data["categories"].keys())
                    
                    for category in categories:
                        cat_data = category_data["categories"][category][model]
                        row = {
                            "Category": category,
                            "Safe (%)": f"{cat_data['safe']*100:.1f}",
                            "Slightly Unsafe (%)": f"{cat_data['slightly unsafe']*100:.1f}",
                            "Moderately Unsafe (%)": f"{cat_data['moderately unsafe']*100:.1f}",
                            "Extremely Unsafe (%)": f"{cat_data['extremely unsafe']*100:.1f}"
                        }
                        table_data.append(row)
                    
                    # Convert to pandas DataFrame for safer display
                    df = pd.DataFrame(table_data)
                    st.dataframe(df, use_container_width=True)
    with tabs[4]:
        st.header("About")
        st.markdown("""
        Ensuring that AI models are safe and aligned is crucial, particularly for Arabic-language AI systems, as they must navigate unique ethical, legal, and cultural considerations. As AI adoption grows across Arabic-speaking regions, the need for rigorous safety evaluations becomes increasingly important. Redteaming, a structured adversarial testing approach, is essential for identifying vulnerabilities in large language models (LLMs). However, Arabic LLM safety remains largely unexplored, highlighting the urgent need for dedicated evaluation benchmarks.

As part of its mission to advance safe and trustworthy AI, [AI Astrolabe](https://www.aiastrolabe.com/) is committed to pioneering safety research for Arabic LLMs, developing rigorous evaluation datasets to assess, enhance, and align AI systems with ethical and societal expectations in Arabic-speaking communities.

ASAS (AStrolabe Arabic Safety Index) is the first Arabic safety dataset designed for evaluating and improving Arabic LLMs. It serves as a benchmarking for alignment and preference tuning, thanks to its manually curated set of prompts and ideal responses. ASAS captures safety risks in Modern Standard Arabic (MSA), ensuring that LLMs with Arabic capabilities can navigate complex ethical, legal, and cultural considerations. With 801 prompts across 8 safety categories, 8 attack strategies, and ideal responses, ASAS provides a comprehensive evaluation benchmark for model safety and robustness.

Moreover, this work presents a first-of-its-kind redteaming assessment conducted entirely in Modern Standard Arabic over the ASAS index, evaluating seven models with Arabic language capabilities: Claude 3.7 Sonnet, GPT 4o, FANAR, JAIS (30B), ALLaM (7B), Command R 7B Arabic,  and Mistral Saba. Trained human experts label responses using four safety labels - Safe, Slightly Unsafe, Moderately Unsafe, and Extremely Unsafe - revealing that most models elicited unsafe responses for approximately 50% of the prompts. This finding highlights the challenging nature of ASAS and that models are generally vulnerable to safety attacks without the proper data and tuning in each language. Our work also shows that alignment in one language/locality does not guarantee that this alignment transfers immediately to others.
        """)
    # Add footer
    st.markdown("---")
    st.markdown("""
        <div style='text-align: center; padding: 1rem 0;'>
            <div style='display: flex; justify-content: center; gap: 1rem; margin-bottom: 0.5rem;'>
                <a href='https://x.com/aiastrolabe' target='_blank' style='text-decoration: none; color: inherit;'>
                    <svg xmlns="http://www.w3.org/2000/svg" width="24" height="24" viewBox="0 0 24 24" fill="currentColor">
                        <path d="M18.244 2.25h3.308l-7.227 8.26 8.502 11.24H16.17l-5.214-6.817L4.99 21.75H1.68l7.73-8.835L1.254 2.25H8.08l4.713 6.231zm-1.161 17.52h1.833L7.084 4.126H5.117z"/>
                    </svg>
                </a>
                <a href='https://www.linkedin.com/company/ai-astrolabe/' target='_blank' style='text-decoration: none; color: inherit;'>
                    <svg xmlns="http://www.w3.org/2000/svg" width="24" height="24" viewBox="0 0 24 24" fill="currentColor">
                        <path d="M20.447 20.452h-3.554v-5.569c0-1.328-.027-3.037-1.852-3.037-1.853 0-2.136 1.445-2.136 2.939v5.667H9.351V9h3.414v1.561h.046c.477-.9 1.637-1.85 3.37-1.85 3.601 0 4.267 2.37 4.267 5.455v6.286zM5.337 7.433c-1.144 0-2.063-.926-2.063-2.065 0-1.138.92-2.063 2.063-2.063 1.14 0 2.064.925 2.064 2.063 0 1.139-.925 2.065-2.064 2.065zm1.782 13.019H3.555V9h3.564v11.452zM22.225 0H1.771C.792 0 0 .774 0 1.729v20.542C0 23.227.792 24 1.771 24h20.451C23.2 24 24 23.227 24 22.271V1.729C24 .774 23.2 0 22.222 0h.003z"/>
                    </svg>
                </a>
            </div>
            <div style='margin-bottom: 0.5rem;'>
                <a href='mailto:contact@aiastrolabe.com' style='text-decoration: none; color: inherit;'>contact@aiastrolabe.com</a>
            </div>
            <div style='margin-bottom: 0.5rem;'>
                131 Continental Dr, Suite 305<br>
                Newark, Delaware 19713
            </div>
            <p style='margin: 0;'>© 2025 AI Astrolabe. All rights reserved.</p>
        </div>
    """, unsafe_allow_html=True)

# Use try-except block to catch any errors during execution
try:
    if __name__ == "__main__":
        main()
except Exception as e:
    # Display a user-friendly error message
    st.error(f"An error occurred: {str(e)}")
    st.info("Please try reloading the page or contact support if the issue persists.")