import gradio as gr
import pandas as pd
import plotly.graph_objects as go
import numpy as np
import base64
from model_handler import (
    ModelHandler,
    TASK_CATEGORIES,
    TASK_DISPLAY_NAMES,
    TASK_METRICS,
)
from data_handler import (
    prepare_leaderboard,
    prepare_detailed_leaderboard,
    get_task_scores,
)

# CSS for styled HTML tables with merged headers (uses Gradio CSS variables)
TABLE_CSS = """
<style>
.detailed-table {
    width: 100%;
    border-collapse: collapse;
    font-size: 14px;
    margin: 10px 0;
    display: table !important;
    visibility: visible !important;
}
.detailed-table thead,
.detailed-table tbody,
.detailed-table tr {
    display: table-row-group;
    visibility: visible !important;
}
.detailed-table tr {
    display: table-row !important;
}
.detailed-table thead tr th {
    background-color: var(--background-fill-secondary) !important;
    color: var(--body-text-color) !important;
    font-weight: 600 !important;
    padding: 10px 8px !important;
    border: 1px solid var(--border-color-primary) !important;
    text-align: center !important;
    display: table-cell !important;
}
.detailed-table tbody tr td {
    padding: 8px 12px !important;
    text-align: center !important;
    border: 1px solid var(--border-color-primary) !important;
    background-color: var(--background-fill-primary) !important;
    color: var(--body-text-color) !important;
    display: table-cell !important;
    visibility: visible !important;
}
.detailed-table tbody tr:hover td {
    background-color: var(--background-fill-secondary) !important;
}
.detailed-table tbody td:first-child,
.detailed-table tbody td:nth-child(2) {
    text-align: left !important;
}
</style>
"""


def df_to_styled_html(df):
    """Convert DataFrame to styled HTML with CSS."""
    table_html = df.to_html(classes="detailed-table", border=1, index=False, na_rep="-")
    return TABLE_CSS + f'<div style="overflow-x: auto;">{table_html}</div>'


SCORE_COLUMNS = [
    "Average", "NER", "POS", "Reading", "Classification",
    "MCQA", "Generation", "Translation", "Exams", "Text Proc.", "MMLU",
]


def create_leaderboard_chart(leaderboard_df, score_col="Average"):
    """Create a horizontal bar chart for top 20 models by the chosen score column."""
    if leaderboard_df.empty or score_col not in leaderboard_df.columns:
        return None

    # Re-sort by chosen column and take top 20
    sorted_df = leaderboard_df.sort_values(score_col, ascending=False).head(20)
    top_20 = sorted_df.copy().sort_values(score_col, ascending=True)

    # Split into two traces for legend
    closed_mask = top_20['Size'] == '-'
    open_mask = ~closed_mask

    fig = go.Figure()
    for mask, label, color in [
        (closed_mask, "Closed Source", "#B54955"),
        (open_mask, "Open Source","#6D9AC7"),
    ]:
        subset = top_20[mask]
        fig.add_trace(go.Bar(
            y=subset["Model"],
            x=subset[score_col],
            orientation="h",
            name=label,
            marker_color=color,
            text=subset[score_col].round(4),
            textposition="outside",
            hovertemplate="<b>%{y}</b><br>" + score_col + ": %{x:.4f}<extra></extra>",
        ))

    fig.update_layout(
        title=f"Top 20 Models by {score_col}",
        xaxis_title=score_col,
        yaxis_title="Model",
        xaxis_range  = [0, 20] if score_col == "Exams" else [0, 1],
        height=600,
        width=None,
        hovermode="y unified",
        showlegend=True,
        legend=dict(x=1, y=1, xanchor="right", yanchor="top", bgcolor="rgba(0,0,0,0)", font=dict(color="black")),
        margin=dict(l=100),
        paper_bgcolor="white",
        plot_bgcolor="white",
        font=dict(color="black"),
        xaxis=dict(
            showgrid=True,
            gridwidth=1,
            gridcolor="rgba(200,200,200,0.15)",
            tickfont=dict(color="black"),
            title_font=dict(color="black"),
        ),
        yaxis=dict(
            dtick=1,
            tickfont=dict(color="black"),
            title_font=dict(color="black"),
            categoryorder = 'total ascending'
        )
    )

    return fig

def _load_logo_data_uri(company):
    """Load a company logo from logos/ folder and return as base64 data URI."""
    import os
    logos_dir = os.path.join(os.path.dirname(__file__), "logos")
    for ext, mime in [("png", "image/png"), ("svg", "image/svg+xml")]:
        path = os.path.join(logos_dir, f"{company}.{ext}")
        if os.path.exists(path):
            with open(path, "rb") as f:
                b64 = base64.b64encode(f.read()).decode()
            return f"data:{mime};base64,{b64}"
    return None


