//
// MotionSolver.swift
// Workouts
//
// Copyright 2026 Rouslan Zenetl. All Rights Reserved.
//

import CoreGraphics
import Foundation

/// Swift port of the Exercise Library's reference solver (`Exercise Library/render.py`):
/// forward kinematics, analytic 2-bone IK, and angle-space tweening with shortest-path
/// lerp and ease-in-out. The math is kept 1:1 with the Python so both renderers produce
/// the same figure from the same data — change them in lockstep.
///
/// Angles are absolute world angles in degrees, y-up (0 = right, 90 = up, −90 = down);
/// points are y-down 320×180 canvas coordinates, so dir(θ) = (cos θ, −sin θ).

/// The four two-bone limbs, keyed by their motion-script names.
enum FigureLimb: String, CaseIterable {
    case armR = "arm_r"
    case armL = "arm_l"
    case legR = "leg_r"
    case legL = "leg_l"

    var isLeft: Bool { self == .armL || self == .legL }

    /// The key a planted extremity uses in a key frame's `pins`.
    var pinKey: String {
        switch self {
        case .armR: "hand_r"
        case .armL: "hand_l"
        case .legR: "foot_r"
        case .legL: "foot_l"
        }
    }

    /// `[upper, lower]` bone lengths for this limb.
    func boneLengths(_ body: ExerciseBodyProfile) -> [Double] {
        switch self {
        case .armR, .armL: [body.upperArm, body.foreArm]
        case .legR, .legL: [body.thigh, body.shin]
        }
    }

    /// The authored `[upper, lower]` angles for this limb, if present in the frame.
    func angles(in frame: MotionKeyFrame) -> [Double]? {
        switch self {
        case .armR: frame.armR
        case .armL: frame.armL
        case .legR: frame.legR
        case .legL: frame.legL
        }
    }
}

/// A key frame resolved to pure angles (pinned limbs replaced by their IK solution),
/// so poses interpolate cleanly in angle space.
struct FigurePose {
    var root: CGPoint
    var spine: [Double]
    var neck: Double
    /// Absent for figures facing the viewer — no nose tick.
    var gaze: Double?
    var limbs: [FigureLimb: [Double]]
    var pins: [String: CGPoint]
    var hold: Double
    var tween: Double
}

/// Trunk FK: pelvis, spine mid, neck joint, and the offset left-limb attachments.
struct FigureAttachments {
    let pelvis: CGPoint
    let mid: CGPoint
    let neck: CGPoint
    let shoulderL: CGPoint
    let hipL: CGPoint

    func point(for limb: FigureLimb) -> CGPoint {
        switch limb {
        case .armR: neck
        case .armL: shoulderL
        case .legR: pelvis
        case .legL: hipL
        }
    }
}

/// A pose resolved to drawable points.
struct FigureGeometry {
    var headCenter: CGPoint
    var headRadius: Double
    /// Nil when the pose has no gaze (face-on view).
    var noseStart: CGPoint?
    var noseEnd: CGPoint?
    /// Quadratic Bézier through pelvis → mid → neck (control = 2·mid − (pelvis+neck)/2).
    var spineStart: CGPoint
    var spineControl: CGPoint
    var spineEnd: CGPoint
    /// Attachment → elbow/knee → extremity, per drawn limb.
    var limbs: [FigureLimb: [CGPoint]]
}

enum MotionSolver {
    /// Unit vector for a y-up angle, in y-down canvas coordinates.
    static func direction(_ degrees: Double) -> CGVector {
        let r = degrees * .pi / 180
        return CGVector(dx: cos(r), dy: -sin(r))
    }

    /// Chain FK: `[start, joint1, joint2, …]` for per-bone absolute angles.
    static func walk(from start: CGPoint, angles: [Double], lengths: [Double]) -> [CGPoint] {
        var points = [start]
        var current = start
        for (angle, length) in zip(angles, lengths) {
            let d = direction(angle)
            current = CGPoint(x: current.x + d.dx * length, y: current.y + d.dy * length)
            points.append(current)
        }
        return points
    }

    /// Y-up world angle (degrees) of the segment a→b.
    static func angle(from a: CGPoint, to b: CGPoint) -> Double {
        atan2(-(b.y - a.y), b.x - a.x) * 180 / .pi
    }

