space-ios/iOSClient/Library/WeScan/Scan/RectangleFeaturesFunnel.swift

//
//  RectangleFeaturesFunnel.swift
//  WeScan
//
//  Created by Boris Emorine on 2/9/18.
//  Copyright © 2018 WeTransfer. All rights reserved.
//

import Foundation
import AVFoundation

/// `RectangleFeaturesFunnel` is used to improve the confidence of the detected rectangles.
/// Feed rectangles to a `RectangleFeaturesFunnel` instance, and it will call the completion block with a rectangle whose confidence is high enough to be displayed.
final class RectangleFeaturesFunnel {
    
    /// `RectangleMatch` is a class used to assign matching scores to rectangles.
    private final class RectangleMatch: NSObject {
        /// The rectangle feature object associated to this `RectangleMatch` instance.
        let rectangleFeature: CIRectangleFeature
        
        /// The score to indicate how strongly the rectangle of this instance matches other recently added rectangles.
        /// A higher score indicates that many recently added rectangles are very close to the rectangle of this instance.
        var matchingScore = 0

        init(rectangleFeature: CIRectangleFeature) {
            self.rectangleFeature = rectangleFeature
        }
        
        override var description: String {
            return "Matching score: \(matchingScore) - Rectangle: \(rectangleFeature)"
        }
        
        /// Whether the rectangle of this instance is within the distance of the given rectangle.
        ///
        /// - Parameters:
        ///   - rectangle: The rectangle to compare the rectangle of this instance with.
        ///   - threshold: The distance used to determinate if the rectangles match in pixels.
        /// - Returns: True if both rectangles are within the given distance of each other.
        func matches(_ rectangle: CIRectangleFeature, withThreshold threshold: CGFloat) -> Bool {
            return rectangleFeature.isWithin(threshold, ofRectangleFeature: rectangle)
        }
    }
    
    /// The queue of last added rectangles. The first rectangle is oldest one, and the last rectangle is the most recently added one.
    private var rectangles = [RectangleMatch]()
    
    /// The maximum number of rectangles to compare newly added rectangles with. Determines the maximum size of `rectangles`. Increasing this value will impact performance.
    let maxNumberOfRectangles = 8
    
    /// The minimum number of rectangles needed to start making comparaisons and determining which rectangle to display. This value should always be inferior than `maxNumberOfRectangles`.
    /// A higher value will delay the first time a rectangle is displayed.
    let minNumberOfRectangles = 3
    
    /// The value in pixels used to determine if two rectangle match or not. A higher value will prevent displayed rectangles to be refreshed. On the opposite, a smaller value will make new rectangles be displayed constantly.
    let matchingThreshold: CGFloat = 40.0
    
    /// The minumum number of matching rectangles (within the `rectangle` queue), to be confident enough to display a rectangle.
    let minNumberOfMatches = 2

    /// Add a rectangle to the funnel, and if a new rectangle should be displayed, the completion block will be called.
    /// The algorithm works the following way:
    /// 1. Makes sure that the funnel has been fed enough rectangles
    /// 2. Removes old rectangles if needed
    /// 3. Compares all of the recently added rectangles to find out which one match each other
    /// 4. Within all of the recently added rectangles, finds the "best" one (@see `bestRectangle(withCurrentlyDisplayedRectangle:)`)
    /// 5. If the best rectangle is different than the currently displayed rectangle, informs the listener that a new rectangle should be displayed
    /// - Parameters:
    ///   - rectangleFeature: The rectangle to feed to the funnel.
    ///   - currentRectangle: The currently displayed rectangle. This is used to avoid displaying very close rectangles.
    ///   - completion: The completion block called when a new rectangle should be displayed.
    func add(_ rectangleFeature: CIRectangleFeature, currentlyDisplayedRectangle currentRectangle: CIRectangleFeature?, completion: (CIRectangleFeature) -> Void) {
        let rectangleMatch = RectangleMatch(rectangleFeature: rectangleFeature)
        rectangles.append(rectangleMatch)
        
        guard rectangles.count >= minNumberOfRectangles else {
            return
        }
        
        if rectangles.count > maxNumberOfRectangles {
            rectangles.removeFirst()
        }
        
        updateRectangleMatches()
        
        guard let bestRectangle = bestRectangle(withCurrentlyDisplayedRectangle: currentRectangle) else {
            return
        }
        
        if let previousRectangle = currentRectangle,
            bestRectangle.rectangleFeature.isWithin(matchingThreshold, ofRectangleFeature: previousRectangle) {
        } else if bestRectangle.matchingScore >= minNumberOfMatches {
            completion(bestRectangle.rectangleFeature)
        }
    }
    
