import subprocess
from time import time, sleep
import matplotlib.pyplot as plt

pivot_strategies = {
    "rightmost": 0,
    "middle": 1,
    "random": 2
}

def run_benchmark(array, pivot_strategy):
    input_data = f"{len(array)}\n{pivot_strategies[pivot_strategy]}\n" + "\n".join(map(str, array)) + "\n"

    start = time()
    process = subprocess.Popen("./dist/quicksort", stdin=subprocess.PIPE)
    process.communicate(input_data.encode())
    if process.returncode != 0:
        raise RuntimeError("Error while benchmarking")
    end = time()

    return end - start

def generate_v_shaped(array_size):
    return [i for i in range(array_size // 2, 0, -1)] + [i for i in range(array_size // 2)]

def plot(sequences, array_size):
    pivot_names = {
        "rightmost": "Pivot prawy",
        "middle": "Pivot środkowy",
        "random": "Pivot losowy"
    }
    plt.figure(figsize=(10, 6))
    for pivot_type, timings in sequences.items():
        plt.plot(array_size, timings, label=pivot_names[pivot_type])
    plt.xlabel("Rozmiar tablicy")
    plt.ylabel("Czas (ms)")
    plt.title("Quicksort")
    plt.legend()
    plt.savefig("./benchmarks/quicksort.png")

if __name__ == "__main__":
    array_sizes = [i for i in range(1000, 15001, 1000)]

    with open("results.txt", "w") as f:
        sequenceTiming = {
                "rightmost": [],
                "middle": [],
                "random": [],
            }
        for pivot_strategy in pivot_strategies:
            print(f"Running with pivot strategy {pivot_strategy}")
            for array_size in array_sizes:
                sequence = generate_v_shaped(array_size)
                durations = []
                for i in range(10):
                    print(f"     - with size of {array_size} {i+1}/10", end="\r")
                    duration = run_benchmark(sequence, pivot_strategy)
                    duration *= 1000
                    durations.append(duration)
                    sleep(0.1)
                print()
                duration = sum(durations) / len(durations)
                sequenceTiming[pivot_strategy].append(duration)
                sleep(0.1)
        plot(sequenceTiming, array_sizes)
    print("Benchmark finished")