Logitech Muse Anchor Tracking, Model Spawning, and Haptic Feedback (RealityKit)

ImmersiveStylusSimpleExampleView.swift

import CoreHaptics
import GameController
import RealityKit
import SwiftUI

// MARK: Immersive View
struct ImmersiveStylusSimpleExampleView: View {
    @State private var stylusManager = StylusManager()

    var body: some View {
        RealityView { content in
            let root = Entity()
            content.add(root)
            stylusManager.rootEntity = root
            await stylusManager.handleControllerSetup()
        }
        .task {
            // Don't forget to add the Accessory Tracking capability
            let configuration = SpatialTrackingSession.Configuration(tracking: [.accessory])
            let session = SpatialTrackingSession()
            await session.run(configuration)
        }
    }
}

// MARK: StylusManager
@MainActor
final class StylusManager {
    var rootEntity: Entity?
    private var hapticEngines: [ObjectIdentifier: CHHapticEngine] = [:]
    private var hapticPlayers: [ObjectIdentifier: CHHapticPatternPlayer] = [:]
    private var inflationHapticPlayers: [ObjectIdentifier: CHHapticPatternPlayer] = [:]
    private var activePressureModel: ModelEntity?
    private var activeModelColor: UIColor = .systemRed
    private var maxScale: Float = 0.5
    private let baseRadius: Float = 0.01
    private var tipIndicators: [AnchorEntity: Entity] = [:]
}

// MARK: Setup Stylus
extension StylusManager {
    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)
        print("📍 Available locations: \(source.accessoryLocations)")
        
        guard let location = source.locationName(named: "aim") else { return }
        
        let anchor = AnchorEntity(
            .accessory(from: source, location: location),
            trackingMode: .predicted,
            physicsSimulation: .none
        )
        root.addChild(anchor)
        
        let key = ObjectIdentifier(stylus)
        
        // Setup haptics if available
        setupHaptics(for: stylus, key: key)
        setupStylusInputs(stylus: stylus, anchor: anchor, key: key)
        addAxisIndicator(to: anchor)
    }
    
}

// MARK: Inputs (Button Handlers + Tip Handler)
extension StylusManager {
    private func setupStylusInputs(stylus: GCStylus, anchor: AnchorEntity, key: ObjectIdentifier) {
        guard let input = stylus.input else { return }
        
        // Handle main button (simple button pressed)
        input.buttons[.stylusPrimaryButton]?.pressedInput.pressedDidChangeHandler = { [weak self] _, _, pressed in
            guard pressed, let self else { return }
            Task { @MainActor in
                self.playHaptic(for: key)
                self.spawnModel(at: anchor, color: .systemBlue, radius: 0.02)
            }
        }
        
        // Handle secondary button pressure changes - create and increase scale based on pressure
        input.buttons[.stylusSecondaryButton]?.pressedInput.valueDidChangeHandler = { [weak self] _, _, pressure in
            guard let self else { return }
            Task { @MainActor in
                if pressure > 0.05 {
                    self.updatePressureModel(at: anchor, pressure: pressure, color: .systemRed, key: key)
                } else if pressure <= 0 {
                    self.dropPressureModel(at: anchor)
                }
            }
        }
        
        // Handle tip pressure changes - create and increase scale based on pressure
        input.buttons[.stylusTip]?.pressedInput.valueDidChangeHandler = { [weak self] _, _, pressure in
            guard let self else { return }
            Task { @MainActor in
                if pressure > 0.05 {
                    self.updatePressureModel(at: anchor, pressure: pressure, color: .systemGreen, key: key)
                } else if pressure <= 0 {
                    self.dropPressureModel(at: anchor)
                }
            }
        }
    }
}

// MARK: Haptics
extension StylusManager {
    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 main button press
            let tapPattern = try CHHapticPattern(events: [
                CHHapticEvent(eventType: .hapticTransient, parameters: [
                    CHHapticEventParameter(parameterID: .hapticIntensity, value: 0.85),
                    CHHapticEventParameter(parameterID: .hapticSharpness, value: 0.5)
                ], relativeTime: 0.0)
            ], parameters: [])
            
            let player = try engine?.makePlayer(with: tapPattern)
            hapticPlayers[key] = player
            
            // Create an "inflation" pattern for tip and secondary pressure increase
            let inflationPattern = try CHHapticPattern(events: [
                CHHapticEvent(eventType: .hapticContinuous, parameters: [
                    CHHapticEventParameter(parameterID: .hapticIntensity, value: 0.55),
                    CHHapticEventParameter(parameterID: .hapticSharpness, value: 1.0)
                ], relativeTime: 0.0, duration: 0.1)
            ], parameters: [])
            
            let inflationPlayer = try engine?.makePlayer(with: inflationPattern)
            inflationHapticPlayers[key] = inflationPlayer
        } 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)")
        }
    }
    
    private func playInflationHaptic(for key: ObjectIdentifier) {
        guard let player = inflationHapticPlayers[key] else { return }
        
        do {
            try player.start(atTime: CHHapticTimeImmediate)
        } catch {
            print("❌ Failed to play inflation haptic: \(error)")
        }
    }
}

