import gradio as gr
import pandas as pd
import plotly.express as px
from dataclasses import dataclass, field
from typing import List, Dict, Tuple, Union
import json
import os
from collections import OrderedDict
import re

@dataclass
class ScorecardCategory:
    name: str
    questions: List[Dict[str, Union[str, List[str]]]]
    scores: Dict[str, int] = field(default_factory=dict)


def extract_category_number(category_name: str) -> int:
    """Extract the category number from the category name."""
    match = re.match(r'^(\d+)\.?\s*.*$', category_name)
    return int(match.group(1)) if match else float('inf')

def sort_categories(categories):
    """Sort categories by their numeric prefix."""
    return sorted(categories, key=extract_category_number)


# def load_scorecard_templates(directory):
#     templates = []
#     for filename in os.listdir(directory):
#         if filename.endswith('.json'):
#             with open(os.path.join(directory, filename), 'r') as file:
#                 data = json.load(file)
#                 templates.append(ScorecardCategory(
#                     name=data['name'],
#                     questions=data['questions']
#                 ))
#     return templates

def get_modality_icon(modality):
    """Return an emoji icon for each modality type."""
    icons = {
        "Text-to-Text": "📝",  # Memo icon for text-to-text
        "Text-to-Image": "🎨",  # Artist palette for text-to-image
        "Image-to-Text": "🔍",  # Magnifying glass for image-to-text
        "Image-to-Image": "🖼️",  # Frame for image-to-image
        "Audio": "🎵",  # Musical note for audio
        "Video": "🎬",  # Clapper board for video
        "Multimodal": "🔄"  # Cycle arrows for multimodal
    }
    return icons.get(modality, "💫")  # Default icon if modality not found

def create_metadata_card(metadata):
    """Create a formatted HTML card for metadata."""
    html = "<div class='card metadata-card'>"
    html += "<div class='card-title'>Model Information</div>"
    html += "<div class='metadata-content'>"
    
    # Handle special formatting for modalities
    modalities = metadata.get("Modalities", [])
    formatted_modalities = ""
    if modalities:
        formatted_modalities = " ".join(
            f"<span class='modality-badge'>{get_modality_icon(m)} {m}</span>"
            for m in modalities
        )
    
    # Order of metadata display (customize as needed)
    display_order = ["Name", "Provider", "Type", "URL"]
    
    # Display ordered metadata first
    for key in display_order:
        if key in metadata:
            value = metadata[key]
            if key == "URL":
                html += f"<div class='metadata-row'><span class='metadata-label'>{key}:</span> "
                html += f"<a href='{value}' target='_blank' class='metadata-link'>{value}</a></div>"
            else:
                html += f"<div class='metadata-row'><span class='metadata-label'>{key}:</span> <span class='metadata-value'>{value}</span></div>"
    
    # Add modalities if present
    if formatted_modalities:
        html += f"<div class='metadata-row'><span class='metadata-label'>Modalities:</span> <div class='modality-container'>{formatted_modalities}</div></div>"
    
    # Add any remaining metadata not in display_order
    for key, value in metadata.items():
        if key not in display_order and key != "Modalities":
            html += f"<div class='metadata-row'><span class='metadata-label'>{key}:</span> <span class='metadata-value'>{value}</span></div>"
    
    html += "</div></div>"
    return html


def load_models_from_json(directory):
    models = {}
    for filename in os.listdir(directory):
        if filename.endswith('.json'):
            with open(os.path.join(directory, filename), 'r') as file:
                model_data = json.load(file)
                model_name = model_data['metadata']['Name']
                models[model_name] = model_data
    
    return OrderedDict(sorted(models.items(), key=lambda x: x[0].lower()))

# Load templates and models
# scorecard_template = load_scorecard_templates('scorecard_templates')
models = load_models_from_json('model_data')

def create_source_html(sources):
    if not sources:
        return ""
    
    html = "<div class='sources-list'>"
    for source in sources:
        icon = source.get("type", "")
        detail = source.get("detail", "")
        name = source.get("name", detail)
        
        html += f"<div class='source-item'>{icon} "
        if detail.startswith("http"):
            html += f"<a href='{detail}' target='_blank'>{name}</a>"
        else:
            html += name
        html += "</div>"
    html += "</div>"
    return html

