import plotly.graph_objects as go
from plotly.subplots import make_subplots
import hashlib
import math

TRIM_PERCENT = 0.10

def get_color(name):
    hash_object = hashlib.md5(name.encode())
    color_index = int(hash_object.hexdigest(), 16) % 360
    return f'hsl({color_index}, 90%, 30%)'

def plot_results(results, output_file, input_file):
    # determine the number of loops we have data for
    total_loops = set()
    for x_values in results.values():
        for xx_values in x_values.values():
            total_loops.update(xx_values.keys())
    total_loops = sorted(total_loops)
    num_loops = len(total_loops)
    loop_to_col = {loop: idx + 1 for idx, loop in enumerate(total_loops)}

    fig = make_subplots(
        rows=4,
        cols=num_loops,
        subplot_titles=[f"{loop}" for loop in total_loops] * 4,
        vertical_spacing=0.1,
        horizontal_spacing=0.05,
        shared_xaxes=True,
        shared_yaxes=False,
    )

    for coro_name, x_values in results.items():
        for speedup, xx_values in x_values.items():
            for loop, experiments in xx_values.items():
                col = loop_to_col[loop]

                x_list = []
                y_latency_list = []
                y_throughput_list = []
                y_max_latency_list = []
                y_num_callbacks_list = []
                latency_hover_text = []
                max_latency_hover_text = []

                for experiment in experiments:
                    completed_callbacks = experiment[0]
                    virtual_run_time = experiment[1]

                    x_val = speedup * 100
                    x_list.append(x_val)

                    num_callbacks = len(completed_callbacks)
                    y_num_callbacks_list.append(num_callbacks)

                    # handle average latency graph
                    if num_callbacks > 0:

                        trim_count = math.floor(num_callbacks * TRIM_PERCENT / 2)
                        sorted_callbacks = sorted(completed_callbacks, key=lambda cb: cb[1])

                        trimmed_callbacks = (
                            sorted_callbacks[trim_count: len(sorted_callbacks) - trim_count]
                            if trim_count > 0 else sorted_callbacks
                        )

                        trimmed_latencies = [cb[1] for cb in trimmed_callbacks]
                        latency = sum(trimmed_latencies) / len(trimmed_latencies)

                        y_latency_list.append(latency)

                        max_cb = max(completed_callbacks, key=lambda cb: cb[1])
                        y_max_latency_list.append(max_cb[1])

                        breakdown = "<br>".join([f"  {cb[0]}: {round(cb[1], 4)}" for cb in trimmed_callbacks])
                        latency_hover_text.append(
                            f"{coro_name}<br>Speedup: {speedup}<br>Trimmed Average Wait ({int(TRIM_PERCENT*100)}%): {round(latency, 4)}<br>Breakdown:<br>{breakdown}"
                        )
                        max_latency_hover_text.append(
                            f"{coro_name}<br>Speedup: {speedup}<br>Max Wait: {round(max_cb[1], 4)}<br>Handle: {max_cb[0]}"
                        )

                    else:
                        latency = 0
                        y_latency_list.append(latency)

                        latency_hover_text.append(f"{coro_name}<br>Speedup: {speedup}<br>No callbacks")
                        max_latency_hover_text.append(f"{coro_name}<br>Speedup: {speedup}<br>No callbacks")

                    # handle throughput graph
                    throughput = num_callbacks / virtual_run_time if virtual_run_time else 0
                    y_throughput_list.append(throughput)

                fig.add_trace(go.Scatter(
                    x=x_list,
                    y=y_latency_list,
                    mode='markers',
                    name=coro_name,
                    marker=dict(color=get_color(coro_name)),
                    hovertext=latency_hover_text,
                    showlegend=True,
                ), row=1, col=col)

                fig.add_trace(go.Scatter(
                    x=x_list,
                    y=y_throughput_list,
                    mode='markers',
                    name=coro_name,
                    marker=dict(color=get_color(coro_name)),
                    showlegend=False,
                ), row=2, col=col)

                fig.add_trace(go.Scatter(
                    x=x_list,
                    y=y_max_latency_list,
                    mode='markers',
                    name=coro_name,
                    marker=dict(color=get_color(coro_name)),
                    hovertext=max_latency_hover_text,
                    showlegend=False,
                ), row=3, col=col)

                fig.add_trace(go.Scatter(
                    x=x_list,
                    y=y_num_callbacks_list,
                    mode='markers',
                    name=coro_name,
                    marker=dict(color=get_color(coro_name)),
                    showlegend=False,
                ), row=4, col=col)

    fig.update_layout(
        height=1080,
        width=1920 * num_loops,
        title_text=f"Coroutine Performance Metrics: {input_file}",
        showlegend=False,
    )
    for col in range(1, num_loops + 1):
        fig.update_xaxes(title_text="speedup (% optimized away)", row=4, col=col)
        fig.update_xaxes(showticklabels=True, col=col)
        fig.update_xaxes(showticklabels=True, row=2, col=col)
        fig.update_xaxes(showticklabels=True, row=3, col=col)
    fig.update_yaxes(title_text="average latency (seconds)", row=1, col=1)
    fig.update_yaxes(title_text="throughput (handles per second)", row=2, col=1)
    fig.update_yaxes(title_text="maximum latency (seconds)", row=3, col=1)
    fig.update_yaxes(title_text="# of callbacks", row=4, col=1)

    fig.write_html(output_file)