Blender4iOS/Sources/Animation/AnimationData.swift

// AnimationData.swift
// Core animation data model — mirrors Blender's FCurve / Keyframe / Action system.

import Observation
import simd

// MARK: - Interpolation & Easing (matches Blender's enum)

enum KeyframeInterpolation: String, CaseIterable, Identifiable {
    case constant = "Constant"      // BEZT_IPO_CONST — holds value until next keyframe
    case linear   = "Linear"        // BEZT_IPO_LIN  — straight line
    case bezier   = "Bézier"        // BEZT_IPO_BEZ  — cubic Bézier
    case bounce   = "Bounce"
    case elastic  = "Elastic"
    case back     = "Back"

    var id: String { rawValue }
}

enum HandleType: String, CaseIterable, Identifiable {
    case free      = "Free"         // independent handle control
    case aligned   = "Aligned"      // handles are co-linear (mirrored angle)
    case vector    = "Vector"       // auto-computed to point at neighbor
    case auto      = "Auto"         // auto-smooth (Blender default)
    case autoClamp = "Auto Clamped" // auto but clamped to prevent overshoot

    var id: String { rawValue }
}

// MARK: - Keyframe

@Observable
final class Keyframe: Identifiable {
    let id = UUID()

    /// Frame number (Blender uses integer frames, we allow float for sub-frame)
    var frame: Float

    /// Value at this keyframe
    var value: Float

    /// Left handle (incoming tangent), in absolute coordinates
    var handleLeft: SIMD2<Float>

    /// Right handle (outgoing tangent), in absolute coordinates
    var handleRight: SIMD2<Float>

    /// Handle type
    var handleType: HandleType = .auto

    /// Interpolation mode TO the next keyframe
    var interpolation: KeyframeInterpolation = .bezier

    /// Selection state
    var isSelected: Bool = false
    var isHandleLeftSelected: Bool = false
    var isHandleRightSelected: Bool = false

    init(frame: Float, value: Float) {
        self.frame = frame
        self.value = value
        // Default handles: horizontal, 1/3 of typical frame spacing
        self.handleLeft  = SIMD2<Float>(frame - 5, value)
        self.handleRight = SIMD2<Float>(frame + 5, value)
    }

    /// Auto-compute handles for smooth interpolation (Blender's "Auto" mode).
    func autoComputeHandles(prev: Keyframe?, next: Keyframe?) {
        guard handleType == .auto || handleType == .autoClamp else { return }

        let dt: Float = 5.0
        if let prev, let next {
            // Slope through neighbors
            let slope = (next.value - prev.value) / (next.frame - prev.frame)
            handleLeft  = SIMD2<Float>(frame - dt, value - slope * dt)
            handleRight = SIMD2<Float>(frame + dt, value + slope * dt)

            // Auto-clamp: prevent overshoot
            if handleType == .autoClamp {
                let minV = min(prev.value, next.value)
                let maxV = max(prev.value, next.value)
                handleLeft.y  = handleLeft.y.clamped(to: minV...maxV)
                handleRight.y = handleRight.y.clamped(to: minV...maxV)
            }
        } else if let prev {
            let slope = (value - prev.value) / (frame - prev.frame)
            handleLeft = SIMD2<Float>(frame - dt, value - slope * dt)
            handleRight = SIMD2<Float>(frame + dt, value)
        } else if let next {
            let slope = (next.value - value) / (next.frame - frame)
            handleLeft = SIMD2<Float>(frame - dt, value)
            handleRight = SIMD2<Float>(frame + dt, value + slope * dt)
        } else {
            handleLeft  = SIMD2<Float>(frame - dt, value)
            handleRight = SIMD2<Float>(frame + dt, value)
        }
    }

    /// Move handle maintaining aligned constraint if needed.
    func moveHandleLeft(to pos: SIMD2<Float>) {
        handleLeft = pos
        if handleType == .aligned {
            let delta = SIMD2<Float>(frame, value) - pos
            handleRight = SIMD2<Float>(frame, value) + delta
        }
    }

    func moveHandleRight(to pos: SIMD2<Float>) {
        handleRight = pos
        if handleType == .aligned {
            let delta = SIMD2<Float>(frame, value) - pos
            handleLeft = SIMD2<Float>(frame, value) + delta
        }
    }
}

extension Comparable {
    func clamped(to range: ClosedRange<Self>) -> Self {
        min(max(self, range.lowerBound), range.upperBound)
    }
}

// MARK: - Animation Track (single F-Curve)

/// Represents one animated channel, e.g. "location.x" or "rotation.z".
/// Mirrors Blender's FCurve.
@Observable
final class AnimationTrack: Identifiable {
    let id = UUID()
    let channelPath: String          // e.g. "location", "rotation_euler"
    let channelIndex: Int            // 0=X, 1=Y, 2=Z
    var keyframes: [Keyframe] = []
    var isMuted: Bool = false
    var isLocked: Bool = false
    var isVisible: Bool = true       // visibility in graph editor