def create_leaderboard():
    scores = []
    for model, data in models.items():
        total_score = 0
        total_questions = 0
        
        for category in data['scores'].values():
            for section in category.values():
                if section['status'] != 'N/A':
                    questions = section.get('questions', {})
                    total_score += sum(1 for q in questions.values() if q)
                    total_questions += len(questions)
        
        score_percentage = (total_score / total_questions * 100) if total_questions > 0 else 0
        scores.append((model, score_percentage))
    
    df = pd.DataFrame(scores, columns=['Model', 'Score Percentage'])
    df = df.sort_values('Score Percentage', ascending=False).reset_index(drop=True)
    
    html = "<div class='card leaderboard-card'>"
    html += "<div class='card-title'>AI Model Social Impact Leaderboard</div>"
    html += "<table class='leaderboard-table'>"
    html += "<tr><th>Rank</th><th>Model</th><th>Score Percentage</th></tr>"
    for i, (_, row) in enumerate(df.iterrows(), 1):
        html += f"<tr><td>{i}</td><td>{row['Model']}</td><td>{row['Score Percentage']:.2f}%</td></tr>"
    html += "</table></div>"
    
    return html

def create_category_chart(selected_models, selected_categories):
    if not selected_models:
        return px.bar(title='Please select at least one model for comparison')
    
    # Sort categories before processing
    selected_categories = sort_categories(selected_categories)
    
    data = []
    for model in selected_models:
        for category in selected_categories:
            if category in models[model]['scores']:
                total_score = 0
                total_questions = 0
                
                for section in models[model]['scores'][category].values():
                    if section['status'] != 'N/A':
                        questions = section.get('questions', {})
                        total_score += sum(1 for q in questions.values() if q)
                        total_questions += len(questions)
                
                score_percentage = (total_score / total_questions * 100) if total_questions > 0 else 0
                data.append({
                    'Model': model,
                    'Category': category,
                    'Score Percentage': score_percentage
                })
    
    df = pd.DataFrame(data)
    if df.empty:
        return px.bar(title='No data available for the selected models and categories')
    
    fig = px.bar(df, x='Model', y='Score Percentage', color='Category',
                 title='AI Model Scores by Category',
                 labels={'Score Percentage': 'Score Percentage'},
                 category_orders={"Category": selected_categories})
    return fig

