Logitech Muse Laser Pen Scene Measurement (RealityKit & ARKit)

ImmersiveLaserPenMeasurementView.swift

import ARKit
import CoreHaptics
import GameController
import RealityKit
import SwiftUI

struct ImmersiveLaserPenMeasurementView: View {
    @State private var stylusManager = StylusLaserPenManager()

    var body: some View {
        RealityView { content in
            let root = Entity()
            content.add(root)
            stylusManager.rootEntity = root
            await stylusManager.handleControllerSetup()
            
            stylusManager.setupLaser()
            let _ = content.subscribe(to: SceneEvents.Update.self) { event in
                stylusManager.updateLaser()
           }
        }
        .task {
            // Don't forget to add the Accessory Tracking and World Sensing capabilities
            let configuration = SpatialTrackingSession.Configuration(tracking: [.accessory])
            let session = SpatialTrackingSession()
            await session.run(configuration)
            await stylusManager.startSceneReconstruction()
        }
    }
}

@Observable
final class MeasurementLabelModel {
    var text: String = "--"
}

struct MeasurementLabelPreviewView: View {
    @Bindable var model: MeasurementLabelModel
    var body: some View {
        MeasurementLabelStaticView(text: model.text)
    }
}

struct MeasurementLabelStaticView: View {
    let text: String
    var body: some View {
        Text(text)
            .font(.system(size: 48))
            .padding(20)
            .glassBackgroundEffect()
            .cornerRadius(12)
    }
}


@MainActor
final class StylusLaserPenManager {
    var rootEntity: Entity?
    var laserBeamEntity: Entity?
    var tipSphereEntity: Entity?
    private var anchors: [StylusAnchor: AnchorEntity] = [:]
    private var hapticEngines: [ObjectIdentifier: CHHapticEngine] = [:]
    private var hapticPlayers: [ObjectIdentifier: CHHapticPatternPlayer] = [:]

    private let arSession = ARKitSession()
    private let sceneReconstruction = SceneReconstructionProvider()
    private var meshEntities = [UUID: Entity]()

    private let unitHeight: Float = 1.0
    private let laserRadius: Float = 0.0015
    private let labelOffset: Float = 0.015
    
    var laserLength: Float = StylusLaserPenManager.maxLaserLength
    static var maxLaserLength: Float = 10.0

    private var currentHitPoint: SIMD3<Float>?
    private var currentRayDir: SIMD3<Float>?
    private var firstMeasurePoint: SIMD3<Float>?

    private var previewCylinder: ModelEntity?
    private var previewLabelEntity: Entity?
    private var previewLabelModel: MeasurementLabelModel?

    private var finalizedMeasurementsRoot = Entity()
    
    enum StylusAnchor: String {
        case aim
        case origin
    }
}

// Setup Stylus
extension StylusLaserPenManager {
    func handleControllerSetup() async {
        // Existing connections
        let styluses = GCStylus.styli

        for stylus in styluses where stylus.productCategory == GCProductCategorySpatialStylus {
            try? await setupAccessory(stylus: stylus)
        }

        NotificationCenter.default.addObserver(
            forName: NSNotification.Name.GCStylusDidConnect, object: nil, queue: .main
        ) { [weak self] note in
            guard let self,
                  let stylus = note.object as? GCStylus,
                  stylus.productCategory == GCProductCategorySpatialStylus else { return }
            Task { @MainActor in
                try? await self.setupAccessory(stylus: stylus)
            }
        }
    }
    
    private func setupAccessory(stylus: GCStylus) async throws {
        guard let root = rootEntity else { return }
        let source = try await AnchoringComponent.AccessoryAnchoringSource(device: stylus)
        
        // List available locations (aim and origin appear to be possible)
        print("📍 Available locations: \(source.accessoryLocations)")

        guard let aimLocation = source.locationName(named: StylusAnchor.aim.rawValue) else { return }

        let aimAnchor = AnchorEntity(
            .accessory(from: source, location: aimLocation),
            trackingMode: .continuous,
            physicsSimulation: .none
        )
        root.addChild(aimAnchor)

        let key = ObjectIdentifier(stylus)

        setupHaptics(for: stylus, key: key)
        setupStylusInputs(stylus: stylus, anchor: aimAnchor, key: key)
        addStylusTipIndicator(to: aimAnchor, color: .red)
        anchors[.aim] = aimAnchor
        
        guard let originLocation = source.locationName(named: StylusAnchor.origin.rawValue) else { return }
        let originAnchor = AnchorEntity(
            .accessory(from: source, location: originLocation),
            trackingMode: .continuous,
            physicsSimulation: .none
        )
        root.addChild(originAnchor)
        anchors[.origin] = originAnchor
    }
    
