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

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

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 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')
    
    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)
        ]

    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>"
            
            category_yes = 0
            category_no = 0
            category_na = 0
            
            for section, details in category_data.items():
                status = details['status']
                sources = details.get('sources', [])
                questions = details.get('questions', {})
                
                # Determine section class based on status
                section_class = "section-na" if status == "N/A" else "section-active"
                card_content += f"<div class='section {section_class}'>"
                card_content += f"<h3>{section}</h3>"
                
                # Add sources if they exist
                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>"
                
                # Process questions
                if questions:
                    card_content += "<div class='questions'>"
                    for question, is_checked in questions.items():
                        if status == "N/A":
                            style_class = "na"
                            icon = "○"  # Circle for N/A items
                            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>"
                
                card_content += "</div>"  # Close section div
            
            # Calculate category score (excluding N/A items)
            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>"  # Close card div
            all_cards_content += card_content

    all_cards_content += "</div>"
    
    # Calculate total score (excluding N/A items)
    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_md, 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;
}
"""

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=[cat.name for cat in scorecard_template],
                                         label="Filter Categories", 
                                         value=[cat.name for cat in scorecard_template],
                                         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.Markdown()
        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()