import { Grid, useBVH, useGLTF, CameraControls, AccumulativeShadows, OrbitControls, Stats } from '@react-three/drei'
import { Camera, Canvas, useFrame, useThree, useLoader, RenderCallback, RootState } from '@react-three/fiber'
import { Text as DreiText, TextProps as DreiTextProps } from '@react-three/drei';
import * as THREE from 'three'
import { FontLoader } from 'three/addons/loaders/FontLoader.js'
import { TextGeometry } from 'three/addons/geometries/TextGeometry.js'
import HelvetikerRegular from "three/examples/fonts/helvetiker_regular.typeface.json"
import { useEffect, useRef, useState, JSX, forwardRef } from 'react'
import Boxing from "./assets/Test_WeiHua_Segment_1.json"

// 133, 3
type PosePoints3D = [number, number, number][]

// F, 133, 3
type AnimePosePoints3D = PosePoints3D[]

interface Skeleton {
    "1": PosePoints3D[],
    "2": PosePoints3D[],
    "3": PosePoints3D[],
}

// 扩展 Drei 的 TextProps 接口
interface TextProps extends DreiTextProps {
    position?: [number, number, number];
    fontSize?: number;
    color?: string;
    billboard?: boolean;
}

// 创建类型安全的 Text 组件
const Text = forwardRef<THREE.Mesh, TextProps>((props, ref) => {
    return <DreiText {...props} ref={ref} />;
});

// const POSE_3D = POSE_3D_ as AnimePosePoints3D
// const POSE_3D_MANY = POSE_3D_MANY_ as AnimePosePoints3D[] // N F 133 3
const BOXING = Boxing as Skeleton

const THREE_ADDONS = {
    FontLoader,
    TextGeometry,
} as const

// Create OpenCV to OpenGL conversion matrix
// OpenCV: X right, Y down, Z forward
// OpenGL: X right, Y up, Z backward
const CV_TO_GL_MAT = new THREE.Matrix4().set(
    1, 0, 0, 0,
    0, -1, 0, 0,
    0, 0, -1, 0,
    0, 0, 0, 1
)

// Z-up to Y-up conversion matrix
// Rotate -90 degrees around X axis to convert from Z-up to Y-up
const Z_UP_TO_Y_UP = new THREE.Matrix4().set(
    -1, 0, 0, 0,
    0, 0, -1, 0,
    0, -1, 0, 0,
    0, 0, 0, 1
)

const Z_UP_TO_Y_UP_PRIME = new THREE.Matrix4().set(
    1, 0, 0, 0,
    0, 0, 1, 0,
    0, 1, 0, 0,
    0, 0, 0, 1
)

// Color definitions for different body parts
const COLOR_SPINE = new THREE.Color(138 / 255, 201 / 255, 38 / 255)  // green, spine & head
const COLOR_ARMS = new THREE.Color(255 / 255, 202 / 255, 58 / 255)   // yellow, arms & shoulders
const COLOR_LEGS = new THREE.Color(25 / 255, 130 / 255, 196 / 255)   // blue, legs & hips
const COLOR_FINGERS = new THREE.Color(255 / 255, 0, 0)           // red, fingers
const COLOR_FACE = new THREE.Color(255 / 255, 200 / 255, 0)        // yellow, face
const COLOR_FOOT = new THREE.Color(255 / 255, 128 / 255, 0)        // orange, foot
const COLOR_HEAD = new THREE.Color(255 / 255, 0, 255 / 255)        // purple, head

// Body bone connections
const BODY_BONES = [
    // legs
    [15, 13], [13, 11], [16, 14], [14, 12], [11, 12], // legs
    [5, 11], [6, 12], [5, 6],                         // torso
    [5, 7], [7, 9], [6, 8], [8, 10],                  // arms
    [1, 2], [0, 1], [0, 2], [1, 3], [2, 4],           // head
    [15, 17], [15, 18], [15, 19],                     // left foot
    [16, 20], [16, 21], [16, 22],                     // right foot
] as const

