import math
import random
import time
import matplotlib.pyplot as plt


def merge_sort(arr, k):
    """
       Perform merge sort on a list, splitting it into k parts.

       Args:
           arr (list): The list to be sorted.
           k (int): The number of parts to split the list into.

       Returns:
           list: The sorted list.

       Examples:
           >>> merge_sort([12, 11, 13, 5, 6, 7], 3)
           [5, 6, 7, 11, 12, 13]

       Tests:
            >>> array = [0, 7, -2, 11, 5, 120, 0, 11, 6, 7, 11, -2, -2]
            >>> merge_sort(array, 2)
            [-2, -2, -2, 0, 0, 5, 6, 7, 7, 11, 11, 11, 120]

            >>> array = [0, 7, -2, 11, 5, 120, 0, 11, 6, 7, 11, -2, -2]
            >>> merge_sort(array, 3)
            [-2, -2, -2, 0, 0, 5, 6, 7, 7, 11, 11, 11, 120]

            >>> array = [0, 7, -2, 11, 5, 120, 0, 11, 6, 7, 11, -2, -2]
            >>> merge_sort(array, 13)
            [-2, -2, -2, 0, 0, 5, 6, 7, 7, 11, 11, 11, 120]

            >>> array = [20,1]
            >>> merge_sort(array, 2)
            [1, 20]

            >>> array = [20,1]
            >>> merge_sort(array, 20)
            [1, 20]

            >>> array = []
            >>> merge_sort(array, 2)
            []
            >>> merge_sort(array, 5)
            []
    """
    if len(arr) <= 1:
        return arr

    temp = [0 for _ in range(len(arr))]
    recursive_merge(arr, 0, len(arr), temp, k)
    return arr


def recursive_merge(arr, start, end, temp, k):
    """
        Merge multiple sorted arrays into a single sorted array.

        Args:
            arr (list): List to be sorted.
            start (int): Startindex of the subarray to be sorted.
            end (int): Endindex of the subarray to be sorted.
            temp (list): list of length at least end-start.
            k (int): The number of parts to split the list into.


        Tests:
            >>> A = [8, 10, 4]
            >>> tmp = [0, 0, 0]
            >>> recursive_merge(A, 1, 3, tmp, 2)
            >>> A
            [8, 4, 10]

            >>> A = []
            >>> tmp = []
            >>> recursive_merge(A, 0, 0, tmp, 2)
            >>> A
            []

            >>> A = [10, 5, 4]
            >>> tmp = [0,0,0]
            >>> recursive_merge(A, 0, 3, tmp, 2)
            >>> A
            [4, 5, 10]
    """
    length = end-start

    # if array has length 1, we are done (cause it is sorted)
    if length <= 1:
        return

    # split array in parts of size n/k (rounded up)
    part_size = math.ceil(length / k)
    start_idxs = [i for i in range(start, end, part_size)]
    end_idxs = [min(i + part_size, end) for i in start_idxs]

    # sort all parts, using recursion
    for i in range(len(start_idxs)):
        recursive_merge(arr, start_idxs[i], end_idxs[i], temp, k)

    # merge everything
    pointers = [0] * len(start_idxs)
    idx = start
    while True:
        smallest = None
        smallest_index = None

        # search for the smallest element among the k-pointers
        for i in range(len(start_idxs)):
            position = start_idxs[i] + pointers[i]
            if position < end_idxs[i]:
                if smallest is None or arr[position] < smallest:
                    smallest = arr[position]
                    smallest_index = i

        # if all pointers are at the end of their array -> break out
        if smallest is None:
            break

        # add smallest element to result end increase given pointer by 1
        pointers[smallest_index] += 1
        temp[idx] = smallest
        idx += 1
    for i in range(start, end):
        arr[i] = temp[i]
    return


def time_trial(range_n, calc_k):
    """
           Performs timed tests and returns time measurements.

           Args:
               range_n (range): specifies the input sizes.
               calc_k (lambda function): computes k given n.

           Returns:
               list: list of execution times.

    """
    results = []
    for n in range_n:
        arr = [random.randint(0, 10000) for _ in range(n)]
        start_time = time.time()
        merge_sort(arr, calc_k(n))
        end_time = time.time()

        results.append(end_time - start_time)
        if not all([arr[i] <= arr[i+1] for i in range(n-1)]):
            raise Exception("Didn't Sort Correctly!")
    return results


if __name__ == "__main__":
    range_n = range(100, 5101, 500)
    results_2 = time_trial(range_n, lambda n: 2)
    results_3 = time_trial(range_n, lambda n: 3)
    results_log = time_trial(range_n, lambda n: math.ceil(math.log(n)))
    results_sqrt = time_trial(range_n, lambda n: math.ceil(math.sqrt(n)))
    results_nfourth = time_trial(range_n, lambda n: n // 4)

    fig, (ax1, ax2) = plt.subplots(1, 2)
    fig.suptitle('MergeSort splitting into k parts')
    fig.set_size_inches(15, 5)

    # linear scale
    ax1.plot(range_n, results_2, "b", label="k=2")
    ax1.plot(range_n, results_3, "r", label="k=3")
    ax1.plot(range_n, results_log, 'm', label="k=log(n)")
    ax1.plot(range_n, results_sqrt, 'g', label="k=sqrt(n)")
    ax1.plot(range_n, results_nfourth, 'y', label="k=n/4")
    ax1.legend(loc="upper left")
    ax1.set_title('Linear Scale')
    ax1.grid(True)

    # log scale
    ax2.plot(range_n, results_2, "b", label="k=2")
    ax2.plot(range_n, results_3, "r", label="k=3")
    ax2.plot(range_n, results_log, 'm', label="k=log(n)")
    ax2.plot(range_n, results_sqrt, 'g', label="k=sqrt(n)")
    ax2.plot(range_n, results_nfourth, 'y', label="k=n/4")
    ax2.legend(loc="upper left")
    ax2.set_title('Log Scale')
    ax2.set_yscale("log")
    ax2.grid(True)

    plt.savefig('results.png')
    plt.show()