    /// Display color matching Blender convention
    var color: SIMD4<Float> {
        switch channelIndex {
        case 0: SIMD4(0.86, 0.20, 0.27, 1) // X = red
        case 1: SIMD4(0.30, 0.69, 0.31, 1) // Y = green
        case 2: SIMD4(0.28, 0.46, 0.86, 1) // Z = blue
        default: SIMD4(0.7, 0.7, 0.7, 1)   // W/custom = grey
        }
    }

    var displayName: String {
        let axis = ["X", "Y", "Z", "W"]
        let idx = channelIndex < axis.count ? axis[channelIndex] : "\(channelIndex)"
        return "\(channelPath).\(idx)"
    }

    init(channelPath: String, channelIndex: Int) {
        self.channelPath = channelPath
        self.channelIndex = channelIndex
    }

    /// Sorted keyframes (must be sorted for evaluation).
    var sortedKeyframes: [Keyframe] {
        keyframes.sorted { $0.frame < $1.frame }
    }

    /// Insert a keyframe, auto-computing handles.
    func insertKeyframe(frame: Float, value: Float) {
        // Remove existing keyframe at same frame
        keyframes.removeAll { abs($0.frame - frame) < 0.001 }

        let kf = Keyframe(frame: frame, value: value)
        keyframes.append(kf)
        recomputeAutoHandles()
    }

    /// Evaluate the curve at a given frame using cubic Bézier interpolation.
    func evaluate(at frame: Float) -> Float {
        let sorted = sortedKeyframes
        guard !sorted.isEmpty else { return 0 }

        // Before first keyframe
        if frame <= sorted.first!.frame { return sorted.first!.value }
        // After last keyframe
        if frame >= sorted.last!.frame { return sorted.last!.value }

        // Find surrounding keyframes
        for i in 0..<(sorted.count - 1) {
            let kf0 = sorted[i]
            let kf1 = sorted[i + 1]
            if frame >= kf0.frame && frame <= kf1.frame {
                return interpolate(from: kf0, to: kf1, at: frame)
            }
        }

        return sorted.last!.value
    }

    // MARK: - Bézier interpolation

    private func interpolate(from kf0: Keyframe, to kf1: Keyframe, at frame: Float) -> Float {
        switch kf0.interpolation {
        case .constant:
            return kf0.value

        case .linear:
            let t = (frame - kf0.frame) / (kf1.frame - kf0.frame)
            return kf0.value + (kf1.value - kf0.value) * t

        case .bezier:
            return bezierEvaluate(kf0: kf0, kf1: kf1, frame: frame)

        case .bounce:
            let t = (frame - kf0.frame) / (kf1.frame - kf0.frame)
            let bt = bounceEaseOut(t)
            return kf0.value + (kf1.value - kf0.value) * bt

        case .elastic:
            let t = (frame - kf0.frame) / (kf1.frame - kf0.frame)
            let et = elasticEaseOut(t)
            return kf0.value + (kf1.value - kf0.value) * et

        case .back:
            let t = (frame - kf0.frame) / (kf1.frame - kf0.frame)
            let bt = backEaseOut(t)
            return kf0.value + (kf1.value - kf0.value) * bt
        }
    }

    /// Evaluate cubic Bézier — solves for t given x, then evaluates y.
    /// Matches Blender's BKE_fcurve_correct_bezpart approach.
    private func bezierEvaluate(kf0: Keyframe, kf1: Keyframe, frame: Float) -> Float {
        let p0 = SIMD2<Float>(kf0.frame, kf0.value)
        let p1 = kf0.handleRight
        let p2 = kf1.handleLeft
        let p3 = SIMD2<Float>(kf1.frame, kf1.value)

        // Find t for the given frame (x coordinate) using Newton-Raphson
        let t = solveBezierT(x: frame, p0x: p0.x, p1x: p1.x, p2x: p2.x, p3x: p3.x)

        // Evaluate y at found t
        let oneMinusT = 1.0 - t
        let y = oneMinusT * oneMinusT * oneMinusT * p0.y
              + 3 * oneMinusT * oneMinusT * t * p1.y
              + 3 * oneMinusT * t * t * p2.y
              + t * t * t * p3.y

        return y
    }

    /// Newton-Raphson solver to find parameter t for a given x on the Bézier.
    private func solveBezierT(x: Float, p0x: Float, p1x: Float, p2x: Float, p3x: Float) -> Float {
        // Initial guess: linear
        var t = (x - p0x) / max(p3x - p0x, 0.0001)
        t = t.clamped(to: 0...1)

        // 8 iterations of Newton-Raphson (matches Blender's precision)
        for _ in 0..<8 {
            let omt = 1.0 - t
            let currentX = omt * omt * omt * p0x
                         + 3 * omt * omt * t * p1x
                         + 3 * omt * t * t * p2x
                         + t * t * t * p3x

            let dx = 3 * omt * omt * (p1x - p0x)
                   + 6 * omt * t * (p2x - p1x)
                   + 3 * t * t * (p3x - p2x)

            if abs(dx) < 1e-8 { break }
            t -= (currentX - x) / dx
            t = t.clamped(to: 0...1)
        }

        return t
    }

    // MARK: - Easing functions