    private func addStylusTipIndicator(to anchor: AnchorEntity, color: UIColor) {
        let tipSphere = ModelEntity(
            mesh: .generateSphere(radius: 0.0015),
            materials: [UnlitMaterial(color: .red)]
        )
        // Add to root and store reference instead of attaching to anchor
        rootEntity?.addChild(tipSphere)
        tipSphereEntity = tipSphere
    }

    private func setupStylusInputs(stylus: GCStylus, anchor: AnchorEntity, key: ObjectIdentifier) {
        guard let input = stylus.input else { return }
        
        input.buttons[.stylusPrimaryButton]?.pressedInput.pressedDidChangeHandler = { [weak self] _, _, pressed in
            guard pressed, let self else { return }
            Task { @MainActor in
                self.playHaptic(for: key)
                self.toggleLaserBeam()
            }
        }
        
        input.buttons[.stylusSecondaryButton]?.pressedInput.pressedDidChangeHandler = { [weak self] _, _, pressed in
            print("stylusSecondaryButton pressedInput pressedDidChangeHandler \(pressed)")
            guard pressed, let self else { return }
            Task { @MainActor in
                self.playHaptic(for: key)
                self.handleMeasurementTap()
            }
        }
    }
}

// MARK: SceneReconstructionProvider
extension StylusLaserPenManager {
    func startSceneReconstruction() async {
        guard SceneReconstructionProvider.isSupported else {
            print("⚠️ SceneReconstructionProvider not supported")
            return
        }
        do {
            try await arSession.run([sceneReconstruction])
        } catch {
            print("❌ Failed to start ARKitSession: \(error)")
            return
        }

        Task { @MainActor [weak self] in
            guard let self else { return }
            for await update in self.sceneReconstruction.anchorUpdates {
                let meshAnchor = update.anchor
                guard let shape = try? await ShapeResource.generateStaticMesh(from: meshAnchor) else { continue }

                switch update.event {
                case .added:
                    let entity = Entity()
                    entity.transform = Transform(matrix: meshAnchor.originFromAnchorTransform)
                    entity.components.set(CollisionComponent(shapes: [shape], isStatic: true))
                    self.meshEntities[meshAnchor.id] = entity
                    self.rootEntity?.addChild(entity)

                case .updated:
                    guard let entity = self.meshEntities[meshAnchor.id] else { continue }
                    entity.transform = Transform(matrix: meshAnchor.originFromAnchorTransform)
                    if var collision = entity.components[CollisionComponent.self] {
                        collision.shapes = [shape]
                        entity.components.set(collision)
                    }

                case .removed:
                    self.meshEntities[meshAnchor.id]?.removeFromParent()
                    self.meshEntities.removeValue(forKey: meshAnchor.id)
                }
            }
        }
    }
}

// MARK: Haptics
extension StylusLaserPenManager {
    private func setupHaptics(for stylus: GCStylus, key: ObjectIdentifier) {
        guard let deviceHaptics = stylus.haptics else { return }
        
        // Create haptic engine
        let engine = deviceHaptics.createEngine(withLocality: .default)
        do {
            try engine?.start()
            hapticEngines[key] = engine
            
            // Create a simple "tap" pattern for button presses
            let pattern = try CHHapticPattern(events: [
                CHHapticEvent(eventType: .hapticTransient, parameters: [
                    CHHapticEventParameter(parameterID: .hapticIntensity, value: 0.8),
                    CHHapticEventParameter(parameterID: .hapticSharpness, value: 0.5)
                ], relativeTime: 0.0)
            ], parameters: [])
            
            let player = try engine?.makePlayer(with: pattern)
            hapticPlayers[key] = player
        } catch {
            print("❌ Failed to setup haptics: \(error)")
        }
    }
    
    private func playHaptic(for key: ObjectIdentifier) {
        guard let player = hapticPlayers[key] else { return }
        
        do {
            try player.start(atTime: CHHapticTimeImmediate)
        } catch {
            print("❌ Failed to play haptic: \(error)")
        }
    }
}

// MARK: - Laser Beam
extension StylusLaserPenManager {
    func setupLaser() {
        print("🔴 Setting up laser beam")
        let cylinderMesh = MeshResource.generateCylinder(height: unitHeight, radius: laserRadius)
        let material = UnlitMaterial(color: .red)
        let entity = ModelEntity(mesh: cylinderMesh, materials: [material])
        entity.components.set(OpacityComponent(opacity: 0.0))
        // Start at max length; scale along Y only to keep radius constant.
        entity.scale = SIMD3<Float>(1, laserLength, 1)
        laserBeamEntity = entity
        rootEntity?.addChild(entity)
        print("✅ Laser beam entity created and added to scene")
    }