// Body bone colors
const BODY_BONE_COLORS = [
    COLOR_LEGS, COLOR_LEGS, COLOR_LEGS, COLOR_LEGS, COLOR_LEGS,
    COLOR_SPINE, COLOR_SPINE, COLOR_SPINE,
    COLOR_ARMS, COLOR_ARMS, COLOR_ARMS, COLOR_ARMS,
    COLOR_HEAD, COLOR_HEAD, COLOR_HEAD, COLOR_HEAD, COLOR_HEAD,
    COLOR_FOOT, COLOR_FOOT, COLOR_FOOT,
    COLOR_FOOT, COLOR_FOOT, COLOR_FOOT,
] as const

// Hand bone connections (in pairs of [start, end] indices)
const HAND_BONES = [
    // right hand
    [91, 92], [92, 93], [93, 94], [94, 95],           // right thumb
    [91, 96], [96, 97], [97, 98], [98, 99],           // right index
    [91, 100], [100, 101], [101, 102], [102, 103],    // right middle
    [91, 104], [104, 105], [105, 106], [106, 107],    // right ring
    [91, 108], [108, 109], [109, 110], [110, 111],    // right pinky
    // left hand
    [112, 113], [113, 114], [114, 115], [115, 116],   // left thumb
    [112, 117], [117, 118], [118, 119], [119, 120],   // left index
    [112, 121], [121, 122], [122, 123], [123, 124],   // left middle
    [112, 125], [125, 126], [126, 127], [127, 128],   // left ring
    [112, 129], [129, 130], [130, 131], [131, 132]    // left pinky
] as const

const DEFAULT_TRANSFORMATION_MATRIX = [
    1, 0, 0, 0,
    0, 1, 0, 0,
    0, 0, 1, 0,
    0, 0, 0, 1,
] as const
const DEFAULT_NEAR = 0.05
const DEFAULT_FAR = 1
const CAMERA_EXTRINSIC_MATRIX_MAP: Record<string, number[]> = {
    "5602": [
        -0.76138617, 0.16195241, 0.62774399, -0.73832425, -0.40758532,
        -0.87257285, -0.26924121, 0.55561231, 0.5041481, -0.46085577,
        0.73037432, 3.28663985, 0., 0., 0.,
        1.
    ] as const,
    "5603": [
        -0.44966325, -0.02806161, -0.89275725, -0.24260155, 0.21789368,
        -0.97275592, -0.07917234, 0.41197614, -0.8662132, -0.23012705,
        0.44352704, 4.37769458, 0., 0., 0.,
        1.
    ] as const,
    "5604": [
        0.9551736, -0.02735669, 0.29477958, 0.5908996, -0.06148677,
        -0.99234061, 0.10714238, 0.87629594, 0.28959069, -0.12046462,
        -0.94953963, 4.1617598, 0., 0., 0.,
        1.
    ] as const,
    "5605": [
        0.57140284, 0.03118154, -0.82007713, -0.22679438, 0.11021058,
        -0.99314166, 0.0390292, 0.83904835, -0.81323577, -0.11268257,
        -0.5709205, 4.17730229, 0., 0., 0.,
        1.
    ] as const,
}
const CAMERA_INTRINSIC_MATRIX_MAP: Record<string, number[]> = {
    "5602": [
        2686.00393399, 0., 1470.49106388, 0.,
        2699.92688641, 765.51266883, 0., 0.,
        1.
    ] as const,
    "5603": [
        2791.38378418, 0., 1258.11161208, 0.,
        2790.67070205, 788.14860021, 0., 0.,
        1.
    ] as const,
    "5604": [
        2789.25680621, 0., 1231.13101573, 0.,
        2787.08458106, 677.69385417, 0., 0.,
        1.
    ] as const,
    "5605": [
        2644.68145454, 0., 1285.34889127, 0.,
        2644.9700955, 627.20808584, 0., 0.,
        1.
    ] as const,
}

const IMAGE_WIDTH = 2560
const IMAGE_HEIGHT = 1440

const intrinsicToFov = (intrinsic: number[], image_size: { width: number, height: number }) => {
    console.assert(intrinsic.length === 9, "intrinsic must be a 3x3 matrix")
    const fx = intrinsic[0]
    const fy = intrinsic[4]
    const cx = intrinsic[2]
    const cy = intrinsic[5]
    // in degrees
    const fov_x = 2 * Math.atan(image_size.width / (2 * fx)) * (180 / Math.PI)
    const fov_y = 2 * Math.atan(image_size.height / (2 * fy)) * (180 / Math.PI)
    return { fov_x, fov_y }
}

