//
//  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 textures
    let inputTextureDescriptor = MTLTextureDescriptor.texture2DDescriptor(pixelFormat: .rgba8Unorm, width: width, height: height, mipmapped: false)
    inputTextureDescriptor.usage = .shaderRead
    inputTextureDescriptor.storageMode = .managed
    let inputTexture = device.makeTexture(descriptor: inputTextureDescriptor)!
    let bytesPerPixel = 4
    let bytesPerRow = width * bytesPerPixel
    inputTexture.replace(region: wholeImageRegion, mipmapLevel: 0, withBytes: inputBuffer.baseAddress!, bytesPerRow: bytesPerRow)
    
    let outputTextureDescriptor = MTLTextureDescriptor.texture2DDescriptor(pixelFormat: inputTexture.pixelFormat, width: width, height: height, mipmapped: false)
    outputTextureDescriptor.usage = [.shaderRead, .shaderWrite]
    outputTextureDescriptor.storageMode = .managed
    let outputTexture = device.makeTexture(descriptor: outputTextureDescriptor)!
    


    // MARK: - Build the GPU command
    let library = device.makeDefaultLibrary()!
    let kernelFunction = library.makeFunction(name: "redTintTexture")!
    
    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.setTexture(inputTexture, index: 0)
    commandEncoder.setTexture(outputTexture, index: 1)
        
    let threadsPerGrid = MTLSize(width: width, height: height, depth: 1)
    let threadsPerThreadgroup = MTLSize(width: 8, height: 8, depth: 1)
    commandEncoder.dispatchThreads(threadsPerGrid, threadsPerThreadgroup: threadsPerThreadgroup)
    commandEncoder.endEncoding()
    
    // MARK: - Synchronise output memory in GPU with CPU memory
    let blitCommandEnccoder = commandBuffer.makeBlitCommandEncoder()!
    blitCommandEnccoder.synchronize(resource: outputTexture)
    blitCommandEnccoder.endEncoding()
    
    
    // MARK: - Copy the result data into application memory when GPU has finished
    commandBuffer.addCompletedHandler { (commandBuffer) in
        let outputBuffer = UnsafeMutableBufferPointer<UInt32>.allocate(capacity: width * height)
        defer {
            outputBuffer.deallocate()
        }
        outputTexture.getBytes(outputBuffer.baseAddress!, bytesPerRow: bytesPerRow, from: wholeImageRegion, mipmapLevel: 0)
        let image = NSImage(rawBuffer: UnsafeBufferPointer(outputBuffer), width: width, height: height)
        displayImage(image: image)
    }
    // MARK: - Start processing
    commandBuffer.commit()
    }
}