    /// Determines which rectangle is best to displayed.
    /// The criteria used to find the best rectangle is its matching score.
    /// If multiple rectangles have the same matching score, we use a tie breaker to find the best rectangle (@see breakTie(forRectangles:)).
    /// Parameters:
    ///   - currentRectangle: The currently displayed rectangle. This is used to avoid displaying very close rectangles.
    /// Returns: The best rectangle to display given the current history.
    private func bestRectangle(withCurrentlyDisplayedRectangle currentRectangle: CIRectangleFeature?) -> RectangleMatch? {
        var bestMatch: RectangleMatch?
        
        rectangles.reversed().forEach { (rectangle) in
            guard let best = bestMatch else {
                bestMatch = rectangle
                return
            }
            
            if rectangle.matchingScore > best.matchingScore {
                bestMatch = rectangle
                return
            } else if rectangle.matchingScore == best.matchingScore {
                guard let currentRectangle = currentRectangle else {
                    return
                }
                
                bestMatch = breakTie(between: best, rect2: rectangle, currentRectangle: currentRectangle)
            }
        }
        
        return bestMatch
    }
    
    /// Breaks a tie between two rectangles to find out which is best to display.
    /// The first passed rectangle is returned if no other criteria could be used to break the tie.
    /// If the first passed rectangle (rect1) is close to the currently displayed rectangle, we pick it.
    /// Otherwise if the second passed rectangle (rect2) is close to the currently displayed rectangle, we pick this one.
    /// Finally, if none of the passed in rectangles are close to the currently displayed rectangle, we arbitrary pick the first one.
    /// - Parameters:
    ///   - rect1: The first rectangle to compare.
    ///   - rect2: The second rectangle to compare.
    ///   - currentRectangle: The currently displayed rectangle. This is used to avoid displaying very close rectangles.
    /// - Returns: The best rectangle to display between two rectangles with the same matching score.
    private func breakTie(between rect1: RectangleMatch, rect2: RectangleMatch, currentRectangle: CIRectangleFeature) -> RectangleMatch {
        if rect1.rectangleFeature.isWithin(matchingThreshold, ofRectangleFeature: currentRectangle) {
            return rect1
        } else if rect2.rectangleFeature.isWithin(matchingThreshold, ofRectangleFeature: currentRectangle) {
            return rect2
        }
        
        return rect1
    }
    
    /// Loops through all of the rectangles of the queue, and gives them a score depending on how many they match. @see `RectangleMatch.matchingScore`
    private func updateRectangleMatches() {
        resetMatchingScores()
        
        for (i, currentRect) in rectangles.enumerated() {
            for (j, rect) in rectangles.enumerated() {
                if j > i && currentRect.matches(rect.rectangleFeature, withThreshold: matchingThreshold) {
                    currentRect.matchingScore += 1
                    rect.matchingScore += 1
                }
            }
        }
    }
    
    /// Resets the matching score of all of the rectangles in the queue to 0
    private func resetMatchingScores() {
        rectangles = rectangles.map { (rectangle) -> RectangleMatch in
            rectangle.matchingScore = 0
            return rectangle
        }
    }
    
}