    /// Analytic 2-bone IK: `[upper, lower]` angles reaching from `start` toward
    /// `target`, choosing the elbow/knee solution nearest the authored guess.
    static func ik2(start: CGPoint, target: CGPoint, upper a: Double, lower b: Double, guess: [Double]) -> [Double] {
        let dx = target.x - start.x
        let dyUp = -(target.y - start.y)
        // Clamp the reach inside the chain's annulus so acos stays defined.
        let d = max(abs(a - b) + 0.5, min(a + b - 0.01, hypot(dx, dyUp)))
        let base = atan2(dyUp, dx) * 180 / .pi
        let alpha = acos((a * a + d * d - b * b) / (2 * a * d)) * 180 / .pi
        let guessElbow = walk(from: start, angles: guess, lengths: [a, b])[1]
        var best: (distance: Double, angles: [Double])?
        for sign in [1.0, -1.0] {
            let upperAngle = base + sign * alpha
            let elbow = walk(from: start, angles: [upperAngle], lengths: [a])[1]
            let distance = Double(hypot(elbow.x - guessElbow.x, elbow.y - guessElbow.y))
            if best == nil || distance < best!.distance {
                best = (distance, [upperAngle, angle(from: elbow, to: target)])
            }
        }
        return best!.angles
    }

    static func attachments(root: CGPoint, spine: [Double], body: ExerciseBodyProfile) -> FigureAttachments {
        let mid = walk(from: root, angles: [spine[0]], lengths: [body.spine1])[1]
        let neck = walk(from: mid, angles: [spine[1]], lengths: [body.spine2])[1]
        let ox = body.leftOffset.count > 0 ? body.leftOffset[0] : 6
        let oy = body.leftOffset.count > 1 ? body.leftOffset[1] : 2
        return FigureAttachments(
            pelvis: root, mid: mid, neck: neck,
            shoulderL: CGPoint(x: neck.x + ox, y: neck.y + oy),
            hipL: CGPoint(x: root.x + ox, y: root.y + oy)
        )
    }

    /// Resolve a key frame to pure angles: replace each pinned limb's authored angles
    /// with its IK solution (the authored angles only pick the bend direction).
    static func normalize(_ kf: MotionKeyFrame, body: ExerciseBodyProfile) -> FigurePose {
        let root = CGPoint(x: kf.root[0], y: kf.root[1])
        var pins: [String: CGPoint] = [:]
        for (key, xy) in kf.pins ?? [:] where xy.count == 2 {
            pins[key] = CGPoint(x: xy[0], y: xy[1])
        }
        let at = attachments(root: root, spine: kf.spine, body: body)
        var limbs: [FigureLimb: [Double]] = [:]
        for limb in FigureLimb.allCases {
            guard let authored = limb.angles(in: kf) else { continue }
            if let pin = pins[limb.pinKey] {
                let lengths = limb.boneLengths(body)
                limbs[limb] = ik2(start: at.point(for: limb), target: pin,
                                  upper: lengths[0], lower: lengths[1], guess: authored)
            } else {
                limbs[limb] = authored
            }
        }
        return FigurePose(root: root, spine: kf.spine, neck: kf.neck, gaze: kf.gaze,
                          limbs: limbs, pins: pins,
                          hold: kf.hold ?? 0.5, tween: kf.tween ?? 0.6)
    }

    /// Ease-in-out: 3t² − 2t³.
    static func ease(_ t: Double) -> Double {
        3 * t * t - 2 * t * t * t
    }

    /// Shortest-path angular interpolation, so limbs swing in natural arcs.
    static func lerpAngle(_ a: Double, _ b: Double, _ t: Double) -> Double {
        var delta = (b - a + 180).truncatingRemainder(dividingBy: 360)
        if delta < 0 { delta += 360 }
        return a + (delta - 180) * t
    }

