plotting

def boxplot_by_dimension(
    result: Result, dimension: str, name: str | None = None, **box_kwargs: Any
) -> go.Box:
    """
    Create a box plot for a given dimension from responses in a Run Result.

    Args:
        result (Result): The Result object containing the response data
        dimension (str): The dimension to plot as histogram values
        **box_kwargs: Additional keyword arguments to pass to go.Box

    Returns:
        plotly.graph_objects.Box: Box plot of the specified dimension

    Example:
        >>> result = Result(...)
        >>> fig = boxplot_by_dimension(result, "time_to_first_token")
        >>> fig.show()
    """
    x = result.get_dimension(dimension)

    return go.Box(x=x, name=name or result.run_name, **box_kwargs)

def histogram_by_dimension(
    result: Result,
    dimension: str,
    filter_dimension: str | None = None,
    filter_value: Any = None,
    **histogram_kwargs: Any,
) -> go.Histogram:
    """
    Create a histogram plot for a given dimension from the result data.

    Args:
        result (Result): The Result object containing the response data
        dimension (str): The dimension to plot as histogram values
        filter_dimension (str | None, optional): Dimension to filter the data on. Defaults to None.
        filter_value (Any, optional): Value to filter the filter_dimension by. Defaults to None.
        **histogram_kwargs: Additional keyword arguments to pass to go.Histogram

    Returns:
        plotly.graph_objects.Histogram: Histogram plot of the specified dimension

    Example:
        >>> result = Result(...)
        >>> fig = histogram_by_dimension(result, "time_to_first_token")
        >>> fig.show()
    """
    x = result.get_dimension(
        dimension, filter_dimension=filter_dimension, filter_value=filter_value
    )
    name = histogram_kwargs.pop("name", None) or result.run_name
    histnorm = histogram_kwargs.pop("histnorm", None) or "probability"

    return go.Histogram(x=x, name=name, histnorm=histnorm, **histogram_kwargs)

def plot_heatmap(
    result: Result,
    dimension: str,
    n_bins_x: int | None,
    n_bins_y: int | None,
    output_path: Path | None = None,
    show_scatter=False,
) -> go.Figure:
    """
    Generate a heatmap visualization of token counts and a specified dimension.

    Args:
        result (Result): The Result object containing the response data
        dimension (str): The dimension to plot as the z-axis/color values
        n_bins_x (int): Number of bins for the x-axis histogram
        n_bins_y (int): Number of bins for the y-axis histogram
        show_scatter (bool, optional): Whether to overlay scatter points. Defaults to False.

    Returns:
        plotly.graph_objs._figure.Figure: The generated heatmap figure

    Raises:
        ValueError: If the specified dimension is not found in the result data
    """

    num_tokens_input_list = result.get_dimension("num_tokens_input")
    num_tokens_output_list = result.get_dimension("num_tokens_output")
    try:
        z = result.get_dimension(dimension)
    except AttributeError:
        raise ValueError(f"Dimension {dimension} not found in result")

    fig = px.density_heatmap(
        x=num_tokens_input_list,
        y=num_tokens_output_list,
        z=z,
        marginal_x="histogram",
        marginal_y="histogram",
        template="plotly_white",
        histfunc="avg",
        nbinsx=n_bins_x,
        nbinsy=n_bins_y,
        labels={
            "x": "Number of input tokens",
            "y": "Number of output tokens",
        },
        title=f'LLMeter: average {dimension.replace("_", " ").capitalize()}',
        width=800,
        height=800,
    )
    fig.layout["coloraxis"]["colorbar"]["title"] = "average"  # type: ignore
    fig.update_coloraxes(colorbar_tickformat=".1s")
    fig.update_coloraxes(colorbar_ticksuffix="s")
    fig.update_layout(coloraxis_colorscale="oryel")
    fig.update_traces(marker_color="grey", selector=dict(type="histogram"))
    fig.update_traces(texttemplate="%{z:.3s}s <br> %{x:.2s}, %{y:.2s}")
    fig.update_traces(
        texttemplate="%{y:.3s}", selector=dict(type="histogram", bingroup="x")
    )
    fig.update_traces(
        texttemplate="%{x:.3s}", selector=dict(type="histogram", bingroup="y")
    )
    if show_scatter:
        fig.add_trace(
            go.Scatter(
                x=num_tokens_input_list,
                y=num_tokens_output_list,
                mode="markers",
                showlegend=False,
                marker=dict(
                    symbol="x",
                    opacity=0.7,
                    color="white",
                    size=8,
                    line=dict(width=1),
                ),
            )
        )
    return fig

