src/plots.py

"""
Visualization functions for DeathMath Leaderboard analytics.
Provides bar charts and radar plots for model comparison.
"""

import itertools as it
import random

import matplotlib.pyplot as plt
import numpy as np
import plotly.graph_objects as go
from matplotlib.figure import Figure

from src.leaderboard import build_leaderboard_df


def create_plot(selected_models: list[str]) -> Figure:
    """
    Create bar chart visualization comparing selected models.

    Args:
        selected_models: List of model names to display

    Returns:
        Matplotlib figure with bar chart comparing metrics across models
    """
    models_df = build_leaderboard_df()

    if not selected_models or models_df.empty:
        fig, ax = plt.subplots(figsize=(10, 6))
        ax.text(
            0.5,
            0.5,
            "Нет данных для отображения",
            horizontalalignment="center",
            verticalalignment="center",
            transform=ax.transAxes,
            fontsize=14,
        )
        ax.set_axis_off()
        return fig

    models_to_show = models_df[models_df["model"].isin(selected_models)]

    if models_to_show.empty:
        fig, ax = plt.subplots(figsize=(10, 6))
        ax.text(
            0.5,
            0.5,
            "Выбранные модели не найдены в данных",
            horizontalalignment="center",
            verticalalignment="center",
            transform=ax.transAxes,
            fontsize=14,
        )
        ax.set_axis_off()
        return fig

    fig, ax = plt.subplots(figsize=(12, 8))

    bar_width = 0.25
    models_count = len(models_to_show)
    indices = np.arange(models_count)
    colors = ["#1f77b4", "#ff7f0e", "#2ca02c"]

    ax.bar(indices - bar_width, models_to_show["math_score"], bar_width, label="RussianMath Score", color=colors[0])
    ax.bar(indices, models_to_show["physics_score"], bar_width, label="RussianPhysics Score", color=colors[1])
    ax.bar(indices + bar_width, models_to_show["score"], bar_width, label="Combined Score", color=colors[2])

    ax.set_xlabel("Модели")
    ax.set_ylabel("Баллы")
    ax.set_title("Сравнение производительности моделей на DeathMath benchmark")
    ax.set_xticks(indices)
    ax.set_xticklabels(models_to_show["model"], rotation=45, ha="right")
    ax.legend()
    ax.set_ylim(0, 1.0)
    ax.grid(axis="y", linestyle="--", alpha=0.7)

    plt.tight_layout()

    return fig


def create_radar_plot(selected_models: list[str]) -> go.Figure:
    """
    Create radar chart visualization comparing selected models.

    Args:
        selected_models: List of model names to display

    Returns:
        Plotly figure with interactive radar chart
    """
    models = build_leaderboard_df()
    metrics = ["math_score", "physics_score", "score"]
    metric_labels = ["RussianMath", "RussianPhysics", "Combined"]

    min_colour_distance_between_models = 100
    seed = 42

    def generate_colours(min_distance: int, seed: int) -> dict[str, tuple[int, int, int]]:
        """Generate distinct colors for each model."""
        colour_mapping = {}
        all_models = selected_models

        for i in it.count():
            min_colour_distance = min_distance - i
            retries_left = 10 * len(all_models)

            for model_id in all_models:
                random.seed(hash(model_id) + i + seed)
                r, g, b = 0, 0, 0
                too_bright, similar_to_other_model = True, True

                while (too_bright or similar_to_other_model) and retries_left > 0:
                    r, g, b = tuple(random.randint(0, 255) for _ in range(3))
                    too_bright = np.min([r, g, b]) > 200
                    similar_to_other_model = any(
                        np.abs(np.array(colour) - np.array([r, g, b])).sum() < min_colour_distance
                        for colour in colour_mapping.values()
                    )
                    retries_left -= 1

                colour_mapping[model_id] = (r, g, b)
            if len(colour_mapping) == len(all_models):
                break

        return colour_mapping

    colour_mapping = generate_colours(min_colour_distance_between_models, seed)
    fig = go.Figure()

    for _, model_data in models.iterrows():
        model_name = model_data["model"]
        if model_name not in selected_models:
            continue

        values = [model_data[metric] for metric in metrics]
        color = f"rgb{colour_mapping[model_name]}"

        fig.add_trace(
            go.Scatterpolar(
                r=values,
                theta=metric_labels,
                name=model_name,
                fill="toself",
                fillcolor=f"rgba{colour_mapping[model_name] + (0.6,)}",
                line={"color": color},
            )
        )

    fig.update_layout(
        polar={"radialaxis": {"visible": True, "range": [0, 1]}},
        showlegend=True,
        title="Сравнение моделей на DeathMath",
        template="plotly_dark",
    )

    return fig