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

# Load scorecard templates
scorecard_template = load_scorecard_templates('scorecard_templates')

# Function to read JSON files and populate models dictionary
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
    
    # Sort the models alphabetically by name
    return OrderedDict(sorted(models.items(), key=lambda x: x[0].lower()))

# Load models from JSON files
models = load_models_from_json('model_data')

css = """
.container {
    display: flex;
    flex-wrap: wrap;
    justify-content: space-between;
}
.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;
}
.question {
    margin-bottom: 20px;
    padding: 15px;
    border-radius: 5px;
}
.question h3 {
    margin-top: 0;
    color: #2c3e50;
}
.question-yes {
    background-color: #e6ffe6;
}
.question-no {
    background-color: #ffe6e6;
}
.question-na {
    background-color: #fffde6;
}
.status {
    font-weight: bold;
}
details {
    margin-top: 10px;
}
summary {
    cursor: pointer;
    color: #3498db;
    font-weight: bold;
}
summary:hover {
    text-decoration: underline;
}
.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;
}
.leaderboard-table tr:last-child td {
    border-bottom: none;
}
@media (max-width: 768px) {
    .card {
        width: 100%;
    }
}
"""

def create_leaderboard():
    scores = []
    for model, data in models.items():
        total_score = 0
        total_questions = 0
        for category in data['scores']:
            for question, details in data['scores'][category].items():
                if details['status'] == 'Yes':
                    total_score += 1
                total_questions += 1
        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_questions = len(models[model]['scores'][category])
                yes_count = sum(1 for q in models[model]['scores'][category].values() if q['status'] == 'Yes')
                score_percentage = (yes_count / 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 model not in models:
        return [gr.update(visible=True, value="Please select a model to view details.")] + [gr.update(visible=False)] * 2

    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 in scorecard_template:
        if category.name in selected_categories and 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 question, details in category_data.items():
                status = details['status']
                source = details.get('source', 'N/A')
                
                if status == 'Yes':
                    bg_class = 'question-yes'
                    category_yes += 1
                    total_yes += 1
                elif status == 'No':
                    bg_class = 'question-no'
                    category_no += 1
                    total_no += 1
                else:
                    bg_class = 'question-na'
                    category_na += 1
                    total_na += 1

                card_content += f"<div class='question {bg_class}'>"
                card_content += f"<h3>{question}</h3>\n\n"
                card_content += f"<p><span class='status'>{status}</span></p>\n\n<p><strong>Source:</strong> {source}</p>\n\n"
                
                if details.get('applicable_evaluations'):
                    card_content += "<details><summary>View Applicable Evaluations</summary>\n\n"
                    card_content += "<ul>"
                    for eval in details['applicable_evaluations']:
                        card_content += f"<li>{eval}</li>"
                    card_content += "</ul>\n"
                    card_content += "</details>\n\n"
                else:
                    card_content += "<details><summary>View Applicable Evaluations</summary>\n\n"
                    card_content += "<p>No applicable evaluations.</p>\n"
                    card_content += "</details>\n\n"
                
                card_content += "</div>"
            
            category_score = category_yes / (category_yes + category_no) * 100 if (category_yes + category_no) > 0 else 0
            card_content += f"<div class='category-score'>Category Score: {category_score:.2f}% (Yes: {category_yes}, No: {category_no}, N/A: {category_na})</div>"
            card_content += "</div>"
            all_cards_content += card_content

    all_cards_content += "</div>"
    
    total_score = total_yes / (total_yes + total_no) * 100 if (total_yes + total_no) > 0 else 0
    total_score_md = f"<div class='total-score'>Total Score: {total_score:.2f}% (Yes: {total_yes}, No: {total_no}, N/A: {total_na})</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)
    ]

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_chart = 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_chart), 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)
        scorecard_updates = update_detailed_scorecard(selected_model, 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()] + scorecard_updates

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", 
                                    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()
    
    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()