def create_spend_report_chart(leaderboard_df):
    """Create scatter plot of spend vs accuracy with company logos as markers."""
    try:
        spend_df = pd.read_csv("full_price_report.csv")
        spend_df.columns = ["Model", "Total Tokens", "Total Cost ($)"]
    except FileNotFoundError:
        return None

    if leaderboard_df.empty or spend_df.empty:
        return None

    # Merge leaderboard and spend data
    merged = spend_df.merge(
        leaderboard_df[["Model", "Average", "Size"]],
        on="Model",
        how="inner"
    )

    if merged.empty:
        return None

    # Extract company name before stripping prefix
    merged["Company"] = merged["Model"].str.split("/").str[0]
    merged["Model"] = merged["Model"].str.split("/").str[-1]
    # print('[MODELS]',merged["Model"].tolist())
    # Pre-load logos (one per company)
    logo_cache = {}
    for company in merged["Company"].unique():
        uri = _load_logo_data_uri(company)
        if uri:
            logo_cache[company] = uri

    fig = go.Figure()
    fig.add_trace(go.Scatter(
        x=merged["Total Cost ($)"],
        y=merged["Average"],
        mode="markers+text",
        marker=dict(size=18, color="rgba(0,0,0,0)", line=dict(width=0)),
        text=merged["Model"],
        textposition=['middle right', 'middle left', 'middle left', 'middle left', 'middle left', 'middle left', 'middle right', 'middle right', 'middle right', 'middle right', 'middle right', 'middle left', 'middle right', 'bottom left', 'middle left', 'middle left', 'middle right'],
        textfont=dict(size=11, color="black"),
        hovertemplate="<b>%{text}</b><br>Cost: $%{x:.2f}<br>Avg Score: %{y:.4f}<extra></extra>",
        showlegend=False,
    ))

    fig.update_layout(
        title="Model Evaluation: Cost vs Accuracy<br><sub>Hover or read labels for full details</sub>",
        paper_bgcolor="white",
        plot_bgcolor="white",
        font=dict(color="black"),
        height=650,
        width=None,
        autosize=True,
        xaxis_title="Total Cost ($)",
        yaxis=dict(range=[0, 1],
                    showgrid=False,
                    anchor="free",
                    position=0,
                    side="left",
                    showline=True),
        showlegend=False,
        margin=dict(l=40, r=40, t=80, b=40),
    )

    fig.update_xaxes(
        showgrid=False,
        tickprefix="$",
        tickmode="auto",
        autorange="reversed",
        zeroline=False,
    )

    # Build logo data for JS overlay
    import json
    logo_points = []
    for _, row in merged.iterrows():
        uri = logo_cache.get(row["Company"])
        if uri:
            logo_points.append({
                "x": row["Total Cost ($)"],
                "y": row["Average"],
                "src": uri,
            })

    chart_div = fig.to_html(include_plotlyjs=False, full_html=False, div_id="spend-chart")
    logo_size = 13

    inner_html = f"""
<!DOCTYPE html>
<html><head>
<script src="https://cdn.plot.ly/plotly-2.27.0.min.js"></script>
<style>html,body{{margin:0;padding:0;width:100%;height:100%;overflow:hidden;background:white;}}</style>
</head><body>
{chart_div}
<script>
(function() {{
  var logos = {json.dumps(logo_points)};
  var logoSize = {logo_size};

  function positionLogos() {{
    var gd = document.getElementById('spend-chart');
    if (!gd || !gd._fullLayout) return;
    var xa = gd._fullLayout.xaxis;
    var ya = gd._fullLayout.yaxis;
    var plotRect = gd.querySelector('.nsewdrag');
    if (!plotRect) return;
    var gdRect = gd.getBoundingClientRect();
    var pRect = plotRect.getBoundingClientRect();
    var offsetX = pRect.left - gdRect.left;
    var offsetY = pRect.top - gdRect.top;

    gd.querySelectorAll('.logo-overlay').forEach(function(el) {{ el.remove(); }});

    logos.forEach(function(pt) {{
      var px = xa.l2p(xa.d2l(pt.x)) + offsetX;
      var py = ya.l2p(ya.d2l(pt.y)) + offsetY;
      var img = document.createElement('img');
      img.className = 'logo-overlay';
      img.src = pt.src;
      img.style.cssText = 'position:absolute;pointer-events:none;' +
        'width:' + logoSize + 'px;height:' + logoSize + 'px;' +
        'left:' + (px - logoSize/2) + 'px;top:' + (py - logoSize/2) + 'px;' +
        'border-radius:50%;';
      gd.style.position = 'relative';
      gd.appendChild(img);
    }});
  }}

  var checkReady = setInterval(function() {{
    var gd = document.getElementById('spend-chart');
    if (gd && gd._fullLayout) {{
      clearInterval(checkReady);
      positionLogos();
      gd.on('plotly_relayout', positionLogos);
      gd.on('plotly_afterplot', positionLogos);
      new ResizeObserver(positionLogos).observe(gd);
    }}
  }}, 100);
}})();
</script>
</body></html>
"""

    import html as html_mod
    escaped = html_mod.escape(inner_html)
    return f'<iframe srcdoc="{escaped}" style="width:100%;height:700px;border:none;"></iframe>'


