//
//  AppDelegate.swift
//  ComputeShader
//
//  Created by Alex Hoppen on 26.07.19.
//  Copyright © 2019 Alex Hoppen. All rights reserved.
//

import Cocoa

@NSApplicationMain
class AppDelegate: NSObject, NSApplicationDelegate {
    func applicationDidFinishLaunching(_ aNotification: Notification) {
        // MARK: - Image dimensions
        
        let width = 800
        let height = 600
        let wholeImageRegion = MTLRegionMake2D(0, 0, width, height)
        
        // MARK: - Copy raw image to memory
        let image = NSImage(named: "Stanford.jpg")!
        let inputBuffer = UnsafeMutableBufferPointer<UInt32>.allocate(capacity: width * height)
        defer {
            inputBuffer.deallocate()
        }
        image.writeRawIntoBuffer(inputBuffer, width: width, height: height)
        
        // MARK: - Get the GPU to execute commands on
        let device = MTLCreateSystemDefaultDevice()!
        
        // MARK: - Build Metal buffers
        let inputMtlBuffer = device.makeBuffer(bytes: inputBuffer.baseAddress!, length: inputBuffer.count * MemoryLayout<UInt32>.size, options: [])!
        let outputMtlBuffer = device.makeBuffer(length: inputBuffer.count * MemoryLayout<UInt32>.size, options: [])!
        

        // MARK: - Build the GPU command
        let library = device.makeDefaultLibrary()!
        let kernelFunction = library.makeFunction(name: "redTintBuffer")!
        
        let commandQueue = device.makeCommandQueue()!
        let commandBuffer = commandQueue.makeCommandBuffer()!
        
        // MARK: - The compute transformation
        let commandEncoder = commandBuffer.makeComputeCommandEncoder()!
        let pipelineState = try! device.makeComputePipelineState(function: kernelFunction)
        commandEncoder.setComputePipelineState(pipelineState)
        commandEncoder.setBuffer(inputMtlBuffer, offset: 0, index: 0)
        commandEncoder.setBuffer(outputMtlBuffer, offset: 0, index: 1)
        
        let threadsPerGrid = MTLSize(width: width * height, height: 1, depth: 1)
        let threadsPerThreadgroup = MTLSize(width: 64, height: 1, depth: 1)
        commandEncoder.dispatchThreads(threadsPerGrid, threadsPerThreadgroup: threadsPerThreadgroup)
        commandEncoder.endEncoding()
        
        // MARK: - Copy the result data into application memory when GPU has finished
        commandBuffer.addCompletedHandler { (commandBuffer) in
            let outputBuffer = UnsafeBufferPointer<UInt32>(start: outputMtlBuffer.contents().assumingMemoryBound(to: UInt32.self), count: outputMtlBuffer.length / MemoryLayout<UInt32>.size)
            let image = NSImage(rawBuffer: outputBuffer, width: width, height: height)
            displayImage(image: image)
        }
        // MARK: - Start processing
        commandBuffer.commit()
    }
}