def update_detailed_scorecard(model, selected_categories):
    if not model:
        return [
            gr.update(value="Please select a model to view details.", visible=True),
            gr.update(visible=False),
            gr.update(visible=False)
        ]
    
    print("Selected categories:", selected_categories)
    print("Available categories in model:", list(models[model]['scores'].keys()))

    # Sort categories before processing
    selected_categories = sort_categories(selected_categories)
    metadata_html = create_metadata_card(models[model]['metadata'])


    # metadata_md = f"## Model Metadata for {model}\n\n"
    # for key, value in models[model]['metadata'].items():
    #     metadata_md += f"**{key}:** {value}\n\n"

    total_yes = 0
    total_no = 0
    total_na = 0

    all_cards_content = "<div class='container'>"
    for category_name in selected_categories:
        if category_name in models[model]['scores']:
            category_data = models[model]['scores'][category_name]
            card_content = f"<div class='card'><div class='card-title'>{category_name}</div>"
            
            # Sort sections within each category
            sorted_sections = sorted(category_data.items(), 
                                  key=lambda x: float(re.match(r'^(\d+\.?\d*)', x[0]).group(1)))
            
            category_yes = 0
            category_no = 0
            category_na = 0
            
            for section, details in sorted_sections:
                status = details['status']
                sources = details.get('sources', [])
                questions = details.get('questions', {})
                
                section_class = "section-na" if status == "N/A" else "section-active"
                status_class = status.lower()
                status_icon = "●" if status == "Yes" else "○" if status == "N/A" else "×"
                
                card_content += f"<div class='section {section_class}'>"
                card_content += f"<div class='section-header'><h3>{section}</h3>"
                card_content += f"<span class='status-badge {status_class}'>{status_icon} {status}</span></div>"
                
                if sources:
                    card_content += "<div class='sources-list'>"
                    for source in sources:
                        icon = source.get("type", "")
                        detail = source.get("detail", "")
                        name = source.get("name", detail)
                        
                        card_content += f"<div class='source-item'>{icon} "
                        if detail.startswith("http"):
                            card_content += f"<a href='{detail}' target='_blank'>{name}</a>"
                        else:
                            card_content += name
                        card_content += "</div>"
                    card_content += "</div>"
                
                if questions:
                    yes_count = sum(1 for v in questions.values() if v)
                    total_count = len(questions)
                    
                    card_content += "<details class='question-accordion'>"
                    if status == "N/A":
                        card_content += f"<summary>View {total_count} N/A items</summary>"
                    else:
                        card_content += f"<summary>View details ({yes_count}/{total_count} completed)</summary>"
                    
                    card_content += "<div class='questions'>"
                    for question, is_checked in questions.items():
                        if status == "N/A":
                            style_class = "na"
                            icon = "○"
                            category_na += 1
                            total_na += 1
                        else:
                            if is_checked:
                                style_class = "checked"
                                icon = "✓"
                                category_yes += 1
                                total_yes += 1
                            else:
                                style_class = "unchecked"
                                icon = "✗"
                                category_no += 1
                                total_no += 1
                        
                        card_content += f"<div class='question-item {style_class}'>{icon} {question}</div>"
                    card_content += "</div></details>"
                
                card_content += "</div>"
            
            if category_yes + category_no > 0:
                category_score = category_yes / (category_yes + category_no) * 100
                card_content += f"<div class='category-score'>Category Score: {category_score:.2f}% (Yes: {category_yes}, No: {category_no}, N/A: {category_na})</div>"
            elif category_na > 0:
                card_content += f"<div class='category-score'>Category Score: N/A (All {category_na} items not applicable)</div>"
            
            card_content += "</div>"
            all_cards_content += card_content

    all_cards_content += "</div>"
    
    if total_yes + total_no > 0:
        total_score = total_yes / (total_yes + total_no) * 100
        total_score_md = f"<div class='total-score'>Total Score: {total_score:.2f}% (Yes: {total_yes}, No: {total_no}, N/A: {total_na})</div>"
    else:
        total_score_md = "<div class='total-score'>No applicable scores (all items N/A)</div>"
    
    return [
        gr.update(value=metadata_html, visible=True),
        gr.update(value=all_cards_content, visible=True),
        gr.update(value=total_score_md, visible=True)
    ]

