//
//  BaseViewController.swift
//  MachineLearningDemo
//
//  Created by Pavel Odstrčilík on 20.04.2021.
//

import Foundation
import UIKit
import Vision
import SnapKit

open class BaseViewController: UIViewController, UINavigationControllerDelegate, UIImagePickerControllerDelegate {
    let imagePicker = UIImagePickerController()
    let imageView = UIImageView()
    let resultVisualEffectView = UIVisualEffectView()
    let resultLabel = UILabel()
    let activityIndicatorView = UIActivityIndicatorView(style: .large)
    let informationLabel = UILabel()
    var mlModel: MLModel?

    lazy var classificationRequest: VNCoreMLRequest? = {
        do {
            guard let mlModelUnwrapped = mlModel else {
                return nil
            }
            let model = try VNCoreMLModel(for: mlModelUnwrapped)
            let mlRequest = VNCoreMLRequest(model: model,
                                            completionHandler: { [weak self] request,
                                                                             error in
                                                self?.processClassification(request: request, error: error)
                                            })
            mlRequest.imageCropAndScaleOption = .scaleFill
            return mlRequest
        } catch {
            fatalError("Failed to load model: \(error)")
        }
    }()


    public init(model: MLModel) {
        super.init(nibName: nil, bundle: nil)
        mlModel = model
        prepareBarButtonItems()
        prepareImagePicker()
        prepareImageView()
        prepareResultVisualEffectView()
        prepareResultLabel()
        prepareActivityIndicatorView()
        prepareInformationLabel()
        self.title = "Machine Learning Demo"
        view.backgroundColor = .systemBackground
    }
    
    required public init?(coder: NSCoder) {
        fatalError("init(coder:) has not been implemented")
    }

    open func prepareBarButtonItems() {
        let photoImage = UIImage(systemName: "camera.fill")
        let photoBarButtonItem = UIBarButtonItem(image: photoImage, style: .plain,
                                                 target: self, action: #selector(showCamera))
        let galleryImage = UIImage(systemName: "photo")
        let galleryBarButtonItem = UIBarButtonItem(image: galleryImage, style: .plain,
                                                   target: self, action: #selector(showImagePicker))
        self.navigationItem.rightBarButtonItems = [photoBarButtonItem, galleryBarButtonItem]
    }

    open func prepareImagePicker() {
        imagePicker.delegate = self
        imagePicker.allowsEditing = false
    }

    @objc open func showImagePicker() {
        guard UIImagePickerController.isSourceTypeAvailable(.savedPhotosAlbum) else {
            return
        }
        imagePicker.sourceType = .savedPhotosAlbum
        self.present(imagePicker, animated: true, completion: nil)
    }

    @objc open func showCamera() {
        guard UIImagePickerController.isSourceTypeAvailable(.camera) else {
            return
        }
        imagePicker.sourceType = .camera
        self.present(imagePicker, animated: true, completion: nil)
    }

    open func prepareImageView() {
        view.addSubview(imageView)
        imageView.snp.makeConstraints { (make) in
            make.top.equalTo(view.safeAreaLayoutGuide.snp.top)
            make.bottom.equalTo(view.safeAreaLayoutGuide.snp.bottom)
            make.leading.trailing.equalToSuperview()
        }
        imageView.contentMode = .scaleAspectFit
    }

    // MARK: - UIImagePickerControllerDelegate
    public func imagePickerController(_ picker: UIImagePickerController,
                                      didFinishPickingMediaWithInfo
                                        info: [UIImagePickerController.InfoKey: Any]) {
        picker.dismiss(animated: true, completion: nil)
        if let image = info[UIImagePickerController.InfoKey.originalImage] as? UIImage {
            imageView.image = image
            updateClassification(image: image)
        }
    }

    open func processClassification(request: VNRequest, error: Error?) {
        DispatchQueue.main.async {
            self.activityIndicatorView.stopAnimating()

            guard let result = request.results as? [VNClassificationObservation] else {
                return
            }

            if result.isEmpty {
                self.resultLabel.text = "Nothing recognized"
            } else {
                guard let resultUnwrapped = result.first else {
                    return
                }
                let classificationResult = String(format: " (%.2f) %@", resultUnwrapped.confidence, resultUnwrapped.identifier)
                self.resultLabel.text = "Classification: \(classificationResult)"
            }
        }
    }

    open func updateClassification(image: UIImage) {
        informationLabel.isHidden = true
        resultLabel.text = ""
        activityIndicatorView.startAnimating()
        let orientation = UInt32(image.imageOrientation.rawValue)
        let imageOrientation = CGImagePropertyOrientation(rawValue:
                                                            orientation) ?? .up

        guard let image = CIImage(image: image) else {
            return
        }

        DispatchQueue.global(qos: .userInitiated).async {
            let handler = VNImageRequestHandler(ciImage: image,
                                                orientation: imageOrientation)
            do {
                if let classificationRequestUnwrapped = self.classificationRequest {
                    try handler.perform([classificationRequestUnwrapped])
                }
            } catch {
                print("Error \(error.localizedDescription)")
            }
        }
    }

    open func prepareResultVisualEffectView() {
        resultVisualEffectView.effect = UIBlurEffect(style: .regular)
        view.addSubview(resultVisualEffectView)
        resultVisualEffectView.snp.makeConstraints { (make) in
            make.bottom.equalTo(view.safeAreaLayoutGuide.snp.bottom).offset(-20)
            make.leading.equalToSuperview().offset(50)
            make.trailing.equalToSuperview().offset(-50)
            make.height.equalTo(100)
        }

        resultVisualEffectView.layer.cornerRadius = 10
        resultVisualEffectView.layer.masksToBounds = true
    }

    open func prepareResultLabel() {
        resultVisualEffectView.contentView.addSubview(resultLabel)
        resultLabel.snp.makeConstraints { (make) in
            make.edges.equalToSuperview().inset(UIEdgeInsets(top: 0, left: 10, bottom: 0, right: 10))
        }
        resultLabel.textAlignment = .center
        resultLabel.font = .systemFont(ofSize: 20)
    }

    open func prepareActivityIndicatorView() {
        resultVisualEffectView.contentView.addSubview(activityIndicatorView)
        activityIndicatorView.snp.makeConstraints { (make) in
            make.edges.equalToSuperview()
        }
    }

    open func prepareInformationLabel() {
        informationLabel.text = "Pick a photo from gallery or take a photo"
        view.addSubview(informationLabel)
        informationLabel.snp.makeConstraints { (make) in
            make.leading.equalToSuperview().offset(20)
            make.trailing.equalToSuperview().offset(-20)
            make.center.equalToSuperview()
        }
        informationLabel.textAlignment = .center
        informationLabel.textColor = .systemGray
    }
}