def generate_about_markdown(task_scores: dict) -> str:
    """Generate About section markdown with task scores and metrics."""
    # Create mapping of display name to metric
    display_to_metric = {}
    for task_key, display_name in TASK_DISPLAY_NAMES.items():
        metric = TASK_METRICS.get(task_key, "unknown")
        display_to_metric[display_name] = metric

    # Create metric descriptions
    metric_descriptions = {
        "ner_accuracy": "Accuracy",
        "ud_pos_regex_acc": "Accuracy",
        "bleu": "BLEU",
        "exact_match_mcqa": "Accuracy",
        "armenian_exam_score": "Custom Exam Score (0-20)",
        "armenian_mmlu_pro_score": "Accuracy",
        "punctuation_accuracy": "Accuracy",
        "space_accuracy": "Accuracy",
    }

    markdown = """# About ArmBench-LLM

ArmBench-LLM is a comprehensive benchmark for evaluating large language models on Armenian language tasks.

## Task Categories

"""
    for category, task_keys in TASK_CATEGORIES.items():
        markdown += f"- **{category}**\n"
        for task_key in task_keys:
            display_name = TASK_DISPLAY_NAMES.get(task_key, task_key)
            metric = display_to_metric.get(display_name, "unknown")
            metric_desc = metric_descriptions.get(metric, metric)
            markdown += f"  - {display_name} ({metric_desc})\n"

    markdown += """

## Submission Guide

To submit your model for evaluation:

1. **Evaluate your model** using lighteval with Armenian tasks at [GitHub](https://github.com/Metric-AI-Lab/lighteval)

2. **Save your results.json** - the lighteval output format is automatically parsed

3. **Add the tag and results**:
   - Add the `ArmBench-LLM` tag to your model card
   - Upload `results.json` to your model repository

4. Click "Refresh Data" to see your results on the leaderboard

## Contributing

You can contribute to this benchmark in several ways:

- Provide API credits for evaluating additional API-based models.
- Citing our work in your research and publications.
- Contributing to the development of the benchmark itself with data or with evaluation results.

## Citation

If you use ArmBench-LLM in your research, please cite it as follows:

```bibtex
@misc{armbench-llm,
  title={ArmBench-LLM: Benchmarking LLMs on Armenian Language Tasks},
  author={Metric-AI-Lab},
  year={2026},
  howpublished={\\url{https://huggingface.co/spaces/Metricam/ArmBench-LLM}},
  note={Benchmark for evaluating large language models on Armenian language tasks}
}
```

## About Metric

Metric is an AI Research Lab in Yerevan, Armenia. Contact: info@metric.am

*This is a non-commercial research project.*
"""
    return markdown


# Global state
global_data = {}


def load_spend_report():
    """Load spend report from CSV file with selected columns."""
    try:
        df = pd.read_csv("full_price_report.csv")
        # Select and rename columns
        df.columns = ["Model", "Total Tokens", "Total Cost ($)"]
        # Round numeric columns for better readability
        df["Total Tokens"] = df["Total Tokens"].round(0).astype(int)
        df["Total Cost ($)"] = df["Total Cost ($)"].round(4)
        # Sort by total cost ascending
        df = df.sort_values(by="Total Cost ($)", ascending=True).reset_index(drop=True)
        return df
    except FileNotFoundError:
        return pd.DataFrame({"Error": ["full_price_report.csv not found"]})
    except Exception as e:
        return pd.DataFrame({"Error": [str(e)]})


