import random
import pandas as pd
import time
import sys
import os
import matplotlib.pyplot as plt

# Увеличиваем лимит рекурсии для BST на отсортированных данных (может достичь глубины N)
sys.setrecursionlimit(20000)

# =========================================================
# 1. СВЯЗНЫЙ СПИСОК (LinkedListPhoneBook)
# =========================================================
def ll_insert(head, name, phone):
    if head is None:
        return {'name': name, 'phone': phone, 'next': None}
    
    curr = head
    while True:
        if curr['name'] == name:
            curr['phone'] = phone  # Обновление существующей записи
            break
        if curr['next'] is None:
            curr['next'] = {'name': name, 'phone': phone, 'next': None}
            break
        curr = curr['next']
    return head

def ll_find(head, name):
    curr = head
    while curr:
        if curr['name'] == name:
            return curr['phone']
        curr = curr['next']
    return None

def ll_delete(head, name):
    if head is None:
        return None
    if head['name'] == name:
        return head['next']
    
    curr = head
    while curr['next']:
        if curr['next']['name'] == name:
            curr['next'] = curr['next']['next']
            break
        curr = curr['next']
    return head

def ll_list_all(head):
    res = []
    curr = head
    while curr:
        res.append((curr['name'], curr['phone']))
        curr = curr['next']
    res.sort(key=lambda x: x[0])
    return res

# =========================================================
# 2. ХЕШ-ТАБЛИЦА
# =========================================================
HT_SIZE = 10007  # Простое число для равномерного распределения

def ht_init():
    return [None] * HT_SIZE

def _ht_idx(name):
    return hash(name) % HT_SIZE

def ht_insert(buckets, name, phone):
    idx = _ht_idx(name)
    buckets[idx] = ll_insert(buckets[idx], name, phone)
    return buckets

def ht_find(buckets, name):
    return ll_find(buckets[_ht_idx(name)], name)

def ht_delete(buckets, name):
    idx = _ht_idx(name)
    buckets[idx] = ll_delete(buckets[idx], name)
    return buckets

def ht_list_all(buckets):
    res = []
    for bucket in buckets:
        curr = bucket
        while curr:
            res.append((curr['name'], curr['phone']))
            curr = curr['next']
    res.sort(key=lambda x: x[0])
    return res

# =========================================================
# 3. ДВОИЧНОЕ ДЕРЕВО ПОИСКА (BST)
# =========================================================
def bst_insert(root, name, phone):
    if root is None:
        return {'name': name, 'phone': phone, 'left': None, 'right': None}
    
    if name < root['name']:
        root['left'] = bst_insert(root['left'], name, phone)
    elif name > root['name']:
        root['right'] = bst_insert(root['right'], name, phone)
    else:
        root['phone'] = phone
    return root

def bst_find(root, name):
    curr = root
    while curr:
        if name == curr['name']:
            return curr['phone']
        elif name < curr['name']:
            curr = curr['left']
        else:
            curr = curr['right']
    return None

def bst_delete(root, name):
    if root is None:
        return None
    
    if name < root['name']:
        root['left'] = bst_delete(root['left'], name)
    elif name > root['name']:
        root['right'] = bst_delete(root['right'], name)
    else:
        # Узел найден
        if root['left'] is None:
            return root['right']
        if root['right'] is None:
            return root['left']
        
        # Два потомка: находим минимальный в правом поддереве
        min_node = root['right']
        while min_node['left'] is not None:
            min_node = min_node['left']
        
        root['name'] = min_node['name']
        root['phone'] = min_node['phone']
        root['right'] = bst_delete(root['right'], min_node['name'])
    return root

def bst_list_all(root):
    if root is None:
        return []
    return bst_list_all(root['left']) + [(root['name'], root['phone'])] + bst_list_all(root['right'])