css = """
.container {
    display: flex;
    flex-wrap: wrap;
    justify-content: space-between;
}
.container.svelte-1hfxrpf.svelte-1hfxrpf {
    height: 0%;
}
.card {
    width: calc(50% - 20px);
    border: 1px solid #e0e0e0;
    border-radius: 10px;
    padding: 20px;
    margin-bottom: 20px;
    background-color: #ffffff;
    box-shadow: 0 4px 6px rgba(0,0,0,0.1);
    transition: all 0.3s ease;
}
.card:hover {
    box-shadow: 0 6px 8px rgba(0,0,0,0.15);
    transform: translateY(-5px);
}
.card-title {
    font-size: 1.4em;
    font-weight: bold;
    margin-bottom: 15px;
    color: #333;
    border-bottom: 2px solid #e0e0e0;
    padding-bottom: 10px;
}
.sources-list {
    margin: 10px 0;
}
.source-item {
    margin: 5px 0;
    padding: 5px;
    background-color: #f8f9fa;
    border-radius: 4px;
}
.question-item {
    margin: 5px 0;
    padding: 8px;
    border-radius: 4px;
}
.question-item.checked {
    background-color: #e6ffe6;
}
.question-item.unchecked {
    background-color: #ffe6e6;
}
.category-score, .total-score {
    background-color: #f0f8ff;
    border: 1px solid #b0d4ff;
    border-radius: 5px;
    padding: 10px;
    margin-top: 15px;
    font-weight: bold;
    text-align: center;
}
.total-score {
    font-size: 1.2em;
    background-color: #e6f3ff;
    border-color: #80bdff;
}
.leaderboard-card {
    width: 100%;
    max-width: 800px;
    margin: 0 auto;
}
.leaderboard-table {
    width: 100%;
    border-collapse: collapse;
}
.leaderboard-table th, .leaderboard-table td {
    padding: 10px;
    text-align: left;
    border-bottom: 1px solid #e0e0e0;
}
.leaderboard-table th {
    background-color: #f2f2f2;
    font-weight: bold;
}
.section {
    margin-bottom: 20px;
    padding: 15px;
    border-radius: 5px;
    background-color: #f8f9fa;
}
@media (max-width: 768px) {
    .card {
        width: 100%;
    }
}
.dark {
    background-color: #1a1a1a;
    color: #e0e0e0;

    .card {
        background-color: #2a2a2a;
        border-color: #444;
    }
    .card-title {
        color: #fff;
        border-bottom-color: #444;
    }
    .source-item {
        background-color: #2a2a2a;
    }
    .question-item.checked {
        background-color: #1a3a1a;
    }
    .question-item.unchecked {
        background-color: #3a1a1a;
    }
    .section {
        background-color: #2a2a2a;
    }
    .category-score, .total-score {
        background-color: #2c3e50;
        border-color: #34495e;
    }
    .leaderboard-table th {
        background-color: #2c3e50;
    }
}

.section-na {
    opacity: 0.6;
}

.question-item.na {
    background-color: #f0f0f0;
    color: #666;
}

.dark .question-item.na {
    background-color: #2d2d2d;
    color: #999;
}

.section-header {
    display: flex;
    justify-content: space-between;
    align-items: center;
    margin-bottom: 10px;
}

.status-badge {
    font-size: 0.9em;
    padding: 4px 8px;
    border-radius: 12px;
    font-weight: 500;
}

.status-badge.yes {
    background-color: #e6ffe6;
    color: #006600;
}

.status-badge.no {
    background-color: #ffe6e6;
    color: #990000;
}

.status-badge.n\/a {
    background-color: #f0f0f0;
    color: #666666;
}

.question-accordion {
    margin-top: 10px;
}

.question-accordion summary {
    cursor: pointer;
    padding: 8px;
    background-color: #f8f9fa;
    border-radius: 4px;
    margin-bottom: 10px;
    font-weight: 500;
}

.question-accordion summary:hover {
    background-color: #e9ecef;
}

.dark .status-badge.yes {
    background-color: #1a3a1a;
    color: #90EE90;
}

.dark .status-badge.no {
    background-color: #3a1a1a;
    color: #FFB6B6;
}

.dark .status-badge.n\/a {
    background-color: #2d2d2d;
    color: #999999;
}

.dark .question-accordion summary {
    background-color: #2a2a2a;
}

.dark .question-accordion summary:hover {
    background-color: #333333;
}
.metadata-card {
    margin-bottom: 30px;
    width: 100% !important;
}

.metadata-content {
    display: flex;
    flex-direction: column;
    gap: 12px;
}

.metadata-row {
    display: flex;
    align-items: flex-start;
    gap: 10px;
    line-height: 1.5;
}

.metadata-label {
    font-weight: 600;
    min-width: 100px;
    color: #555;
}

.metadata-value {
    color: #333;
}

.metadata-link {
    color: #007bff;
    text-decoration: none;
}

.metadata-link:hover {
    text-decoration: underline;
}

.modality-container {
    display: flex;
    flex-wrap: wrap;
    gap: 8px;
}

.modality-badge {
    display: inline-flex;
    align-items: center;
    gap: 4px;
    padding: 4px 10px;
    background-color: #f0f7ff;
    border: 1px solid #cce3ff;
    border-radius: 15px;
    font-size: 0.9em;
    color: #0066cc;
}

.dark .metadata-label {
    color: #aaa;
}

.dark .metadata-value {
    color: #ddd;
}

.dark .metadata-link {
    color: #66b3ff;
}

.dark .modality-badge {
    background-color: #1a2733;
    border-color: #2c3e50;
    color: #99ccff;
}
"""

first_model = next(iter(models.values()))
category_choices = list(first_model['scores'].keys())