    /// Interpolate two normalized poses. A pin survives the tween only if planted in
    /// BOTH neighboring key frames (so planted limbs hold exactly and releasing pins
    /// release naturally); a limb absent from either side is dropped for the tween.
    static func lerp(_ a: FigurePose, _ b: FigurePose, _ t: Double) -> FigurePose {
        var limbs: [FigureLimb: [Double]] = [:]
        for (limb, va) in a.limbs {
            guard let vb = b.limbs[limb] else { continue }
            limbs[limb] = zip(va, vb).map { lerpAngle($0, $1, t) }
        }
        var pins: [String: CGPoint] = [:]
        for (key, pa) in a.pins {
            guard let pb = b.pins[key] else { continue }
            pins[key] = CGPoint(x: pa.x + (pb.x - pa.x) * t, y: pa.y + (pb.y - pa.y) * t)
        }
        return FigurePose(
            root: CGPoint(x: a.root.x + (b.root.x - a.root.x) * t,
                          y: a.root.y + (b.root.y - a.root.y) * t),
            spine: zip(a.spine, b.spine).map { lerpAngle($0, $1, t) },
            neck: lerpAngle(a.neck, b.neck, t),
            gaze: (a.gaze != nil && b.gaze != nil) ? lerpAngle(a.gaze!, b.gaze!, t) : nil,
            limbs: limbs, pins: pins, hold: a.hold, tween: a.tween
        )
    }

    /// Normalized pose → drawable points. Limbs with an active pin are re-solved so
    /// planted hands/feet hold that point exactly through tweens.
    static func geometry(of pose: FigurePose, body: ExerciseBodyProfile, hide: Set<String>) -> FigureGeometry {
        let at = attachments(root: pose.root, spine: pose.spine, body: body)
        let head = walk(from: at.neck, angles: [pose.neck], lengths: [body.neck])[1]
        let r = body.headR
        // Nose tick: 7pt outward from the head rim along the gaze direction
        // (omitted for face-on poses without a gaze).
        var noseStart: CGPoint?
        var noseEnd: CGPoint?
        if let gaze = pose.gaze {
            let d = direction(gaze)
            noseStart = CGPoint(x: head.x + d.dx * r, y: head.y + d.dy * r)
            noseEnd = CGPoint(x: head.x + d.dx * (r + 7), y: head.y + d.dy * (r + 7))
        }

        var limbs: [FigureLimb: [CGPoint]] = [:]
        for (limb, authored) in pose.limbs where !hide.contains(limb.rawValue) {
            let lengths = limb.boneLengths(body)
            var angles = authored
            if let pin = pose.pins[limb.pinKey] {
                angles = ik2(start: at.point(for: limb), target: pin,
                             upper: lengths[0], lower: lengths[1], guess: angles)
            }
            limbs[limb] = walk(from: at.point(for: limb), angles: angles, lengths: lengths)
        }

        return FigureGeometry(
            headCenter: head,
            headRadius: r,
            noseStart: noseStart,
            noseEnd: noseEnd,
            spineStart: at.pelvis,
            spineControl: CGPoint(x: 2 * at.mid.x - (at.pelvis.x + at.neck.x) / 2,
                                  y: 2 * at.mid.y - (at.pelvis.y + at.neck.y) / 2),
            spineEnd: at.neck,
            limbs: limbs
        )
    }
}

/// The full looping animation for one motion: normalized key poses plus
/// continuous-time sampling — hold at each key frame, then an eased angle-space tween
/// to the next; the last frame tweens back to the first.
struct MotionTimeline {
    let poses: [FigurePose]
    let duration: Double

    init?(motion: ExerciseMotion, body: ExerciseBodyProfile) {
        let poses = motion.frames.map { MotionSolver.normalize($0, body: body) }
        let duration = poses.reduce(0) { $0 + $1.hold + $1.tween }
        guard !poses.isEmpty, duration > 0 else { return nil }
        self.poses = poses
        self.duration = duration
    }

    /// The pose at wall-clock `time`, looping every `duration` seconds.
    func pose(at time: Double) -> FigurePose {
        var t = time.truncatingRemainder(dividingBy: duration)
        if t < 0 { t += duration }
        for (i, pose) in poses.enumerated() {
            if t < pose.hold { return pose }
            t -= pose.hold
            if t < pose.tween {
                let next = poses[(i + 1) % poses.count]
                return MotionSolver.lerp(pose, next, MotionSolver.ease(t / pose.tween))
            }
            t -= pose.tween
        }
        return poses[0]
    }
}