//
//  BinarySearch.swift
//  SwiftSoup-iOS
//
//  Created by Garth Snyder on 2/28/19.
//
//  Adapted from https://stackoverflow.com/questions/31904396/swift-binary-search-for-standard-array
//

import Foundation

extension Collection {

  /// Generalized binary search algorithm for ordered Collections
  ///
  /// Behavior is undefined if the collection is not properly sorted.
  ///
  /// This is only O(logN) for RandomAccessCollections; Collections in
  /// general may implement offsetting of indexes as an O(K) operation. (E.g.,
  /// Strings are like this).
  ///
  /// - Note: If you are using this for searching only (not insertion), you
  ///     must always test the element at the returned index to ensure that
  ///     it's a genuine match. If the element is not present in the array,
  ///     you will still get a valid index back that represents the location
  ///     where it should be inserted. Also check to be sure the returned
  ///     index isn't off the end of the collection.
  ///
  /// - Parameter predicate: Reports the ordering of a given Element relative
  ///     to the desired Element. Typically, this is <.
  ///
  /// - Returns: Index N such that the predicate is true for all elements up to
  ///     but not including N, and is false for all elements N and beyond

  func binarySearch(predicate: (Element) -> Bool) -> Index {
    var low = startIndex
    var high = endIndex
    while low != high {
      let mid = index(low, offsetBy: distance(from: low, to: high) / 2)
      if predicate(self[mid]) {
        low = index(after: mid)
      } else {
        high = mid
      }
    }
    return low
  }

  /// Binary search lookup for ordered Collections using a KeyPath
  /// relative to Element.
  ///
  /// Behavior is undefined if the collection is not properly sorted.
  ///
  /// This is only O(logN) for RandomAccessCollections; Collections in
  /// general may implement offsetting of indexes as an O(K) operation. (E.g.,
  /// Strings are like this).
  ///
  /// - Note: If you are using this for searching only (not insertion), you
  ///     must always test the element at the returned index to ensure that
  ///     it's a genuine match. If the element is not present in the array,
  ///     you will still get a valid index back that represents the location
  ///     where it should be inserted. Also check to be sure the returned
  ///     index isn't off the end of the collection.
  ///
  /// - Parameter keyPath: KeyPath that extracts the Element value on which
  ///     the Collection is presorted. Must be Comparable and Equatable.
  ///     ordering is presumed to be <, however that is defined for the type.
  ///
  /// - Returns: The index of a matching element, or nil if not found. If
  ///     the return value is non-nil, it is always a valid index.

  func indexOfElement<T>(withValue value: T, atKeyPath keyPath: KeyPath<Element, T>) -> Index?
  where T: Comparable & Equatable {
    let ix = binarySearch { $0[keyPath: keyPath] < value }
    guard ix < endIndex else { return nil }
    guard self[ix][keyPath: keyPath] == value else { return nil }
    return ix
  }

  func element<T>(withValue value: T, atKeyPath keyPath: KeyPath<Element, T>) -> Element?
  where T: Comparable & Equatable {
    if let ix = indexOfElement(withValue: value, atKeyPath: keyPath) {
      return self[ix]
    }
    return nil
  }

  func elements<T>(withValue value: T, atKeyPath keyPath: KeyPath<Element, T>) -> [Element]
  where T: Comparable & Equatable {
    guard let start = indexOfElement(withValue: value, atKeyPath: keyPath) else { return [] }
    var end = index(after: start)
    while end < endIndex && self[end][keyPath: keyPath] == value {
      end = index(after: end)
    }
    return Array(self[start..<end])
  }
}