with gr.Blocks(css=css) as demo:
    gr.Markdown("# AI Model Social Impact Scorecard Dashboard")
    
    with gr.Row():
        tab_selection = gr.Radio(["Leaderboard", "Category Analysis", "Detailed Scorecard"], 
                                label="Select Tab", value="Leaderboard")
    
    with gr.Row():
        model_chooser = gr.Dropdown(choices=[""] + list(models.keys()),
                                  label="Select Model for Details", 
                                  value="",
                                  interactive=True, visible=False)
        model_multi_chooser = gr.Dropdown(choices=list(models.keys()),
                                        label="Select Models for Comparison", 
                                        multiselect=True, interactive=True, visible=False)
        category_filter = gr.CheckboxGroup(choices=category_choices,
                                         label="Filter Categories", 
                                         value=category_choices,
                                         visible=False)
    
    with gr.Column(visible=True) as leaderboard_tab:
        leaderboard_output = gr.HTML()
    
    with gr.Column(visible=False) as category_analysis_tab:
        category_chart = gr.Plot()
    
    with gr.Column(visible=False) as detailed_scorecard_tab:
        model_metadata = gr.HTML()
        all_category_cards = gr.HTML()
        total_score = gr.Markdown()

    # Initialize the dashboard with the leaderboard
    leaderboard_output.value = create_leaderboard()
    
    def update_dashboard(tab, selected_models, selected_model, selected_categories):
            leaderboard_visibility = gr.update(visible=False)
            category_chart_visibility = gr.update(visible=False)
            detailed_scorecard_visibility = gr.update(visible=False)
            model_chooser_visibility = gr.update(visible=False)
            model_multi_chooser_visibility = gr.update(visible=False)
            category_filter_visibility = gr.update(visible=False)

            if tab == "Leaderboard":
                leaderboard_visibility = gr.update(visible=True)
                leaderboard_html = create_leaderboard()
                return [leaderboard_visibility, category_chart_visibility, detailed_scorecard_visibility,
                        model_chooser_visibility, model_multi_chooser_visibility, category_filter_visibility,
                        gr.update(value=leaderboard_html), gr.update(), gr.update(), gr.update(), gr.update()]
            
            elif tab == "Category Analysis":
                category_chart_visibility = gr.update(visible=True)
                model_multi_chooser_visibility = gr.update(visible=True)
                category_filter_visibility = gr.update(visible=True)
                category_plot = create_category_chart(selected_models or [], selected_categories)
                return [leaderboard_visibility, category_chart_visibility, detailed_scorecard_visibility,
                        model_chooser_visibility, model_multi_chooser_visibility, category_filter_visibility,
                        gr.update(), gr.update(value=category_plot), gr.update(), gr.update(), gr.update()]
            
            elif tab == "Detailed Scorecard":
                detailed_scorecard_visibility = gr.update(visible=True)
                model_chooser_visibility = gr.update(visible=True)
                category_filter_visibility = gr.update(visible=True)
                if selected_model:
                    scorecard_updates = update_detailed_scorecard(selected_model, selected_categories)
                else:
                    scorecard_updates = [
                        gr.update(value="Please select a model to view details.", visible=True),
                        gr.update(visible=False),
                        gr.update(visible=False)
                    ]
                return [leaderboard_visibility, category_chart_visibility, detailed_scorecard_visibility,
                        model_chooser_visibility, model_multi_chooser_visibility, category_filter_visibility,
                        gr.update(), gr.update()] + scorecard_updates

    # Set up event handlers
    tab_selection.change(
        fn=update_dashboard,
        inputs=[tab_selection, model_multi_chooser, model_chooser, category_filter],
        outputs=[leaderboard_tab, category_analysis_tab, detailed_scorecard_tab,
                model_chooser, model_multi_chooser, category_filter,
                leaderboard_output, category_chart, model_metadata,
                all_category_cards, total_score]
    )

    model_chooser.change(
        fn=update_dashboard,
        inputs=[tab_selection, model_multi_chooser, model_chooser, category_filter],
        outputs=[leaderboard_tab, category_analysis_tab, detailed_scorecard_tab,
                model_chooser, model_multi_chooser, category_filter,
                leaderboard_output, category_chart, model_metadata,
                all_category_cards, total_score]
    )

    model_multi_chooser.change(
        fn=update_dashboard,
        inputs=[tab_selection, model_multi_chooser, model_chooser, category_filter],
        outputs=[leaderboard_tab, category_analysis_tab, detailed_scorecard_tab,
                model_chooser, model_multi_chooser, category_filter,
                leaderboard_output, category_chart, model_metadata,
                all_category_cards, total_score]
    )

    category_filter.change(
        fn=update_dashboard,
        inputs=[tab_selection, model_multi_chooser, model_chooser, category_filter],
        outputs=[leaderboard_tab, category_analysis_tab, detailed_scorecard_tab,
                model_chooser, model_multi_chooser, category_filter,
                leaderboard_output, category_chart, model_metadata,
                all_category_cards, total_score]
    )

# Launch the app
if __name__ == "__main__":
    demo.launch()