def refresh_data():
    global global_data
    model_handler = ModelHandler()

    df = model_handler.get_llm_benchmark_data()

    detailed_results = model_handler.get_detailed_results()

    leaderboard = prepare_leaderboard(df)
   
    task_scores = get_task_scores(detailed_results)

    global_data = {
        "leaderboard": leaderboard,
        "detailed": prepare_detailed_leaderboard(detailed_results, leaderboard),
        "task_scores": task_scores,
    }
    return (
        create_leaderboard_chart(global_data["leaderboard"]),
        global_data["leaderboard"],
        df_to_styled_html(global_data["detailed"]),
        create_spend_report_chart(global_data["leaderboard"]),
    )


def main():
    global global_data

    model_handler = ModelHandler()

    df = model_handler.get_llm_benchmark_data()

    detailed_results = model_handler.get_detailed_results()

    leaderboard = prepare_leaderboard(df)

    task_scores = get_task_scores(detailed_results)

    global_data = {
        "leaderboard": leaderboard,
        "detailed": prepare_detailed_leaderboard(detailed_results, leaderboard),
        "task_scores": task_scores,
    }

    with gr.Blocks(title="ArmBench-LLM", theme=gr.themes.Base(),fill_width=True) as demo:
        gr.Markdown("# ArmBench-LLM: Benchmarking LLMs on Armenian Language Tasks")
        gr.Markdown(
            """
            Evaluating large language models on Armenian language tasks.
            Developed by [Metric](https://metric.am/).
            """
        )

        with gr.Tabs():
            with gr.TabItem("Leaderboard"):
                gr.Markdown("## Leaderboard")
                gr.Markdown(
                    """
                    **Task Categories:**
                    - **NER**: Named Entity Recognition (FiNER, PioNER)
                    - **POS**: Part-of-Speech Tagging
                    - **Reading Comprehension**: Reading Comprehension (SQuAD, Belebele, DREAM, Hartak, MS MARCO)
                    - **Classification**: Text Classification (Topic-14, Sentiment)
                    - **MCQA**: Multiple Choice QA (INCLUDE, Syndarin, Scientific)
                    - **Generation**: Text Generation (Email Summary, Conversation Summary, Simple QA, Paraphrase)
                    - **Translation**: English-Armenian Translation
                    - **Exams**: Armenian Unified Exams (Math, Literature, History)
                    - **Text Proc.**: Text Processing (Punctuation, Space Fix)
                    - **MMLU**: MMLU-Pro Armenian
                    """
                )

                leaderboard_table = gr.DataFrame(
                    value=global_data["leaderboard"], label="LLM Leaderboard"
                )

                gr.Markdown(
                    "> **Note:** We observed significant reliability issues with Claude 4.6 and 4.5 during evaluation. "
                    "To preserve the integrity of the benchmark, all Claude 4.6 and 4.5 models have been excluded from the leaderboard. "
                    "You can learn more about the issue in [this GitHub issue](https://github.com/anthropics/claude-code/issues/17357)."
                )

                score_dropdown = gr.Dropdown(
                    choices=SCORE_COLUMNS,
                    value="Average",
                    label="Score to display",
                )

                with gr.Row():
                    leaderboard_chart = gr.Plot(
                        value=create_leaderboard_chart(global_data["leaderboard"]),
                        container=False,
                    )

                def _update_chart(score_col):
                    return create_leaderboard_chart(global_data["leaderboard"], score_col)

                score_dropdown.change(
                    fn=_update_chart,
                    inputs=[score_dropdown],
                    outputs=[leaderboard_chart],
                )

                with gr.Accordion("Detailed Scores", open=False):
                    detailed_table = gr.HTML(
                        value=df_to_styled_html(global_data["detailed"])
                    )

            with gr.TabItem("About"):
                gr.Markdown(generate_about_markdown(global_data.get("task_scores", {})))

                gr.Image("logo.png", width=200, show_label=False)

            with gr.TabItem("Spend Report"):
                gr.Markdown("## Model Evaluation Spend Report")
                gr.Markdown(
                    """
                    Cost analysis for evaluating each model on ArmBench-LLM tasks.

                    **Notes:**
                    - Token counts are approximations based on API usage patterns
                    - Spend data is reported only for closed models as well as open-source models having more than or equal to 30B parameters
                    """
                )

                gr.DataFrame(value=load_spend_report())

                with gr.Row():
                    spend_chart = gr.HTML(
                        value=create_spend_report_chart(global_data["leaderboard"]),
                    )

        refresh_button = gr.Button("Refresh Data")
        refresh_button.click(
            fn=refresh_data,
            outputs=[
                leaderboard_chart,
                leaderboard_table,
                detailed_table,
                spend_chart,
            ],
        )
    demo.launch(server_name="0.0.0.0", server_port=7860, ssr_mode=False)


if __name__ == "__main__":
    main()