    func updateLaser() {
        guard let entity = laserBeamEntity else { return }
        guard let aimAnchor = anchors[.aim], let originAnchor = anchors[.origin] else { return }
        guard
            let startCol = originAnchor.transformMatrix(relativeTo: .immersiveSpace)?.columns.3,
            let aimCol   = aimAnchor.transformMatrix(relativeTo: .immersiveSpace)?.columns.3
        else { return }

        let start = SIMD3<Float>(startCol.x, startCol.y, startCol.z)
        let aim   = SIMD3<Float>(aimCol.x, aimCol.y, aimCol.z)

        var dir = aim - start
        let dirLen = simd_length(dir)
        if dirLen < 1e-5 { return }
        dir /= dirLen

        // Raycast against the scene (in world space).
        let hits = entity.scene?.raycast(
            origin: aim,
            direction: dir,
            length: StylusLaserPenManager.maxLaserLength,
            query: .nearest,
            mask: .all,
            relativeTo: nil
        ) ?? []

        let newLength = hits.first?.distance ?? StylusLaserPenManager.maxLaserLength
        laserLength = newLength

        // Cache current collision point for measurement preview.
        if hits.first != nil {
            currentHitPoint = aim + dir * newLength
            currentRayDir = dir
        } else {
            currentHitPoint = nil
            currentRayDir = nil
        }

        // Scale unit-height cylinder along +Y to match the desired length.
        entity.scale = SIMD3<Float>(1, max(newLength, 0.0001), 1)

        // Place cylinder's center halfway along the ray.
        entity.position = aim + dir * (newLength * 0.5)

        // Rotate so local +Y aligns with the ray direction.
        let up = SIMD3<Float>(0, 1, 0)
        let dotUD = max(-1.0, min(1.0, simd_dot(up, dir)))
        let angle = acosf(dotUD)
        let axis: SIMD3<Float> = simd_length_squared(simd_cross(up, dir)) < 1e-10
            ? SIMD3<Float>(1, 0, 0) // parallel or antiparallel; choose X
            : simd_normalize(simd_cross(up, dir))
        entity.transform.rotation = simd_quaternion(angle, axis)

        // Move tip indicator to ray end (aim-based).
        if let sphere = tipSphereEntity {
            sphere.position = aim + dir * newLength
        }

        // Live update measurement preview if active.
        if let startPoint = firstMeasurePoint, let endPoint = currentHitPoint {
            updateMeasurementPreview(from: startPoint, to: endPoint)
        }
    }
    
    func toggleLaserBeam() {
        guard let entity = laserBeamEntity else {
            print("⚠️ toggleLaserBeam called but laserBeamEntity is nil")
            return
        }
        
        guard let opacityComponent = entity.components[OpacityComponent.self] else {
            print("⚠️ toggleLaserBeam called but OpacityComponent is nil")
            return
        }
        
        if opacityComponent.opacity == 1 {
            entity.components.set(OpacityComponent(opacity: 0))
        } else {
            entity.components.set(OpacityComponent(opacity: 1))
        }
    }
}

// MARK: - Measurement
extension StylusLaserPenManager {
    private func handleMeasurementTap() {
        guard let root = rootEntity else { return }

        if finalizedMeasurementsRoot.parent == nil {
            root.addChild(finalizedMeasurementsRoot)
        }

        if firstMeasurePoint == nil {
            guard let hit = currentHitPoint else {
                print("⚠️ No hit to start measurement")
                return
            }
            firstMeasurePoint = hit
            let sphere = makeMarkerSphere(color: .systemBlue, radius: 0.006)
            sphere.position = hit
            root.addChild(sphere)
            ensurePreviewEntities()
            updateMeasurementPreview(from: hit, to: hit)
            return
        }

        guard let start = firstMeasurePoint, let end = currentHitPoint else {
            print("⚠️ No hit to finalize measurement")
            return
        }

        finalizeMeasurement(from: start, to: end)

        previewCylinder?.removeFromParent()
        previewCylinder = nil
        previewLabelEntity?.removeFromParent()
        previewLabelEntity = nil
        previewLabelModel = nil
        firstMeasurePoint = nil
    }
    
