#!/usr/bin/env python3

# https://eax.me/heapsort/

import sys
import argparse
from random import random

def heap_insert(arr, size, x):
    """
    Insert a new element into a max-heap.

    This function assumes that the first 'size' elements of 'arr' already form a valid max-heap,
    and it inserts the new element 'x' into this heap, maintaining the heap property.

    Args:
        arr: The array containing the heap
        size: Current number of elements in the heap
        x: The element to insert

    Raises:
        IndexError: If attempting to insert into a heap that would exceed array bounds
    """
    if size >= len(arr):
        raise IndexError("Cannot insert into a full heap")

    arr[size] = x
    _sift_up(arr, size, size + 1)

def _sift_up(arr, idx, size):
    """
    Move an element up in the heap until the heap property is restored.

    Args:
        arr: The array being transformed into a heap
        idx: Index of the element to sift up
        size: Current size of the heap
    """
    if idx == 0:
        return

    parent = (idx - 1) // 2
    if arr[idx] > arr[parent]:
        arr[idx], arr[parent] = arr[parent], arr[idx]
        _sift_up(arr, parent, size)

def heap_extract(arr, size):
    """
    Extract the maximum element from a max-heap.

    This function removes and returns the maximum element (the root) from the heap,
    and restores the heap property.

    Args:
        arr: The array containing the heap
        size: Current number of elements in the heap

    Returns:
        The maximum element or None
    """
    if size == 0:
        return None

    max_element = arr[0]

    if size > 1:
        arr[0] = arr[size - 1]
        _sift_down(arr, 0, size - 1)

    return max_element

def _sift_down(arr, idx, size):
    """
    Move an element down in the heap until the heap property is restored.

    Args:
        arr: The array being transformed into a heap
        idx: Index of the element to sift down
        size: Current size of the heap portion of the array
    """
    largest = idx
    left = 2 * idx + 1
    right = 2 * idx + 2

    if left < size and arr[left] > arr[largest]:
        largest = left

    if right < size and arr[right] > arr[largest]:
        largest = right

    if largest != idx:
        arr[idx], arr[largest] = arr[largest], arr[idx]
        _sift_down(arr, largest, size)

def _heapify(arr):
    """
    Transform an array into a max-heap in-place, without using additional memory.

    This function reorganizes the elements in the array to satisfy the max-heap
    property: for each node i, the value at arr[i] is greater than or equal to
    the values at its children arr[2*i+1] and arr[2*i+2] (if they exist).

    Args:
        arr: The array to transform into a max-heap
    """
    n = len(arr)
    for i in range(n // 2 - 1, -1, -1):
        _sift_down(arr, i, n)

def heapsort(arr):
    """
    Sort an array in-place using the heap sort algorithm.

    This is an implementation of the heap sort algorithm that works in-place.
    It first transforms the array into a max-heap and then repeatedly extracts
    the maximum element.

    Args:
        arr: The array to be sorted in ascending order
    """

    _heapify(arr)
    for i in range(len(arr) - 1, 0, -1):
        arr[0], arr[i] = arr[i], arr[0]
        _sift_down(arr, 0, i)


def random_array(size):
    return [ int(random()*10) for i in range(0, size) ]


parser = argparse.ArgumentParser(description="Run heapsort tests")
parser.add_argument('--tests', type=int, default=10,
                    help='number of tests to run (default: 10)')
args = parser.parse_args()

for ntest in range(1, args.tests + 1):
    print(f"Test #{ntest}")

    external_array = random_array(int(1 + random()*ntest))
    for size in range(0, len(external_array)):
        heap_insert(external_array, size, int(random()*10))

    try:
        heap_insert(external_array, len(external_array), 5)
        print("ERROR: Should have raised IndexError")
        sys.exit(1)
    except IndexError:
        pass # OK

    last_val = None
    for size in range(len(external_array), 0, -1):
        val = heap_extract(external_array, size)
        if last_val is not None and val > last_val:
            print("ERROR: Heap property violation")
            sys.exit(1)
        last_val = val

    test_array = random_array(int(1 + random()*ntest))
    _heapify(test_array)
    n = len(test_array)
    is_heap = True
    for i in range(n):
        left = 2 * i + 1
        right = 2 * i + 2

        if left < n and test_array[i] < test_array[left]:
            print(f"Heap property violation at index {i}: {test_array[i]} < {test_array[left]}")
            is_heap = False

        if right < n and test_array[i] < test_array[right]:
            print(f"Heap property violation at index {i}: {test_array[i]} < {test_array[right]}")
            is_heap = False

    if not is_heap:
        print("ERROR: Array does not satisfy heap property")
        sys.exit(1)

    arr = random_array(int(1 + random()*ntest));
    arr1 = arr.copy()
    arr2 = arr.copy()
    heapsort(arr1)
    arr2.sort()
    if arr1 != arr2:
        print("ERROR: Array not sorted correctly")
        sys.exit(1)

print("All tests passed successfully!")