# frozen_string_literal: true

class Array
  # Returns a new array created by sorting +self+, using heap sort algorithm.
  #
  # Comparisons for the sort will be done using the <=> operator or using an optional code block.
  #
  # The block must implement a comparison between +a+ and +b+ and return an integer less than 0 when +b+ follows +a+,
  # +0+ when +a+ and +b+ are equivalent, or an integer greater than 0 when +a+ follows +b+.
  #
  # The result is not guaranteed to be stable. When the comparison of two elements returns +0+, the order of the
  # elements is unpredictable.
  #
  # @return [Array] the sorted array
  def heap_sort(&block)
    dup.heap_sort!(&block)
  end

  # Sorts +self+ in place, using heap sort algorithm.
  #
  # Comparisons for the sort will be done using the <=> operator or using an optional code block.
  #
  # The block must implement a comparison between +a+ and +b+ and return an integer less than 0 when +b+ follows +a+,
  # +0+ when +a+ and +b+ are equivalent, or an integer greater than 0 when +a+ follows +b+.
  #
  # The result is not guaranteed to be stable. When the comparison of two elements returns +0+, the order of the
  # elements is unpredictable.
  #
  # @return [Array] +self+
  def heap_sort!(&block)
    return self if length <= 1
    heapify(&block)
    (length - 1).downto(1) do |end_|
      swap(end_, 0)
      sift_down(0, end_ - 1, &block)
    end
    self
  end

  # Returns a new array created by sorting +self+ with heap sort algorithm, using a set of keys generated by mapping the
  # values in self through the given block.
  #
  # The result is not guaranteed to be stable. When the comparison of two elements returns +0+, the order of the
  # elements is unpredictable.
  #
  # If no block is given, an Enumerator is returned instead.
  #
  # @return [Array] if a block is given, the sorted array
  # @return [Enumerator] if no block is given, an Enumerator
  def heap_sort_by(&block)
    if block_given?
      dup.heap_sort_by!(&block)
    else
      to_enum :heap_sort_by
    end
  end

  # Sorts +self+ in place with heap sort algorithm, using a set of keys generated by mapping the values in self through
  # the given block.
  #
  # The result is not guaranteed to be stable. When the comparison of two elements returns +0+, the order of the
  # elements is unpredictable.
  #
  # If no block is given, an Enumerator is returned instead.
  #
  # @return [Array] if a block is given, the sorted array
  # @return [Enumerator] if no block is given, an Enumerator
  def heap_sort_by!(&_block)
    if block_given?
      heap_sort! do |a, b|
        yield(a) <=> yield(b)
      end
    else
      to_enum :heap_sort_by!
    end
  end

  private

  # @private
  def heapify(&block)
    end_ = length - 1
    (length / 2 - 1).downto(0) do |start|
      sift_down(start, end_, &block)
    end
  end

  # @private
  def sift_down(start, end_, &block)
    root = start
    loop do
      child = root * 2 + 1
      break if child > end_
      child += 1 if check_less_than_sibling(child, end_, &block)
      break if check_max_heap_order(root, child, &block)
      root = child
    end
  end

  # @private
  def check_less_than_sibling(child, end_, &block)
    child + 1 <= end_ && sort_compare(self[child], self[child + 1], &block) == -1
  end

  # @private
  def check_max_heap_order(root, child, &block)
    return true unless sort_compare(self[root], self[child], &block) == -1
    swap root, child
    false
  end
end