    private func ensurePreviewEntities() {
        guard let root = rootEntity else { return }
        if previewCylinder == nil {
            let cylinder = ModelEntity(mesh: .generateCylinder(height: 1.0, radius: 0.002),
                                       materials: [UnlitMaterial(color: .white)])
            root.addChild(cylinder)
            previewCylinder = cylinder
        }
        if previewLabelEntity == nil {
            let labelModel = MeasurementLabelModel()
            previewLabelModel = labelModel
            let label = Entity()
            label.components.set(ViewAttachmentComponent(rootView: MeasurementLabelPreviewView(model: labelModel)))
            root.addChild(label)
            previewLabelEntity = label
        }
    }

    private func updateMeasurementPreview(from: SIMD3<Float>, to: SIMD3<Float>) {
        guard let cylinder = previewCylinder, let label = previewLabelEntity else { return }
        let vector = to - from
        let length = max(simd_length(vector), 0.0001)
        let direction = simd_normalize(vector)
        let mid = (from + to) * 0.5

        cylinder.scale = [1, length, 1]
        cylinder.position = mid
        cylinder.transform.rotation = rotationAligningUp(to: direction)

        previewLabelModel?.text = formattedDistance(length)
        let forward = orientLabelAlongMeasurement(label, from: from, to: to)
        label.position = mid + forward * labelOffset
    }

    private func finalizeMeasurement(from: SIMD3<Float>, to: SIMD3<Float>) {
        let vector = to - from
        let length = simd_length(vector)
        let direction = simd_normalize(vector)
        let mid = (from + to) * 0.5

        let group = Entity()

        let sphereA = makeMarkerSphere(color: .systemBlue, radius: 0.006)
        sphereA.position = from
        group.addChild(sphereA)

        let sphereB = makeMarkerSphere(color: .systemBlue, radius: 0.006)
        sphereB.position = to
        group.addChild(sphereB)

        let cylinder = ModelEntity(mesh: .generateCylinder(height: 1.0, radius: 0.002),
                                   materials: [UnlitMaterial(color: .white)])
        cylinder.scale = [1, max(length, 0.0001), 1]
        cylinder.position = mid
        cylinder.transform.rotation = rotationAligningUp(to: direction)
        group.addChild(cylinder)

        let label = Entity()
        label.components.set(ViewAttachmentComponent(rootView: MeasurementLabelStaticView(text: formattedDistance(length))))
        let forward = orientLabelAlongMeasurement(label, from: from, to: to)
        label.position = mid + forward * labelOffset
        group.addChild(label)

        finalizedMeasurementsRoot.addChild(group)
    }

    private func makeMarkerSphere(color: UIColor, radius: Float) -> ModelEntity {
        ModelEntity(mesh: .generateSphere(radius: radius),
                    materials: [UnlitMaterial(color: color)])
    }

    private func rotationAligningUp(to direction: SIMD3<Float>) -> simd_quatf {
        let up = SIMD3<Float>(0, 1, 0)
        let dot = max(-1.0, min(1.0, simd_dot(up, direction)))
        let angle = acosf(dot)
        let axis: SIMD3<Float> = simd_length_squared(simd_cross(up, direction)) < 1e-10
            ? SIMD3<Float>(1, 0, 0)
            : simd_normalize(simd_cross(up, direction))
        return simd_quaternion(angle, axis)
    }

    private func formattedDistance(_ meters: Float) -> String {
        if meters >= 1.0 {
            return String(format: "%.2f m", meters)
        } else {
            return String(format: "%.1f cm", meters * 100.0)
        }
    }

    // Align label ALONG the measurement, left-to-right relative to viewer
    // We choose the label's right vector as the measurement direction, but if the
    // viewer is looking such that this would read right-to-left, we flip it.
    private func orientLabelAlongMeasurement(_ label: Entity, from: SIMD3<Float>, to: SIMD3<Float>) -> SIMD3<Float> {
        let worldUp = SIMD3<Float>(0, 1, 0)
        var right = simd_normalize(to - from) // text flows along +X of label
        // Approximate viewer forward from stylus ray; if unavailable, use camera-like forward (0,0,-1)
        let viewerForward = currentRayDir.map { -$0 } ?? SIMD3<Float>(0, 0, -1)
        // Ensure left-to-right: if viewer sees the line decreasing left-to-right, flip it
        if simd_dot(simd_cross(worldUp, viewerForward), right) < 0 {
            right = -right
        }
        var forward = simd_normalize(simd_cross(right, worldUp))
        if simd_length_squared(forward) < 1e-6 {
            // Degenerate if near horizontal alignment; fallback to face viewer directly
            forward = simd_normalize(viewerForward)
        }
        var up = simd_normalize(simd_cross(forward, right))
        if simd_dot(up, worldUp) < 0 {
            up = -up
            forward = -forward
        }
        let basis = simd_float3x3(columns: (right, up, forward))
        label.transform.rotation = simd_quaternion(basis)
        return forward
    }
}