import ARKit
import SceneKit

let horizontalPoints = 256 / 2
let verticalPoints = 192 / 2
var depthNodes = [SCNNode]()
var parentDebugNodes = SCNNode()
var sceneView: ARSCNView!

// Somewhere during setup
func setup() {
    let configuration = ARWorldTrackingConfiguration()
    configuration.frameSemantics = .smoothedSceneDepth
    sceneView.session.run(configuration)
    sceneView.scene.rootNode.addChildNode(parentDebugNodes)

    let sizeGeomPredictions = 0.005
    let geom = SCNBox(width: sizeGeomPredictions, height: sizeGeomPredictions, length: sizeGeomPredictions, chamferRadius: 0)
    geom.firstMaterial?.diffuse.contents = UIColor.green
    
    for _ in 0..<(horizontalPoints * verticalPoints) {
        let node = SCNNode(geometry: geom)
        self.parentDebugNodes.addChildNode(node)
        self.depthNodes.append(node)
    }
}


func session(_ session: ARSession, didUpdate frame: ARFrame) {
    guard let smoothedDepth = frame.smoothedSceneDepth?.depthMap else {
        return
    }
    let capturedImage = frame.capturedImage

    let lockFlags = CVPixelBufferLockFlags.readOnly
    CVPixelBufferLockBaseAddress(smoothedDepth, lockFlags)
    defer {
      CVPixelBufferUnlockBaseAddress(smoothedDepth, lockFlags) 
    }

    let baseAddress = CVPixelBufferGetBaseAddressOfPlane(smoothedDepth, 0)!
    let depthByteBuffer = baseAddress.assumingMemoryBound(to: Float32.self)
  
    // The `.size` accessor simply read the CVPixelBuffer's width and height in pixels.
    //
    // They are the same ratio:
    // 1920 x 1440 = 1440 x 1920 = 0.75
    let depthMapSize = smoothedDepth.size
    // 192 x 256 = 0.75
    let capturedImageSize = capturedImage.size

    var cameraIntrinsics = frame.camera.intrinsics
    let depthResolution = simd_float2(x: Float(depthMapSize.x), y: Float(depthMapSize.y))
    let scaleRes = simd_float2(x: Float(capturedImageSize.x) / depthResolution.x,
                               y: Float(capturedImageSize.y) / depthResolution.y )
    // Make the camera intrinsics be with respect to Depth.
    cameraIntrinsics[0][0] /= scaleRes.x
    cameraIntrinsics[1][1] /= scaleRes.y

    cameraIntrinsics[2][0] /= scaleRes.x
    cameraIntrinsics[2][1] /= scaleRes.y

    // This will be the long size, because of the rotation
    let horizontalStep = Float(depthMapSize.x) / Float(self.horizontalPoints)
    let halfHorizontalStep = horizontalStep / 2
    // This will be the short size, because of the rotation
    let verticalStep = Float(depthMapSize.y) / Float(self.verticalPoints)
    let halfVerticalStep = verticalStep / 2

     for h in 0..<horizontalPoints {
        for v in 0..<verticalPoints {
            let x = Float(h) * horizontalStep + halfHorizontalStep
            let y = Float(v) * verticalStep + halfVerticalStep
            let depthMapPoint = simd_float2(x, y)

            // Sample depth
            let metricDepth = sampleDepthRaw(depthByteBuffer, size: depthMapSize, at: .init(depthMapPoint))
                    
            let wp = worldPoint(depthMapPixelPoint: depthMapPoint,
                                depth: metricDepth,
                                cameraIntrinsics: cameraIntrinsics,
                                // This is crucial: you need to always use the view matrix for Landscape Right.
                                viewMatrixInverted: frame.camera.viewMatrix(for: .landscapeRight).inverse)
            let node = self.depthNodes[v * horizontalPoints + h]
            node.simdWorldPosition = wp
        }
    }
}