const calculaterCubeVersices = (position: number[], dimensions: number[]) => {
    const [cx, cy, cz] = position
    const [width, height, depth] = dimensions
    const halfWidth = width / 2
    const halfHeight = height / 2
    const halfDepth = depth / 2

    return [
        [cx - halfWidth, cy - halfHeight, cz - halfDepth],
        [cx + halfWidth, cy - halfHeight, cz - halfDepth],
        [cx + halfWidth, cy + halfHeight, cz - halfDepth],
        [cx - halfWidth, cy + halfHeight, cz - halfDepth],
        [cx - halfWidth, cy - halfHeight, cz + halfDepth],
        [cx + halfWidth, cy - halfHeight, cz + halfDepth],
        [cx + halfWidth, cy + halfHeight, cz + halfDepth],
        [cx - halfWidth, cy + halfHeight, cz + halfDepth]
    ]
}


const Scene = () => {

    // 定义立方体绘制组件
    const Cube = () => {

        const vertices = calculaterCubeVersices([0, -0.205 + 0.9, 0.90], [0.5, 1.8, 0.5])

        return (
            <>
                <mesh position={[0, -0.205 + 0.9, 0.90]}> {/** 原点位置，相对于六面体中心偏移 */}
                    <boxGeometry args={[0.5, 1.8, 0.5]} /> {/** 边长，宽：0.5米，高：1.8米，深度：0.5米 */}
                    <meshStandardMaterial
                        color="#007BFF" // 蓝色
                        opacity={0.2} // 半透明
                        transparent={true}
                        wireframe={true} // 线框模式，仅显示边框
                    />
                </mesh>

                {/* 顶点标记（带坐标系和文本） */}
                {vertices.map(([x, y, z], index) => (
                    <group key={index} position={[x, y, z]}>
                        {/* 顶点处的小标记点 */}
                        <mesh>
                            <sphereGeometry args={[0.0, 16, 16]} />
                            <meshBasicMaterial color="#ff0000" />
                        </mesh>

                        {/* 坐标轴（X:红, Y:绿, Z:蓝） */}
                        <axesHelper args={[0.2]} />

                        {/* 沙袋上的坐标文本（始终面向相机） */}
                        {/* <Text
                            position={[0, 0.1, 0]}  // 文本在顶点上方
                            fontSize={0.08}
                            color="#007Bff"
                            billboard
                        >
                            {`P${index}: (${x.toFixed(2)}, ${y.toFixed(2)}, ${z.toFixed(2)})`}
                        </Text> */}
                    </group>
                ))}

            </>
        )
    }

    function Floor() {
        return (
            <mesh rotation-x={-Math.PI / 2} position-y={-0.5} receiveShadow>
                <planeGeometry args={[15, 15]} />
                <meshStandardMaterial color="#ccc" />
            </mesh>
        )
    }
    const Axes = () => {
        return <axesHelper args={[15]} />
    }
    interface CameraViewFromExtrinsicProps {
        extrinsic: number[] | THREE.Matrix4
        aspect?: number
        name?: string
        near?: number
        far?: number
        fov?: number
        textSize?: number
    }
    // https://threejs.org/docs/#examples/en/loaders/FontLoader
    // https://www.ilyameerovich.com/simple-3d-text-meshes-in-three-js/

    const CameraViewFromExtrinsic = ({ extrinsic, name, near, far, fov, textSize, aspect }: CameraViewFromExtrinsicProps) => {
        let Rt: THREE.Matrix4
        if (extrinsic instanceof THREE.Matrix4) {
            Rt = extrinsic
        } else if (Array.isArray(extrinsic)) {
            console.assert(extrinsic.length === 16, "extrinsic must be a 4x4 matrix")
            Rt = new THREE.Matrix4()
            // @ts-expect-error 16 elements
            Rt.set(...extrinsic)
        } else {
            throw new Error("extrinsic must be a 4x4 matrix or an array of 16 elements")
        }
        const font = new FontLoader().parse(HelvetikerRegular)
        const camera = new THREE.PerspectiveCamera(fov ?? 60, aspect ?? 4 / 3, near ?? DEFAULT_NEAR, far ?? DEFAULT_FAR)
        const helper = <cameraHelper args={[camera]} />
        camera.applyMatrix4(Rt)

        const textRef = useRef<THREE.Mesh>(null)
        const { camera: viewCamera } = useThree()

        useFrame(() => {
            if (textRef.current) {
                textRef.current.lookAt(viewCamera.position)
            }
        })

        let text: JSX.Element | null = null
        if (name) {
            const geo = new THREE_ADDONS.TextGeometry(name ?? "", { font, size: textSize ?? 0.1, depth: 0.001 })
            const position = new THREE.Vector3()
            position.setFromMatrixPosition(Rt)

            text = (
                <mesh ref={textRef} position={position}>
                    <primitive object={geo} />
                    <meshStandardMaterial color="black" />
                </mesh>
            )
        }
        return (
            <group>
                {text}
                <primitive object={camera} />
                {helper}
            </group>
        )
    }

    const preProcessExtrinsic = (extrinsic: number[] | THREE.Matrix4) => {
        let Rt: THREE.Matrix4
        if (extrinsic instanceof THREE.Matrix4) {
            Rt = extrinsic
        } else if (Array.isArray(extrinsic)) {
            console.assert(extrinsic.length === 16, "extrinsic must be a 4x4 matrix")
            Rt = new THREE.Matrix4()
            // @ts-expect-error 16 elements
            Rt.set(...extrinsic)
        } else {
            throw new Error("extrinsic must be a 4x4 matrix or an array of 16 elements")
        }


        // Then handle OpenCV to OpenGL camera convention
        const cameraCvt = CV_TO_GL_MAT.clone()

        // Convert from Z-up to Y-up first (this affects world coordinates)
        // const worldCvt = Z_UP_TO_Y_UP.clone()

        // Final transformation:
        // 1. Convert world from Z-up to Y-up
        // 2. Apply the camera transform
        // 3. Convert camera coordinates from OpenCV to OpenGL
        const final = new THREE.Matrix4()
        final
            .multiply(cameraCvt)
            .multiply(Rt)
        // .multiply(worldCvt)

        // Invert to get the camera-to-world transform
        final.invert()
        return final
    }

    interface Human3DSkeletonProps {
        skeleton: AnimePosePoints3D
        startFrame?: number
        jointRadius?: number
        boneRadius?: number
        showJoints?: boolean
        showBones?: boolean
        frameRate?: number
    }

    /**
     * 3D人体骨架动画组件
     * @param skeleton 动画帧序列，每帧为133个关节点的三维坐标
     * @param startFrame 起始帧
     * @param jointRadius 关节点球体半径
     * @param boneRadius 骨骼圆柱体半径
     * @param showJoints 是否显示关节点
     * @param showBones 是否显示骨骼
     * @param frameRate 动画帧率
     */
    const Human3DSkeleton = ({
        skeleton,
        startFrame = 0,
        jointRadius = 0.01,
        boneRadius = 0.005,
        showJoints = true,
        showBones = true,
        frameRate = 24
    }: Human3DSkeletonProps) => {
        // 当前动画帧索引
        const [frameIndex, setFrameIndex] = useState(startFrame)
        const totalFrames = skeleton.length
        // 动画帧推进函数，按frameRate自动推进
        const onFrame: RenderCallback = (totalFrames === 0) ? (state, delta) => { } : (state: RootState, delta: number) => {
            // 根据帧率和delta推进帧索引，实现动画
            setFrameIndex(prevFrame => {
                const nextFrame = prevFrame + frameRate * delta
                // 到末尾自动循环
                return nextFrame >= totalFrames ? 0 : nextFrame
            })
            return null
        }

        // 注册动画帧推进
        useFrame(onFrame)

        // 当前帧的133个关节点坐标
        const currentFrame = Math.floor(frameIndex) % totalFrames
        const joints = skeleton[currentFrame]

        /**
         * 右腕（关键点10）速度计算
         * 速度 = (当前帧位置 - 上一帧位置) / (1 / frameRate)
         * 单位：与坐标单位一致/秒
         */
        let wristSpeed: number | null = null;
        if (currentFrame > 0 && joints && joints.length > 16) {
            const prevJoints = skeleton[(currentFrame - 1 + totalFrames) % totalFrames];
            if (prevJoints && prevJoints.length > 16) {
                const wNow = joints[10];
                const wPrev = prevJoints[10];
                if (wNow && wPrev) {
                    const dx = wNow[0] - wPrev[0];
                    const dy = wNow[1] - wPrev[1];
                    const dz = wNow[2] - wPrev[2];
                    const dist = Math.sqrt(dx * dx + dy * dy + dz * dz);
                    wristSpeed = dist * frameRate; // 单位/秒
                }
            }
        }

        /**
         * 右手大臂（右肩-右肘）和小臂（右肘-右腕）夹角计算
         * COCO/BlazePose: 右肩=6, 右肘=8, 右腕=10
         * 夹角为两向量的夹角，单位度
         */
        let rightUpperArmAngle: number | null = null;
        if (joints && joints.length > 16) {
            const shoulder = joints[6]; // 右肩
            const elbow = joints[8];    // 右肘
            const wrist = joints[10];   // 右腕
            if (shoulder && elbow && wrist) {
                // 大臂向量（肩->肘）
                const v1 = new THREE.Vector3(
                    shoulder[0] - elbow[0],
                    shoulder[1] - elbow[1],
                    shoulder[2] - elbow[2]
                );
                // 小臂向量（腕->肘）
                const v2 = new THREE.Vector3(
                    wrist[0] - elbow[0],
                    wrist[1] - elbow[1],
                    wrist[2] - elbow[2]
                );
                // 计算夹角
                const dot = v1.dot(v2);
                const len1 = v1.length();
                const len2 = v2.length();
                if (len1 > 1e-6 && len2 > 1e-6) {
                    const angle = Math.acos(Math.max(-1, Math.min(1, dot / (len1 * len2))));
                    rightUpperArmAngle = angle * 180 / Math.PI;
                }
            }
        }

        /**
         * 获取关节点颜色
         * @param idx 关节点索引
         */
        const getJointColor = (idx: number) => {
            // 脸部
            if (idx >= 23 && idx <= 90) return COLOR_FACE
            // 手指
            if (idx >= 91 && idx <= 132) return COLOR_FINGERS
            // 脚
            if (idx >= 17 && idx <= 22) return COLOR_FOOT
            // 头部
            if (idx <= 4) return COLOR_HEAD
            // 手臂
            if (idx >= 5 && idx <= 10) return COLOR_ARMS
            // 腿
            if (idx >= 11 && idx <= 16) return COLOR_LEGS
            // 躯干
            return COLOR_SPINE
        }

        /**
         * 关节点坐标转换（如需坐标系变换可在此处理）
         */
        const transformJointPosition = (j: [number, number, number]) => {
            const [x, y, z] = j
            const V = new THREE.Vector3(x, y, z)
            return V
        }

        // 生成关节点球体
        const jointMeshes = showJoints ? joints.map((j, idx) => {
            const position = transformJointPosition(j)
            const color = getJointColor(idx)
            return (
                <mesh key={`joint-${idx}`} position={position}>
                    <sphereGeometry args={[jointRadius, 16, 16]} />
                    <meshStandardMaterial color={color} />
                </mesh>
            )
        }) : null

        // 生成骨骼圆柱体
        const boneMeshes = showBones ? (
            <>
                {/* 身体主骨骼 */}
                {BODY_BONES.map((bone, idx) => {
                    const [startIdx, endIdx] = bone
                    if (startIdx >= joints.length || endIdx >= joints.length) return null
                    // 起止点
                    const startPos = transformJointPosition(joints[startIdx])
                    const endPos = transformJointPosition(joints[endIdx])
                    const color = BODY_BONE_COLORS[idx]
                    // 骨骼中点和长度
                    const midpoint = new THREE.Vector3().addVectors(startPos, endPos).multiplyScalar(0.5)
                    const length = startPos.distanceTo(endPos)
                    // 旋转对齐
                    const direction = new THREE.Vector3().subVectors(endPos, startPos).normalize()
                    const quaternion = new THREE.Quaternion()
                    const up = new THREE.Vector3(0, 1, 0)
                    quaternion.setFromUnitVectors(up, direction)
                    return (
                        <mesh key={`bone-body-${idx}`} position={midpoint} quaternion={quaternion}>
                            <cylinderGeometry args={[boneRadius, boneRadius, length, 8]} />
                            <meshStandardMaterial color={color} />
                        </mesh>
                    )
                })}
                {/* 手部骨骼 */}
                {HAND_BONES.map((bone, idx) => {
                    const [startIdx, endIdx] = bone
                    if (startIdx >= joints.length || endIdx >= joints.length) return null
                    const startPos = transformJointPosition(joints[startIdx])
                    const endPos = transformJointPosition(joints[endIdx])
                    const midpoint = new THREE.Vector3().addVectors(startPos, endPos).multiplyScalar(0.5)
                    const length = startPos.distanceTo(endPos)
                    const direction = new THREE.Vector3().subVectors(endPos, startPos).normalize()
                    const quaternion = new THREE.Quaternion()
                    const up = new THREE.Vector3(0, 1, 0)
                    quaternion.setFromUnitVectors(up, direction)
                    return (
                        <mesh key={`bone-hand-${idx}`} position={midpoint} quaternion={quaternion}>
                            <cylinderGeometry args={[boneRadius, boneRadius, length, 8]} />
                            <meshStandardMaterial color={COLOR_FINGERS} />
                        </mesh>
                    )
                })}
            </>
        ) : null

        // 渲染骨架、关节点、骨骼、右臂夹角和右腕速度文本
        return (
            <group>
                {jointMeshes}
                {boneMeshes}
                {/* 实时显示右手大臂和小臂夹角，文本位于右肘上方 */}
                {rightUpperArmAngle !== null && joints[8] && (
                    <Text
                        position={[1.5, 2.5, 0]}
                        fontSize={0.3}
                        color="#ff00ff"
                        billboard
                    >
                        {`右臂夹角: ${rightUpperArmAngle.toFixed(1)}°`}
                    </Text>
                )}
                {/* 实时显示右腕速度，文本位于右腕上方 */}
                {wristSpeed !== null && joints[10] && (
                    <Text
                        position={[-1.5, 2.5, 0]}
                        fontSize={0.3}
                        color="#00bfff"
                        billboard
                    >
                        {`右腕速度: ${wristSpeed.toFixed(3)}`}
                    </Text>
                )}
            </group>
        )
    }

    // 调小关节点和骨骼的半径
    const skeletons = [
        <Human3DSkeleton jointRadius={0.03} boneRadius={0.014} frameRate={24} skeleton={BOXING[1]} />,
    ]

    const cameras = Object.entries(CAMERA_EXTRINSIC_MATRIX_MAP).map(([key, value]) => {
        const intrinsic = CAMERA_INTRINSIC_MATRIX_MAP[key]
        const { fov_x, fov_y } = intrinsicToFov(intrinsic, { width: IMAGE_WIDTH, height: IMAGE_HEIGHT })
        // make the far reverse proportional to the fov
        const far = (1 / fov_x) * 20
        return <CameraViewFromExtrinsic key={key} name={`${key}(${fov_x.toFixed(1)})`} extrinsic={preProcessExtrinsic(value)} fov={fov_x} aspect={IMAGE_WIDTH / IMAGE_HEIGHT} far={far} />
    })

    // 在场景中添加立方体
    const scene = (<group>
        {/* <OrbitControls /> */}
        <ambientLight intensity={0.05} />
        <directionalLight castShadow position={[3.3, 6, 4.4]} intensity={5} />
        {/* <Floor /> */}
        { }
        <Axes />
        <Cube /> {/** 新增立方体 */}
        {cameras}
        {skeletons}
    </group>)
    return (
        // Note that we don't need to import anything, All three.js objects will be treated
        // as native JSX elements, just like you can just write <div /> or <span /> in
        // regular ReactDOM. The general rule is that Fiber components are available under
        // the camel-case version of their name in three.js.
        <>
            <CameraControls />
            <Stats />
            {scene}
        </>
    )
}

function App() {
    return (
        <Canvas shadows style={{ background: "#e9e9e9", width: "100vw", height: "100vh" }}>
            <Scene />
        </Canvas>
    )
}

export default App