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, dtype=float)  # Ensure values are floats
        
    n = len(values)
    pad_width = rows * cols - n
    
    # Use masked array to handle NaN values
    masked_values = np.ma.array(values)
    masked_values = np.ma.pad(masked_values, (0, pad_width), 'constant', constant_values=np.ma.masked)
    masked_values = masked_values.reshape((rows, cols))

    keys = np.pad(keys, (0, pad_width), 'constant', constant_values='').reshape((rows, cols))

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

    # Use compressed array to count non-masked (finite) values
    upper_solve_rate = round(np.count_nonzero(~masked_values.mask) / n * 100, 2)
    
    fig = go.Figure(data=go.Heatmap(
        z=masked_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,
    )
    
    return fig