package com.performancecoachlab.posedetection.skeleton

import kotlin.math.PI
import kotlin.math.abs
import kotlin.math.atan2

data class Skeleton(
    val timestamp: Long = 0L,
    val leftShoulder: SkeletonCoordinate? = null,
    val rightShoulder: SkeletonCoordinate? = null,
    val leftElbow: SkeletonCoordinate? = null,
    val rightElbow: SkeletonCoordinate? = null,
    val leftWrist: SkeletonCoordinate? = null,
    val rightWrist: SkeletonCoordinate? = null,
    val leftHip: SkeletonCoordinate? = null,
    val rightHip: SkeletonCoordinate? = null,
    val leftKnee: SkeletonCoordinate? = null,
    val rightKnee: SkeletonCoordinate? = null,
    val leftAnkle: SkeletonCoordinate? = null,
    val rightAnkle: SkeletonCoordinate? = null,
    val width: Float,
    val height: Float,
) {
    data class SkeletonCoordinate(
        val x: Float,
        val y: Float
    )

    data class SkeletonAngles(
        val leftShoulder: Double?,
        val rightShoulder: Double?,
        val leftElbow: Double?,
        val rightElbow: Double?,
        val leftHip: Double?,
        val rightHip: Double?,
        val leftKnee: Double?,
        val rightKnee: Double?,
    )

    companion object {
        /** Linearly interpolate between two skeletons. [alpha] in 0..1. */
        fun lerp(a: Skeleton, b: Skeleton, alpha: Float): Skeleton {
            fun lerpCoord(
                c1: SkeletonCoordinate?, c2: SkeletonCoordinate?, t: Float
            ): SkeletonCoordinate? {
                if (c1 == null) return c2
                if (c2 == null) return c1
                return SkeletonCoordinate(
                    x = c1.x + (c2.x - c1.x) * t,
                    y = c1.y + (c2.y - c1.y) * t
                )
            }
            val t = alpha.coerceIn(0f, 1f)
            return Skeleton(
                timestamp = (a.timestamp + ((b.timestamp - a.timestamp) * t).toLong()),
                leftShoulder = lerpCoord(a.leftShoulder, b.leftShoulder, t),
                rightShoulder = lerpCoord(a.rightShoulder, b.rightShoulder, t),
                leftElbow = lerpCoord(a.leftElbow, b.leftElbow, t),
                rightElbow = lerpCoord(a.rightElbow, b.rightElbow, t),
                leftWrist = lerpCoord(a.leftWrist, b.leftWrist, t),
                rightWrist = lerpCoord(a.rightWrist, b.rightWrist, t),
                leftHip = lerpCoord(a.leftHip, b.leftHip, t),
                rightHip = lerpCoord(a.rightHip, b.rightHip, t),
                leftKnee = lerpCoord(a.leftKnee, b.leftKnee, t),
                rightKnee = lerpCoord(a.rightKnee, b.rightKnee, t),
                leftAnkle = lerpCoord(a.leftAnkle, b.leftAnkle, t),
                rightAnkle = lerpCoord(a.rightAnkle, b.rightAnkle, t),
                width = a.width,
                height = a.height,
            )
        }
    }

    fun bones(): List<Pair<SkeletonCoordinate, SkeletonCoordinate>> {
        val lines = emptyList<Pair<SkeletonCoordinate, SkeletonCoordinate>>().toMutableList()
        if (leftShoulder != null && rightShoulder != null) lines += Pair(
            leftShoulder,
            rightShoulder
        )
        if (leftShoulder != null && leftElbow != null) lines += Pair(leftShoulder, leftElbow)
        if (rightShoulder != null && rightElbow != null) lines += Pair(rightShoulder, rightElbow)
        if (leftElbow != null && leftWrist != null) lines += Pair(leftElbow, leftWrist)
        if (rightElbow != null && rightWrist != null) lines += Pair(rightElbow, rightWrist)
        if (leftShoulder != null && leftHip != null) lines += Pair(leftShoulder, leftHip)
        if (rightShoulder != null && rightHip != null) lines += Pair(rightShoulder, rightHip)
        if (leftHip != null && rightHip != null) lines += Pair(leftHip, rightHip)
        if (leftHip != null && leftKnee != null) lines += Pair(leftHip, leftKnee)
        if (rightHip != null && rightKnee != null) lines += Pair(rightHip, rightKnee)
        if (leftKnee != null && leftAnkle != null) lines += Pair(leftKnee, leftAnkle)
        if (rightKnee != null && rightAnkle != null) lines += Pair(rightKnee, rightAnkle)
        return lines.toList()
    }

    fun joints(): List<SkeletonCoordinate> {
        val joints = emptyList<SkeletonCoordinate>().toMutableList()
        if (leftShoulder != null) joints += leftShoulder
        if (rightShoulder != null) joints += rightShoulder
        if (leftElbow != null) joints += leftElbow
        if (rightElbow != null) joints += rightElbow
        if (leftWrist != null) joints += leftWrist
        if (rightWrist != null) joints += rightWrist
        if (leftHip != null) joints += leftHip
        if (rightHip != null) joints += rightHip
        if (leftKnee != null) joints += leftKnee
        if (rightKnee != null) joints += rightKnee
        if (leftAnkle != null) joints += leftAnkle
        if (rightAnkle != null) joints += rightAnkle
        return joints.toList()
    }

    fun getAngle(
        firstJoint: SkeletonCoordinate,
        middleJoint: SkeletonCoordinate,
        lastJoint: SkeletonCoordinate
    ): Double {
        var result = toDegrees(
            atan2(lastJoint.y - middleJoint.y, lastJoint.x - middleJoint.x)
                    - atan2(firstJoint.y - middleJoint.y, firstJoint.x - middleJoint.x)
        )
        result = abs(result)
        if (result > 180) {
            result = 360.0 - result
        }
        return result
    }

    private fun toDegrees(radians: Float): Double {
        return radians * 180.0 / PI
    }

    fun getAngles(): SkeletonAngles {
        return SkeletonAngles(
            leftShoulder = leftShoulder?.let { leftShoulder ->
                leftElbow?.let { leftElbow ->
                    leftHip?.let { leftHip ->
                        getAngle(leftElbow, leftShoulder, leftHip)
                    }
                }
            },
            rightShoulder = rightShoulder?.let { rightShoulder ->
                rightElbow?.let { rightElbow ->
                    rightHip?.let { rightHip ->
                        getAngle(rightElbow, rightShoulder, rightHip)
                    }
                }
            },
            leftElbow = leftShoulder?.let { leftShoulder ->
                leftElbow?.let { leftElbow ->
                    leftWrist?.let { leftWrist ->
                        getAngle(leftShoulder, leftElbow, leftWrist)
                    }
                }
            },
            rightElbow = rightShoulder?.let { rightShoulder ->
                rightElbow?.let { rightElbow ->
                    rightWrist?.let { rightWrist ->
                        getAngle(rightShoulder, rightElbow, rightWrist)
                    }
                }
            },
            leftHip = leftShoulder?.let { leftShoulder ->
                leftHip?.let { leftHip ->
                    leftKnee?.let { leftKnee ->
                        getAngle(leftShoulder, leftHip, leftKnee)
                    }
                }
            },
            rightHip = rightShoulder?.let { rightShoulder ->
                rightHip?.let { rightHip ->
                    rightKnee?.let { rightKnee ->
                        getAngle(rightShoulder, rightHip, rightKnee)
                    }
                }
            },
            leftKnee = leftHip?.let { leftHip ->
                leftKnee?.let { leftKnee ->
                    leftAnkle?.let { leftAnkle ->
                        getAngle(leftHip, leftKnee, leftAnkle)
                    }
                }
            },
            rightKnee = rightHip?.let { rightHip ->
                rightKnee?.let { rightKnee ->
                    rightAnkle?.let { rightAnkle ->
                        getAngle(rightHip, rightKnee, rightAnkle)
                    }
                }
            },
        )
    }

    fun mirror(w: Float = width): Skeleton {
        return Skeleton(
            timestamp = timestamp,
            leftShoulder = leftShoulder?.let { SkeletonCoordinate(w - it.x, it.y) },
            rightShoulder = rightShoulder?.let { SkeletonCoordinate(w - it.x, it.y) },
            leftElbow = leftElbow?.let { SkeletonCoordinate(w - it.x, it.y) },
            rightElbow = rightElbow?.let { SkeletonCoordinate(w - it.x, it.y) },
            leftWrist = leftWrist?.let { SkeletonCoordinate(w - it.x, it.y) },
            rightWrist = rightWrist?.let { SkeletonCoordinate(w - it.x, it.y) },
            leftHip = leftHip?.let { SkeletonCoordinate(w - it.x, it.y) },
            rightHip = rightHip?.let { SkeletonCoordinate(w - it.x, it.y) },
            leftKnee = leftKnee?.let { SkeletonCoordinate(w - it.x, it.y) },
            rightKnee = rightKnee?.let { SkeletonCoordinate(w - it.x, it.y) },
            leftAnkle = leftAnkle?.let { SkeletonCoordinate(w - it.x, it.y) },
            rightAnkle = rightAnkle?.let { SkeletonCoordinate(w - it.x, it.y) },
            width = width,
            height = height,
        )
    }

    fun rotate(angle: Int): Skeleton {
        if (angle % 90 != 0) return this

        val normalizedAngle = ((angle % 360) + 360) % 360
        if (normalizedAngle == 0) return this

        val frameWidth = this.width
        val frameHeight = this.height

        val transform: (SkeletonCoordinate) -> SkeletonCoordinate = when (normalizedAngle) {
            90 -> { coord -> SkeletonCoordinate(coord.y, frameWidth - coord.x) }
            180 -> { coord -> SkeletonCoordinate(frameWidth - coord.x, frameHeight - coord.y) }
            270 -> { coord -> SkeletonCoordinate(frameHeight - coord.y, coord.x) }
            else -> { coord -> coord }
        }

        val newDimensions = when (normalizedAngle) {
            90, 270 -> Pair(frameHeight, frameWidth)
            else -> Pair(frameWidth, frameHeight)
        }

        return Skeleton(
            timestamp = timestamp,
            leftShoulder = leftShoulder?.let(transform),
            rightShoulder = rightShoulder?.let(transform),
            leftElbow = leftElbow?.let(transform),
            rightElbow = rightElbow?.let(transform),
            leftWrist = leftWrist?.let(transform),
            rightWrist = rightWrist?.let(transform),
            leftHip = leftHip?.let(transform),
            rightHip = rightHip?.let(transform),
            leftKnee = leftKnee?.let(transform),
            rightKnee = rightKnee?.let(transform),
            leftAnkle = leftAnkle?.let(transform),
            rightAnkle = rightAnkle?.let(transform),
            width = newDimensions.first,
            height = newDimensions.second
        )
    }
}