def plot_load_test_results(
    load_test_result: LoadTestResult,
    log_scale=True,
) -> dict[str, go.Figure]:
    """
    Generate a collection of plots visualizing different metrics from a load test result.

    Args:
        load_test_result (LoadTestResult): The load test result object containing the data to plot
        log_scale (bool, optional): Whether to use logarithmic scale for the plots. Defaults to True.

    Returns:
        dict: Dictionary containing the following plots:
            - time_to_first_token: Figure showing time to first token vs number of clients
            - time_to_last_token: Figure showing time to last token vs number of clients
            - requests_per_minute: Figure showing requests per minute vs number of clients
            - error_rate: Figure showing error rate vs number of clients
            - average_input_tokens_clients: Figure showing average input tokens per minute vs clients
            - average_output_tokens_clients: Figure showing average output tokens per minute vs clients
    """
    f1 = latency_clients_fig(
        load_test_result, "time_to_first_token", log_scale=log_scale
    )
    f2 = latency_clients_fig(
        load_test_result, "time_to_last_token", log_scale=log_scale
    )
    f3 = rpm_clients_fig(load_test_result, log_scale=log_scale)
    f4 = error_clients_fig(load_test_result, log_scale=log_scale)
    f5 = average_input_tokens_clients_fig(load_test_result, log_scale=log_scale)
    f6 = average_output_tokens_clients_fig(load_test_result, log_scale=log_scale)

    return {
        "time_to_first_token": f1,
        "time_to_last_token": f2,
        "requests_per_minute": f3,
        "error_rate": f4,
        "average_input_tokens_clients": f5,
        "average_output_tokens_clients": f6,
    }

def scatter_histogram_2d(
    result: Result, x_dimension: str, y_dimension: str, n_bins_x: int, n_bins_y: int
) -> go.Figure:
    """
    Generate a scatter plot with histograms for the given dimensions.

    Args:
        result (Result): The Result object containing the data.
        x_dimension (str): The dimension to be plotted on the x-axis.
        y_dimension (str): The dimension to be plotted on the y-axis.
        n_bins_x (int): The number of bins for the x-axis histogram.
        n_bins_y (int): The number of bins for the y-axis histogram.

    Returns:
        plotly.graph_objs._figure.Figure: The generated scatter plot figure.
    """

    try:
        x = result.get_dimension(x_dimension)
        y = result.get_dimension(y_dimension)
    except AttributeError:
        raise ValueError(f"Invalid dimension: {x_dimension} or {y_dimension}")

    fig = px.scatter(
        x=x,
        y=y,
        marginal_x="histogram",
        marginal_y="histogram",
        template="plotly_white",
        range_x=[0, max(x for x in x if x is not None) * 1.1],
        range_y=[0, max(y for y in y if y is not None) * 1.1],
        labels={
            "x": x_dimension.replace("_", " ").capitalize(),
            "y": y_dimension.replace("_", " ").capitalize(),
        },
        width=800,
        height=800,
    )
    fig.update_traces(nbinsx=n_bins_x, autobinx=True, selector={"type": "histogram"})
    fig.update_traces(nbinsy=n_bins_y, autobinx=True, selector={"type": "histogram"})
    return fig