"""
plot_results.py - строит столбчатые диаграммы по итогам экспериментов.
"""
import csv
import os
import matplotlib.pyplot as plt
import numpy as np

CSV = os.path.join("docs", "data", "results.csv")
PLOTS_DIR = os.path.join("docs", "data", "plots")
os.makedirs(PLOTS_DIR, exist_ok=True)


def load_means():
    """Возвращает dict[(structure, mode, op)] = (mean, N)."""
    means = {}
    with open(CSV, encoding="utf-8") as f:
        rows = list(csv.reader(f))
    # ищем секцию "--- СРЕДНИЕ ---"
    start = None
    for i, row in enumerate(rows):
        if row and row[0] == "--- СРЕДНИЕ ---":
            start = i + 2  # пропустить заголовок секции
            break
    for row in rows[start:]:
        if not row:
            continue
        structure, mode, op, n, mean, _trials = row
        means[(structure, mode, op)] = (float(mean), int(n))
    return means


def plot_grouped(means, op, fname, title):
    """Сгруппированные столбцы: для каждой структуры - два столбца (shuffled / sorted)."""
    structures = ["LinkedList", "HashTable", "BST"]
    modes = ["shuffled", "sorted"]
    x = np.arange(len(structures))
    width = 0.36

    fig, ax = plt.subplots(figsize=(8, 5))

    for i, mode in enumerate(modes):
        vals_ms = []
        labels = []
        for s in structures:
            mean, n = means[(s, mode, op)]
            vals_ms.append(mean * 1000)
            labels.append(f"N={n}")
        bars = ax.bar(x + (i - 0.5) * width, vals_ms, width,
                      label=mode, alpha=0.85)
        for bar, lab in zip(bars, labels):
            h = bar.get_height()
            ax.text(bar.get_x() + bar.get_width() / 2, h,
                    f"{h:.2f}\n{lab}", ha="center", va="bottom", fontsize=8)

    ax.set_xticks(x)
    ax.set_xticklabels(structures)
    ax.set_ylabel("Время, мс (среднее по 5 запускам)")
    ax.set_title(title)
    ax.set_yscale("log")
    ax.legend(title="порядок входных данных")
    ax.grid(axis="y", linestyle="--", alpha=0.4)
    plt.tight_layout()
    out = os.path.join(PLOTS_DIR, fname)
    plt.savefig(out, dpi=130)
    plt.close()
    print("saved:", out)


def plot_bst_degradation(means, fname):
    """Отдельный график: BST shuffled vs sorted - даже при меньшем N
    отсортированные данные дают намного большее время."""
    ops = ["insert", "find", "delete"]
    shuffled = [means[("BST", "shuffled", op)][0] * 1000 for op in ops]
    sorted_  = [means[("BST", "sorted",   op)][0] * 1000 for op in ops]
    n_shuf   = means[("BST", "shuffled", "insert")][1]
    n_sort   = means[("BST", "sorted",   "insert")][1]

    x = np.arange(len(ops))
    width = 0.36
    fig, ax = plt.subplots(figsize=(7, 4.5))
    ax.bar(x - width/2, shuffled, width, label=f"shuffled (N={n_shuf})")
    ax.bar(x + width/2, sorted_,  width, label=f"sorted (N={n_sort})")
    for i, v in enumerate(shuffled):
        ax.text(i - width/2, v, f"{v:.2f}", ha="center", va="bottom", fontsize=9)
    for i, v in enumerate(sorted_):
        ax.text(i + width/2, v, f"{v:.2f}", ha="center", va="bottom", fontsize=9)

    ax.set_xticks(x); ax.set_xticklabels(ops)
    ax.set_ylabel("Время, мс")
    ax.set_title("BST: деградация на отсортированных данных\n"
                 "(N для sorted в 5 раз меньше, но время больше)")
    ax.set_yscale("log")
    ax.legend()
    ax.grid(axis="y", linestyle="--", alpha=0.4)
    plt.tight_layout()
    out = os.path.join(PLOTS_DIR, fname)
    plt.savefig(out, dpi=130)
    plt.close()
    print("saved:", out)


def main():
    means = load_means()
    plot_grouped(means, "insert",
                 "insert_compare.png",
                 "Вставка всех записей (лог. шкала)")
    plot_grouped(means, "find",
                 "find_compare.png",
                 "Поиск 110 ключей (лог. шкала)")
    plot_grouped(means, "delete",
                 "delete_compare.png",
                 "Удаление 50 записей (лог. шкала)")
    plot_bst_degradation(means, "bst_degradation.png")


if __name__ == "__main__":
    main()