    private func bounceEaseOut(_ t: Float) -> Float {
        if t < 1.0 / 2.75 {
            return 7.5625 * t * t
        } else if t < 2.0 / 2.75 {
            let t2 = t - 1.5 / 2.75
            return 7.5625 * t2 * t2 + 0.75
        } else if t < 2.5 / 2.75 {
            let t2 = t - 2.25 / 2.75
            return 7.5625 * t2 * t2 + 0.9375
        } else {
            let t2 = t - 2.625 / 2.75
            return 7.5625 * t2 * t2 + 0.984375
        }
    }

    private func elasticEaseOut(_ t: Float) -> Float {
        if t <= 0 { return 0 }
        if t >= 1 { return 1 }
        return powf(2, -10 * t) * sinf((t - 0.075) * (2 * .pi) / 0.3) + 1
    }

    private func backEaseOut(_ t: Float) -> Float {
        let s: Float = 1.70158
        let t2 = t - 1
        return t2 * t2 * ((s + 1) * t2 + s) + 1
    }

    // MARK: - Handle computation

    func recomputeAutoHandles() {
        let sorted = sortedKeyframes
        for i in 0..<sorted.count {
            let prev = i > 0 ? sorted[i - 1] : nil
            let next = i < sorted.count - 1 ? sorted[i + 1] : nil
            sorted[i].autoComputeHandles(prev: prev, next: next)
        }
    }
}

// MARK: - Action (collection of tracks for one object)

/// Mirrors Blender's bAction — a reusable animation data block.
@Observable
final class AnimAction: Identifiable {
    let id = UUID()
    var name: String
    var tracks: [AnimationTrack] = []

    /// Frame range
    var frameStart: Float { tracks.flatMap { $0.sortedKeyframes }.map(\.frame).min() ?? 1 }
    var frameEnd: Float   { tracks.flatMap { $0.sortedKeyframes }.map(\.frame).max() ?? 250 }

    init(name: String) {
        self.name = name
    }

    func track(path: String, index: Int) -> AnimationTrack {
        if let existing = tracks.first(where: { $0.channelPath == path && $0.channelIndex == index }) {
            return existing
        }
        let new = AnimationTrack(channelPath: path, channelIndex: index)
        tracks.append(new)
        return new
    }

    /// Evaluate all tracks at a given frame, returns a dictionary of values.
    func evaluate(at frame: Float) -> [String: [Float]] {
        var result: [String: [Float]] = [:]
        for track in tracks where !track.isMuted {
            var arr = result[track.channelPath] ?? [0, 0, 0]
            if track.channelIndex < arr.count {
                arr[track.channelIndex] = track.evaluate(at: frame)
            }
            result[track.channelPath] = arr
        }
        return result
    }
}

// MARK: - Animation System

@Observable
final class AnimationSystem {
    var actions: [AnimAction] = []
    var activeAction: AnimAction?

    // Playback state
    var currentFrame: Float = 1
    var isPlaying: Bool = false
    var fps: Float = 24
    var frameStart: Float = 1
    var frameEnd: Float = 250
    var isLooping: Bool = true

    // Display
    var showDopeSheet: Bool = true  // false = show graph editor

    init() {
        // Create a default action with sample keyframes
        let action = AnimAction(name: "CubeAction")

        let locX = action.track(path: "location", index: 0)
        locX.insertKeyframe(frame: 1,   value: 0)
        locX.insertKeyframe(frame: 30,  value: 3)
        locX.insertKeyframe(frame: 60,  value: -2)
        locX.insertKeyframe(frame: 100, value: 0)

        let locY = action.track(path: "location", index: 1)
        locY.insertKeyframe(frame: 1,   value: 0)
        locY.insertKeyframe(frame: 50,  value: 4)
        locY.insertKeyframe(frame: 100, value: 0)

        let locZ = action.track(path: "location", index: 2)
        locZ.insertKeyframe(frame: 1,   value: 0)
        locZ.insertKeyframe(frame: 40,  value: 2)
        locZ.insertKeyframe(frame: 80,  value: -1)
        locZ.insertKeyframe(frame: 100, value: 0)

        actions.append(action)
        activeAction = action
    }

    /// Advance one frame (called from CADisplayLink).
    func tick(dt: Float) {
        guard isPlaying else { return }
        currentFrame += dt * fps
        if currentFrame > frameEnd {
            currentFrame = isLooping ? frameStart : frameEnd
            if !isLooping { isPlaying = false }
        }
    }

    /// Apply the current frame's animation to a SceneObject.
    func applyToObject(_ obj: SceneObject) {
        guard let action = activeAction else { return }
        let values = action.evaluate(at: currentFrame)

        if let loc = values["location"] {
            obj.transform.position = SIMD3<Float>(loc[0], loc[1], loc[2])
        }
        if let rot = values["rotation_euler"] {
            obj.transform.rotation = SIMD3<Float>(rot[0], rot[1], rot[2])
        }
        if let scl = values["scale"] {
            obj.transform.scale = SIMD3<Float>(scl[0], scl[1], scl[2])
        }
    }
}