// MARK: Tip Indicator
extension StylusManager {
    private func addAxisIndicator(to anchor: AnchorEntity) {
        let length: Float = 0.01
        let thickness: Float = 0.001
        let originSize: Float = 0.0015
        
        // Origin cube (white)
        let originCube = ModelEntity(
            mesh: .generateBox(size: SIMD3<Float>(originSize, originSize, originSize)),
            materials: [SimpleMaterial(color: .white, isMetallic: false)]
        )
        
        // X axis (red)
        let xAxis = ModelEntity(
            mesh: .generateBox(size: SIMD3<Float>(length, thickness, thickness)),
            materials: [SimpleMaterial(color: .systemRed, isMetallic: false)]
        )
        xAxis.position = SIMD3<Float>(-length / 2, 0, 0)
        
        // Y axis (green)
        let yAxis = ModelEntity(
            mesh: .generateBox(size: SIMD3<Float>(thickness, length, thickness)),
            materials: [SimpleMaterial(color: .systemGreen, isMetallic: false)]
        )
        yAxis.position = SIMD3<Float>(0, -length / 2, 0)
        
        // Z axis (blue)
        let zAxis = ModelEntity(
            mesh: .generateBox(size: SIMD3<Float>(thickness, thickness, length)),
            materials: [SimpleMaterial(color: .systemBlue, isMetallic: false)]
        )
        zAxis.position = SIMD3<Float>(0, 0, -length / 2)
        
        let axisRoot = Entity()
        axisRoot.addChild(originCube)
        axisRoot.addChild(xAxis)
        axisRoot.addChild(yAxis)
        axisRoot.addChild(zAxis)
        
        anchor.addChild(axisRoot)
        
        tipIndicators[anchor] = axisRoot
    }
    
}

// MARK: Basic Model Spawning (Main Button)
extension StylusManager {
    private func spawnModel(at anchor: AnchorEntity, color: UIColor, radius: Float) {
        guard let root = rootEntity else { return }
        
        let worldTransform = anchor.transformMatrix(relativeTo: nil)
        let worldPosition = SIMD3<Float>(worldTransform.columns.3.x,
                                         worldTransform.columns.3.y,
                                         worldTransform.columns.3.z)
        let worldOrientation = anchor.orientation(relativeTo: nil)
        
        let cube = ModelEntity(
            mesh: .generateBox(size: SIMD3<Float>(radius * 2, radius * 2, radius * 2)),
            materials: [SimpleMaterial(color: color, isMetallic: false)]
        )
        cube.position = worldPosition
        cube.orientation = worldOrientation
        
        root.addChild(cube)
    }
}

// MARK: Pressure Model Spawning (Secondary Button & Tip)
extension StylusManager {
    private func updatePressureModel(at anchor: AnchorEntity, pressure: Float, color: UIColor, key: ObjectIdentifier) {
        if activePressureModel == nil {
            // Hide the tip indicator while inflating
            tipIndicators[anchor]?.isEnabled = false
            
            // Create new cube
            let cube = ModelEntity(
                mesh: .generateBox(size: SIMD3<Float>(baseRadius * 2, baseRadius * 2, baseRadius * 2)),
                materials: [SimpleMaterial(color: color, isMetallic: false)]
            )
            anchor.addChild(cube)
            activePressureModel = cube
            activeModelColor = color
            maxScale = 0.5  // Reset max scale for new object
        }
        
        // Scale based on pressure (0.0 to 1.0)
        let scale = 0.5 + (pressure * 5.5)
        
        // Only scale UP, never down
        if scale > maxScale {
            maxScale = scale
            activePressureModel?.scale = SIMD3<Float>(repeating: maxScale)
            
            // Offset by its scaled radius in the -Z direction (forward from stylus tip)
            // This creates the "balloon blowing" effect
            let offset = baseRadius * maxScale
            activePressureModel?.position = SIMD3<Float>(0, 0, -offset)
            
            // Play subtle haptic feedback while inflating
            playInflationHaptic(for: key)
            
            let colorEmoji = color == .systemRed ? "🔴" : "🟢"
            print("\(colorEmoji) Pressure: \(String(format: "%.2f", pressure)), Scale: \(String(format: "%.2f", maxScale))")
        }
    }
    
    private func dropPressureModel(at anchor: AnchorEntity) {
        guard let model = activePressureModel, let root = rootEntity else { return }
        
        // Get the current world position before detaching
        let worldTransform = model.transformMatrix(relativeTo: nil)
        let worldPosition = SIMD3<Float>(worldTransform.columns.3.x,
                                          worldTransform.columns.3.y,
                                          worldTransform.columns.3.z)
        let worldScale = model.scale
        let worldOrientation = model.orientation(relativeTo: nil)
        
        model.removeFromParent()

        let droppedModel = ModelEntity(
            mesh: .generateBox(size: SIMD3<Float>(0.02, 0.02, 0.02)),
            materials: [SimpleMaterial(color: activeModelColor, isMetallic: false)]
        )
        droppedModel.position = worldPosition
        droppedModel.orientation = worldOrientation
        droppedModel.scale = worldScale
        root.addChild(droppedModel)
        
        print("💧 Dropped cube at scale \(worldScale)")
        
        // Show the tip indicator again
        tipIndicators[anchor]?.isEnabled = true
        
        // Clear the reference
        activePressureModel = nil
    }
}