import plotly.graph_objects as go
import plotly.express as px
import numpy as np


def plot_elo_mle(df):
    fig = px.scatter(df, x="model", y="rating", error_y="error_y",
                     error_y_minus="error_y_minus",
                    #  title="Bootstrap of Elo MLE Estimates (BigCodeBench-Complete)"
                     )
    fig.update_layout(xaxis_title="Model", 
                      yaxis_title="Rating",
                      autosize=True,
                    #   width=1300,
                    #   height=900,
                      )
    return fig


def plot_solve_rate(df, task, rows=30, cols=38):
    keys = df["task_id"]
    values = df["solve_rate"]
    
    values = np.array(values)
        
    n = len(values)
    values = np.pad(values, (0, rows * cols - n), 'constant', constant_values=np.nan).reshape((rows, cols))
    keys = np.pad(keys, (0, rows * cols - n), 'constant', constant_values='').reshape((rows, cols))

    hover_text = np.empty_like(values, dtype=object)
    for i in range(rows):
        for j in range(cols):
            if not np.isnan(values[i, j]):
                hover_text[i, j] = f"{keys[i, j]}<br>Solve Rate: {values[i, j]:.2f}"
            else:
                hover_text[i, j] = "NaN"

    upper_solve_rate = round(np.count_nonzero(values)/n*100, 2)
    fig = go.Figure(data=go.Heatmap(
        z=values,
        text=hover_text,
        hoverinfo='text',
        colorscale='teal',
        zmin=0,
        zmax=100
    ))

    fig.update_layout(
        title=f'BigCodeBench-{task}<br><i>Lowest Upper Limit: {upper_solve_rate}%</i>',
        xaxis_nticks=cols,
        yaxis_nticks=rows,
        xaxis=dict(showticklabels=False),
        yaxis=dict(showticklabels=False),
        autosize=True,
        # width=760,
        # height=600,
    )
    
    return fig