import { Grid, useBVH, useGLTF, CameraControls, AccumulativeShadows, OrbitControls, Stats } from '@react-three/drei'
import { Camera, Canvas, useFrame, useThree, useLoader } from '@react-three/fiber'
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 } from 'react'
import POSE_3D_ from "./assets/result_ae_01_ae_08.json"

// F, 133, 3
const POSE_3D: Array<Array<[number, number, number]>> = POSE_3D_ as [number, number, number][][]

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
]

// 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,
]

// 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
]

const DEFAULT_TRANSFORMATION_MATRIX = [
  1, 0, 0, 0,
  0, 1, 0, 0,
  0, 0, 1, 0,
  0, 0, 0, 1,
]
const DEFAULT_NEAR = 0.05
const DEFAULT_FAR = 1
const CAMERA_EXTRINSIC_MATRIX_MAP: Record<string, number[]> = {
  "AE_01": [
    0.2321047, -0.97264263, -0.0096808, -0.96323585, 0.06882254,
    0.02634936, -0.99728089, 0.03661007, 0.97025299, 0.23080732,
    0.07305554, 3.34933242, 0., 0., 0.,
    1.
  ],
  "AE_1A": [
    0.93194049, -0.35571886, -0.07036343, -0.92123075, 0.01084819,
    0.22131041, -0.97514308, 0.24922173, 0.36244895, 0.908012,
    0.21010704, 4.87284891, 0., 0., 0.,
    1.
  ],
  "AE_08": [
    0.66806102, -0.74355508, -0.02864123, -1.10173496, 0.05931037,
    0.09157787, -0.99403007, 0.26760438, 0.74173901, 0.66237402,
    0.10528013, 6.92372493, 0., 0., 0.,
    1.
  ]
} as const
const CAMERA_INTRINSIC_MATRIX_MAP: Record<string, number[]> = {
  "AE_01": [
    1806.82137617, 0., 1230.53175624, 0.,
    1809.75580378, 766.36204406, 0., 0.,
    1.
  ],
  "AE_1A": [
    3467.39715751, 0., 1000.62548655, 0.,
    3473.7168112, 831.64048503, 0., 0.,
    1.
  ],
  "AE_08": [
    2785.43931794, 0., 1254.98272372, 0.,
    2788.10437965, 738.82985324, 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 Scene = () => {
  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 {
    startFrame?: number
    jointRadius?: number
    boneRadius?: number
    showJoints?: boolean
    showBones?: boolean
    frameRate?: number
  }

  const Human3DSkeleton = ({
    startFrame = 0,
    jointRadius = 0.01,
    boneRadius = 0.005,
    showJoints = true,
    showBones = true,
    frameRate = 30
  }: Human3DSkeletonProps) => {
    const [frameIndex, setFrameIndex] = useState(startFrame)
    const totalFrames = POSE_3D.length

    // Use frame to animate through the skeleton poses
    useFrame((state, delta) => {
      // Calculate next frame based on desired frame rate and delta time
      setFrameIndex(prevFrame => {
        // Calculate next frame
        const nextFrame = prevFrame + frameRate * delta
        // Loop back to start if we reach the end
        return nextFrame >= totalFrames ? 0 : nextFrame
      })
    })

    // Get the current frame joints - use Math.floor to get the nearest frame
    const currentFrame = Math.floor(frameIndex) % totalFrames
    const joints = POSE_3D[currentFrame]

    // Function to get appropriate color for a joint index
    const getJointColor = (idx: number) => {
      // Face joints (23-90)
      if (idx >= 23 && idx <= 90) return COLOR_FACE
      // Hand joints (91-132)
      if (idx >= 91 && idx <= 132) return COLOR_FINGERS
      // Foot joints (17-22)
      if (idx >= 17 && idx <= 22) return COLOR_FOOT
      // Head (0-4)
      if (idx <= 4) return COLOR_HEAD
      // Arms (5-10)
      if (idx >= 5 && idx <= 10) return COLOR_ARMS
      // Legs (11-16)
      if (idx >= 11 && idx <= 16) return COLOR_LEGS
      // Default
      return COLOR_SPINE
    }

    // Transform a joint position using the coordinate system conversion
    const transformJointPosition = (j: [number, number, number]) => {
      const [x, y, z] = j
      const V = new THREE.Vector3(x, y, z)
      const worldCvt = Z_UP_TO_Y_UP_PRIME.clone()
      V.applyMatrix4(worldCvt)
      return V
    }

    // Create the joint spheres
    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

    // Create the bone cylinders
    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]

          // Calculate midpoint and length
          const midpoint = new THREE.Vector3().addVectors(startPos, endPos).multiplyScalar(0.5)
          const length = startPos.distanceTo(endPos)

          // Calculate rotation
          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])

          // Calculate midpoint and length
          const midpoint = new THREE.Vector3().addVectors(startPos, endPos).multiplyScalar(0.5)
          const length = startPos.distanceTo(endPos)

          // Calculate rotation
          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}
      </group>
    )
  }

  const scene = (<group>
    {/* <OrbitControls /> */}
    <ambientLight intensity={0.05} />
    <directionalLight castShadow position={[3.3, 6, 4.4]} intensity={5} />
    {/* <Floor /> */}
    {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} />
    })}
    <Axes />
    <Human3DSkeleton jointRadius={0.005} boneRadius={0.0025} frameRate={24} />
  </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