import Cocoa
import CreateML

let defaultIterations = 17
let defaultWorkingDirectory = URL(fileURLWithPath: "/Users/imac/Development/CapCollectorData")
let defaultTrainDirectory = defaultWorkingDirectory.appendingPathComponent("images")
let defaultClassifierFile = defaultWorkingDirectory.appendingPathComponent("classifier.mlmodel")

func readArguments() -> (images: URL, classifier: URL, iterations: Int) {
    let count = CommandLine.argc
    guard count > 1 else {
        // No arguments
        return (defaultTrainDirectory, defaultClassifierFile, defaultIterations)
    }
    // First argument is the image directory
    let imageDir = URL(fileURLWithPath: CommandLine.arguments[1])
    let classifier = imageDir.deletingLastPathComponent().appendingPathComponent("classifier.mlmodel")
    guard count > 2 else {
        // Single argument is the image directory
        return (imageDir, classifier, defaultIterations)
    }
    // Second argument is the iteration count
    guard let iterations = Int(CommandLine.arguments[2]) else {
        print("[ERROR] Invalid iterations argument '\(CommandLine.arguments[2])'")
        exit(-1)
    }
    guard count > 3 else {
        return (imageDir, classifier, iterations)
    }
    // Third argument is the classifier path
    let classifierPath = URL(fileURLWithPath: CommandLine.arguments[3])
    if count > 4 {
        print("[WARNING] Ignoring additional arguments")
    }
    return (imageDir, classifierPath, iterations)
}

let arguments = readArguments()

print("[INFO] Using images in \(arguments.images.path)")
print("[INFO] Training for \(arguments.iterations) iterations")
print("[INFO] Classifier path set to \(arguments.classifier.path)")

var params = MLImageClassifier.ModelParameters(augmentation: [])
params.maxIterations = arguments.iterations

let model = try MLImageClassifier(
    trainingData: .labeledDirectories(at: arguments.images),
    parameters: params)

print("[INFO] Writing classifier...")
try model.write(to: arguments.classifier)

/*
let evaluation = model.evaluation(on: .labeledDirectories(at: trainDirectory))
print("Printing evaluation:")
print(evaluation)
print("Finished evaluation")
 */