from time import perf_counter
import numpy as np
import matplotlib.pyplot as plt
from optimum.amd.ryzenai.models import YoloV8ImageProcessor


def measure_latency(model, image):
    processor = YoloV8ImageProcessor()
    pixel_values = processor(image, return_tensors="np").pixel_values

    latencies = []

    # warm up
    for _ in range(10):
        _ = model(pixel_values)
    # Timed run
    for _ in range(10):
        start_time = perf_counter()
        _ = model(pixel_values)
        latency = perf_counter() - start_time
        latencies.append(latency)
    # Compute run statistics
    time_avg_ms = 1000 * np.mean(latencies)
    time_std_ms = 1000 * np.std(latencies)
    return time_avg_ms, time_std_ms


def plot_latency(cpu_latency, ryzen_latency):
    """
    Creates a bar plot comparing CPU and Ryzen latencies.
    
    Parameters:
    cpu_latency (float): Average latency for the CPU.
    ryzen_latency (float): Average latency for the Ryzen.
    """
    
    # Labels for the bars
    labels = ['RyzenAI', 'CPU']
    
    # Latency values
    latency_values = [ryzen_latency, cpu_latency]
    
    # Colors for the bars
    colors = ['blue', 'green']
    
    # Create a bar plot
    x = np.arange(len(labels))  # the label locations
    width = 0.35  # the width of the bars
    
    fig, ax = plt.subplots()
    bars = ax.bar(x, latency_values, width, label='Latency', color=colors)
    
    # Add some text for labels, title and custom x-axis tick labels, etc.
    ax.set_xlabel('Processor')
    ax.set_ylabel('Average Latency (ms)')
    ax.set_title('Average Latency by Processor')
    ax.set_xticks(x)
    ax.set_xticklabels(labels)
    ax.legend()
    
    # Attach a text label above each bar in *bars*, displaying its height.
    for bar in bars:
        height = bar.get_height()
        ax.annotate(f'{height:.2f}',
                    xy=(bar.get_x() + bar.get_width() / 2, height),
                    xytext=(0, 3),  # 3 points vertical offset
                    textcoords="offset points",
                    ha='center', va='bottom')
    
    fig.tight_layout()
    
    plt.show()