//
//  VisionHandler.swift
//  CapFinder
//
//  Created by User on 12.02.18.
//  Copyright © 2018 User. All rights reserved.
//

import Foundation
import Vision
import CoreML
import UIKit

/// Recognise categories in images
class Classifier: Logger {
    
    static let userDefaultsKey = "classifier.version"
    
    let model: VNCoreMLModel
    
    init(model: VNCoreMLModel) {
        self.model = model
    }
    /**
     Classify an image
     - Parameter image: The image to classify
     - Note: This method should not be scheduled on the main thread.
     */
    func recognize(image: UIImage, completion: @escaping (_ matches: [Int: Float]?) -> Void) {
        guard let ciImage = CIImage(image: image) else {
            error("Unable to create CIImage")
            completion(nil)
            return
        }

        let orientation = CGImagePropertyOrientation(image.imageOrientation)
        let handler = VNImageRequestHandler(ciImage: ciImage, orientation: orientation)
        let request = VNCoreMLRequest(model: model) { request, error in
            let matches = self.process(request: request, error: error)
            completion(matches)
        }
        request.imageCropAndScaleOption = .centerCrop
        do {
            try handler.perform([request])
        } catch {
            self.error("Failed to perform classification: \(error)")
        }
    }
    
    private func process(request: VNRequest, error: Error?) -> [Int : Float]? {
        if let e = error {
            self.error("Unable to classify image: \(e.localizedDescription)")
            return nil
        }
        guard let result = request.results as? [VNClassificationObservation] else {
            self.error("Invalid classifier result: \(String(describing: request.results))")
            return nil
        }
        let matches = result.reduce(into: [:]) { $0[Int($1.identifier)!] = $1.confidence }
        
        log("Classifed image with \(matches.count) classes")
        return matches
    }
}