'''
Copyright:

   Copyright © 2025 bdunahu <bdunahu@operationnull.com>

   Licensed under the Apache License, Version 2.0 (the "License");
   you may not use this file except in compliance with the License.
   You may obtain a copy of the License at

       http://www.apache.org/licenses/LICENSE-2.0

   Unless required by applicable law or agreed to in writing, software
   distributed under the License is distributed on an "AS IS" BASIS,
   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   See the License for the specific language governing permissions and
   limitations under the License.

Commentary:
Code:
'''
from typing import TYPE_CHECKING
if TYPE_CHECKING:
    import collections
import plotly.graph_objects as go
from plotly.subplots import make_subplots
import hashlib
import math

TRIM_PERCENT = 0.10

def get_color(name: str) -> str:
    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: "collections.defaultdict[str, collections.defaultdict[float, collections.defaultdict[int, list[tuple[list[tuple[str, float]], float]]]]]", output_file: str, input_file: str) -> None:
    # 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,
    )

    all_coros = set() # used to not add a coro to the legend twice
    for coro_name, x_values in results.items():
        for speedup, xx_values in x_values.items():
            for loop, experiments in xx_values.items():
                show_legend = coro_name not in all_coros
                all_coros.add(coro_name)
                col = loop_to_col[loop]

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

                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])

                    else:
                        latency = 0
                        y_latency_list.append(latency)

                    # 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)),
                    showlegend=show_legend,
                ), 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)),
                    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=True,
    )
    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)