# =========================================================
# ЭКСПЕРИМЕНТАЛЬНАЯ ЧАСТЬ
# =========================================================
def run_experiments():
    N = 10000
    RECORDS = [(f"User_{i:05d}", f"+7900{i:04d}{i%100:02d}") for i in range(N)]
    
    records_shuffled = RECORDS[:]
    random.shuffle(records_shuffled)
    
    records_sorted = sorted(RECORDS, key=lambda x: x[0])
    
    # Наборы для поиска и удаления
    existing_names = [r[0] for r in random.sample(RECORDS, 100)]
    non_existing_names = [f"None_{i}" for i in range(10)]
    find_names = existing_names + non_existing_names
    delete_names = [r[0] for r in random.sample(RECORDS, 50)]
    
    structures = {
        "LinkedList": (lambda: None, ll_insert, ll_find, ll_delete),
        "HashTable":  (ht_init, ht_insert, ht_find, ht_delete),
        "BST":        (lambda: None, bst_insert, bst_find, bst_delete)
    }
    
    modes = {"случайный": records_shuffled, "отсортированный": records_sorted}
    results = []
    
    print("Запуск экспериментов...")
    trials = 5
    for struct_name, (init_f, ins_f, find_f, del_f) in structures.items():
        for mode_name, data in modes.items():
            print(f"  {struct_name} | {mode_name}")
            for t in range(1, trials + 1):
                # Инициализация
                ds = init_f()
                
                # A. Вставка
                t0 = time.perf_counter()
                for name, phone in data:
                    ds = ins_f(ds, name, phone)
                t_ins = time.perf_counter() - t0
                
                # B. Поиск
                t0 = time.perf_counter()
                for name in find_names:
                    find_f(ds, name)
                t_find = time.perf_counter() - t0
                
                # C. Удаление
                t0 = time.perf_counter()
                for name in delete_names:
                    ds = del_f(ds, name)
                t_del = time.perf_counter() - t0
                
                results.append([struct_name, mode_name, "вставка", t, t_ins])
                results.append([struct_name, mode_name, "поиск", t, t_find])
                results.append([struct_name, mode_name, "удаление", t, t_del])
                
    return results

def save_and_plot(results):
    import os
    import matplotlib.pyplot as plt
    import pandas as pd

    os.makedirs("docs/data", exist_ok=True)

    # 1. Сохранение CSV (как было)
    df = pd.DataFrame(results, columns=["Структура", "Режим", "Операция", "Повторение", "Время (сек)"])
    avg = df.groupby(["Структура", "Режим", "Операция"])["Время (сек)"].mean().reset_index()
    avg["Повторение"] = "СРЕДНЕЕ"
    df_full = pd.concat([df, avg], ignore_index=True)
    df_full.to_csv("docs/data/results.csv", index=False, encoding="utf-8-sig")

    # 2. Улучшенный график: 3 отдельных подграфика + логарифмическая шкала
    fig, axes = plt.subplots(1, 3, figsize=(18, 6))
    operations = ["вставка", "поиск", "удаление"]
    structures_order = ["HashTable", "BST", "LinkedList"]  # Фиксируем порядок для удобства чтения
    colors = {"случайный": "#6C157F", "отсортированный": "#1E299F"}

    for ax, op in zip(axes, operations):
        op_data = avg[avg["Операция"] == op]
        pivot = op_data.pivot(index="Структура", columns="Режим", values="Время (сек)")
        pivot = pivot.reindex(structures_order)  # Ставим структуры в удобном порядке

        pivot.plot(kind="bar", ax=ax, color=[colors["случайный"], colors["отсортированный"]], width=0.75)
        ax.set_title(f"Операция: {op.capitalize()}")
        ax.set_ylabel("Время (сек)")
        ax.set_xticklabels(ax.get_xticklabels(), rotation=0)
        ax.grid(axis="y", alpha=0.3, linestyle="--")
        
        # ЛОГАРИФМИЧЕСКАЯ ШКАЛА: обязательна при разбросе от 0.0001 до 30 сек
        ax.set_yscale("log")
        ax.legend(title="Режим", loc="upper right")

    fig.suptitle("Сравнение производительности структур данных", fontsize=16, y=1.05)
    plt.tight_layout()
    plt.savefig("docs/data/plot.png", dpi=200, bbox_inches="tight")

if __name__ == "__main__":
    res = run_experiments()
    save_and_plot(res)
    print("Эксперимент завершен")