import numpy as np
import matplotlib.pyplot as plt
import pandas as pd

# Графики

def plot_all(df: pd.DataFrame):
    structures = ["BST", "HT", "LL"]
    operations = ["Вставка", "Поиск", "Удаление"]
    op_keys    = ["Вставка СРЕДНЕЕ", "Поиск СРЕДНЕЕ", "Удаление СРЕДНЕЕ"]
    filenames  = ["./insert_plot.png", "./search_plot.png", "./delete_plot.png"]

    avg = df[df["Operation"].isin(op_keys)]

    data = {
        (row.Structure, row.Mode, row.Operation): row.Time
        for row in avg.itertuples(index=False)
    }

    x = np.arange(len(structures))
    width = 0.35

    for op_label, op_key, filename in zip(operations, op_keys, filenames):
        fig, ax = plt.subplots(figsize=(8, 5))

        times_shuffled = [data.get((s, "shufled", op_key), float('nan')) for s in structures]
        times_sorted   = [data.get((s, "sorted",  op_key), 0) for s in structures]

        bars1 = ax.bar(x - width/2, times_shuffled, width, label="Случайный",       color="steelblue")
        bars2 = ax.bar(x + width/2, times_sorted,   width, label="Отсортированный",  color="red")

        for bar in bars1:
            height = bar.get_height()
            ax.text(bar.get_x() + bar.get_width() / 2, height,
                    f"{height:.6f}", ha="center", va="bottom", fontsize=8)

        for bar in bars2:
            height = bar.get_height()
            ax.text(bar.get_x() + bar.get_width() / 2, height,
                    f"{height:.6f}", ha="center", va="bottom", fontsize=8)

        ax.set_yscale("log")
        ax.set_title(op_label)
        ax.set_ylabel("Время (сек)")
        ax.set_xticks(x)
        ax.set_xticklabels(structures)
        ax.legend()
        ax.grid(axis="y", linestyle="--", alpha=0.7)

        plt.tight_layout()
        plt.savefig(filename)
        print(f"График сохранён: {filename}")

def main_plot():
    results = pd.read_csv("./results.csv")
    plot_all(results)

if __name__ == "__main__":
    main_plot()