/*
 * Copyright (c) 2017 Yoshihisa Nitta
 * Released under the MIT license
 * http://opensource.org/licenses/mit-license.php
 */

/*
 * NtUnity.cs version 0.1beta 2017/08/19
 *  http://nw.tsuda.ac.jp/NtKinect/
 *
 * requires:
 *   NtKinectDLL version 1.0 or later
 */


using System.Collections;
using System.Collections.Generic;
using UnityEngine;
using System;
using System.Runtime.InteropServices;

namespace NtUnity {
  public class Kinect {
    public const int
      // Number
      bodyCount = 6,
      jointCount = 25,
      rgbCols = 1920,
      rgbRows = 1080,
      depthCols = 512,
      depthRows = 424,
      // JointType
      JointType_SpineBase= 0,
      JointType_SpineMid= 1,
      JointType_Neck= 2,
      JointType_Head= 3,
      JointType_ShoulderLeft= 4,
      JointType_ElbowLeft= 5,
      JointType_WristLeft= 6,
      JointType_HandLeft= 7,
      JointType_ShoulderRight= 8,
      JointType_ElbowRight= 9,
      JointType_WristRight= 10,
      JointType_HandRight= 11,
      JointType_HipLeft= 12,
      JointType_KneeLeft= 13,
      JointType_AnkleLeft= 14,
      JointType_FootLeft= 15,
      JointType_HipRight= 16,
      JointType_KneeRight= 17,
      JointType_AnkleRight= 18,
      JointType_FootRight= 19,
      JointType_SpineShoulder= 20,
      JointType_HandTipLeft= 21,
      JointType_ThumbLeft= 22,
      JointType_HandTipRight= 23,
      JointType_ThumbRight= 24,
      // TrackingState
      TrackingState_NotTracked= 0,
      TrackingState_Inferred= 1,
      TrackingState_Tracked= 2,
      // FacePoint
      FacePointType_None= -1,
      FacePointType_EyeLeft= 0,
      FacePointType_EyeRight= 1,
      FacePointType_Nose= 2,
      FacePointType_MouthCornerLeft= 3,
      FacePointType_MouthCornerRight= 4,
      FacePointType_Count= ( FacePointType_MouthCornerRight + 1 ) ,
      // a_FaceProperty
      FaceProperty_Happy= 0,
      FaceProperty_Engaged= 1,
      FaceProperty_WearingGlasses= 2,
      FaceProperty_LeftEyeClosed= 3,
      FaceProperty_RightEyeClosed= 4,
      FaceProperty_MouthOpen= 5,
      FaceProperty_MouthMoved= 6,
      FaceProperty_LookingAway= 7,
      FaceProperty_Count= ( FaceProperty_LookingAway + 1 ) ,
      // FaceDetectionResult
      DetectionResult_Unknown= 0,
      DetectionResult_No= 1,
      DetectionResult_Maybe= 2,
      DetectionResult_Yes= 3,
      // HDFace
      HDFaceVerticesSize = 1347,
      // dummy
      NtKinectdummy = 0;

    [DllImport ("NtKinectDLL")] private static extern IntPtr getKinect();
    [DllImport ("NtKinectDLL")] private static extern void stopKinect(IntPtr ptr);

    // OpenCV
    [DllImport ("NtKinectDLL")] private static extern void imshow(IntPtr ptr);

    // Multi Thread
    [DllImport ("NtKinectDLL")] private static extern void acquire(IntPtr ptr);
    [DllImport ("NtKinectDLL")] private static extern void release(IntPtr ptr);

    // Audio
    [DllImport ("NtKinectDLL")] private static extern void setAudio(IntPtr ptr, bool flag);
    [DllImport ("NtKinectDLL")] private static extern float getBeamAngle(IntPtr ptr);
    [DllImport ("NtKinectDLL")] private static extern float getBeamAngleConfidence(IntPtr ptr);
    [DllImport ("NtKinectDLL")] private static extern void openAudio(IntPtr ptr, IntPtr filename);
    [DllImport ("NtKinectDLL")] private static extern void closeAudio(IntPtr ptr);
    [DllImport ("NtKinectDLL")] private static extern bool isOpenedAudio(IntPtr ptr);

    // RGB
    [DllImport ("NtKinectDLL")] private static extern void setRGB(IntPtr ptr);
    [DllImport ("NtKinectDLL")] private static extern int getRGB(IntPtr ptr, IntPtr data);

    // Depth
    [DllImport ("NtKinectDLL")] private static extern void setDepth(IntPtr ptr);
    [DllImport ("NtKinectDLL")] private static extern int getDepth(IntPtr ptr, IntPtr data);

    // Infrared
    [DllImport ("NtKinectDLL")] private static extern void setInfrared(IntPtr ptr);
    [DllImport ("NtKinectDLL")] private static extern int getInfrared(IntPtr ptr, IntPtr data);

    // BodyIndex
    [DllImport ("NtKinectDLL")] private static extern void setBodyIndex(IntPtr ptr);
    [DllImport ("NtKinectDLL")] private static extern int getBodyIndex(IntPtr ptr, IntPtr data);
    
    // Skeleton
    [DllImport ("NtKinectDLL")] private static extern void setSkeleton(IntPtr ptr);
    [DllImport ("NtKinectDLL")] private static extern int getSkeleton(IntPtr ptr, IntPtr skelton, IntPtr state, IntPtr id, IntPtr tid);
    [DllImport ("NtKinectDLL")] private static extern int handState(IntPtr ptr,int id,bool isLeft);
    
    // Face
    [DllImport ("NtKinectDLL")] private static extern void setFace(IntPtr ptr, bool isColorSpace);
    [DllImport ("NtKinectDLL")] private static extern int getFace(IntPtr ptr, IntPtr point,IntPtr rect,IntPtr direction,IntPtr property,IntPtr tid);

    // HDFace
    [DllImport ("NtKinectDLL")] private static extern void setHDFace(IntPtr ptr);
    [DllImport ("NtKinectDLL")] private static extern int getHDFace(IntPtr ptr, IntPtr point, IntPtr tid, IntPtr status);
    
    // Gesture
    [DllImport ("NtKinectDLL")] private static extern void setGestureFile(IntPtr ptr, IntPtr filename);
    [DllImport ("NtKinectDLL")] private static extern int setGestureId(IntPtr ptr, IntPtr name, int id); // id: non-zero
    [DllImport ("NtKinectDLL")] private static extern void setGesture(IntPtr ptr);
    [DllImport ("NtKinectDLL")] private static extern int getDiscreteGesture(IntPtr ptr, IntPtr gid, IntPtr confidence);
    [DllImport ("NtKinectDLL")] private static extern int getContinuousGesture(IntPtr ptr, IntPtr gid, IntPtr progress);
    [DllImport ("NtKinectDLL")] private static extern int getGidMapSize();

    private IntPtr kinect;
    public Vector3[] joint = new Vector3[jointCount];
    public int[] jointState = new int[jointCount];

    public Kinect() {
      kinect = getKinect();
    }
    public void stopKinect() { stopKinect(kinect); }

    // OpenCV
    public void imshow() { imshow(kinect); }
    // Multi Thread
    public void acquire() { acquire(kinect); }
    public void release() { release(kinect); }
    // Audio
    public void setAudio(bool flag) { setAudio(kinect,flag); }
    public float getBeamAngle() { return getBeamAngle(kinect); }
    public float getBeamAngleConfidence() { return getBeamAngleConfidence(kinect); }
    public void openAudio(string filename) {
      System.IntPtr fname = Marshal.StringToHGlobalUni(filename);
      openAudio(kinect,fname);
      Marshal.FreeHGlobal(fname);
    }
    public void closeAudio() { closeAudio(kinect); }
    public bool isOpenedAudio() { return isOpenedAudio(kinect); }
    // RGB
    public void setRGB() { setRGB(kinect); }
    public int getRGB(out byte[] data) {
      data = new byte[rgbRows * rgbCols * 3];
      GCHandle gch = GCHandle.Alloc(data,GCHandleType.Pinned);
      int n = getRGB(kinect,gch.AddrOfPinnedObject());
      gch.Free();
      return n;
    }

    // Depth
    public void setDepth() { setDepth(kinect); }
    public int getDepth(out ushort[] data) {
      data = new ushort[depthRows * depthCols];
      GCHandle gch = GCHandle.Alloc(data,GCHandleType.Pinned);
      int n = getDepth(kinect,gch.AddrOfPinnedObject());
      gch.Free();
      return n;
    }

    // Infrared
    public void setInfrared() { setInfrared(kinect); }
    public int getInfrared(out ushort[] data) {
      data = new ushort[depthRows * depthCols];
      GCHandle gch = GCHandle.Alloc(data,GCHandleType.Pinned);
      int n = getInfrared(kinect,gch.AddrOfPinnedObject());
      gch.Free();
      return n;
    }

    // BodyIndex
    public void setBodyIndex() {setBodyIndex(kinect); }
    public int getBodyIndex(byte[] data) {
      data = new byte[depthRows * depthCols];
      GCHandle gch = GCHandle.Alloc(data,GCHandleType.Pinned);
      int n = getBodyIndex(kinect,gch.AddrOfPinnedObject());
      gch.Free();
      return n;
    }

    // Skeleton
    public void setSkeleton() { setSkeleton(kinect); }
    public int getSkeleton(out float[] skeleton,  out int[] state, out int[] id, out ulong[] tid) {
      skeleton = new float[bodyCount * jointCount * 3];
      state = new int[bodyCount * jointCount];
      id = new int[bodyCount];
      tid = new ulong[bodyCount];
      GCHandle gch = GCHandle.Alloc(skeleton,GCHandleType.Pinned);
      GCHandle gch2 = GCHandle.Alloc(state,GCHandleType.Pinned);
      GCHandle gch3 = GCHandle.Alloc(id,GCHandleType.Pinned);
      GCHandle gch4 = GCHandle.Alloc(tid,GCHandleType.Pinned);
      int n = getSkeleton(kinect,gch.AddrOfPinnedObject(),gch2.AddrOfPinnedObject(),gch3.AddrOfPinnedObject(),gch4.AddrOfPinnedObject());
      gch.Free();
      gch2.Free();
      gch3.Free();
      gch4.Free();
      return n;
    }

    // Face
    public void setFace() { setFace(kinect,true); }
    public int getFace(out float[] point,out float[] rect,out float[] direction,out int[] property,out ulong[] tid) {
      point = new float[bodyCount * FacePointType_Count * 3];
      rect = new float[bodyCount * 4];
      direction = new float[bodyCount * 3];
      property = new int[bodyCount * FaceProperty_Count * 2];
      tid = new ulong[bodyCount];
      GCHandle gch = GCHandle.Alloc(point,GCHandleType.Pinned);
      GCHandle gch2 = GCHandle.Alloc(rect,GCHandleType.Pinned);
      GCHandle gch3 = GCHandle.Alloc(direction,GCHandleType.Pinned);
      GCHandle gch4 = GCHandle.Alloc(property,GCHandleType.Pinned);
      GCHandle gch5 = GCHandle.Alloc(tid,GCHandleType.Pinned);
      int n = getFace(kinect,gch.AddrOfPinnedObject(),gch2.AddrOfPinnedObject(),gch3.AddrOfPinnedObject(),gch4.AddrOfPinnedObject(),gch5.AddrOfPinnedObject());
      gch.Free();
      gch2.Free();
      gch3.Free();
      gch4.Free();
      gch5.Free();
      return n;
    }

    // HDFace
    public void setHDFace() { setHDFace(kinect); }
    public int getHDFace(float[] point,ulong[] tid,int[] status) {
      point = new float[bodyCount * HDFaceVerticesSize * 3];
      tid = new ulong[bodyCount];
      status = new int[bodyCount * 2];
      GCHandle gch = GCHandle.Alloc(point,GCHandleType.Pinned);
      GCHandle gch2 = GCHandle.Alloc(tid,GCHandleType.Pinned);
      GCHandle gch3 = GCHandle.Alloc(status,GCHandleType.Pinned);
      int n = getHDFace(kinect,gch.AddrOfPinnedObject(),gch2.AddrOfPinnedObject(),gch3.AddrOfPinnedObject());
      gch.Free();
      gch2.Free();
      gch3.Free();
      return n;
    }

    // Gesture
    public void setGestureFile(string filename) {
      IntPtr gbd = Marshal.StringToHGlobalUni(filename);
      setGestureFile(kinect,gbd);
      Marshal.FreeHGlobal(gbd);
    }
    public int setGestureId(string name, int id) {
      System.IntPtr g = Marshal.StringToHGlobalUni(name); // discrete
      int n = setGestureId(kinect,g,id);
      Marshal.FreeHGlobal(g);
      return n;
    }
    public void setGesture() { setGesture(kinect); }
    public int getDiscreteGesture(out int[] gid,out float[] confidence) {
      gid = new int[bodyCount * getGidMapSize()];
      confidence = new float[bodyCount * getGidMapSize()];
      GCHandle gch = GCHandle.Alloc(gid,GCHandleType.Pinned);
      GCHandle gch2 = GCHandle.Alloc(confidence,GCHandleType.Pinned);
      int n = getDiscreteGesture(kinect,gch.AddrOfPinnedObject(),gch2.AddrOfPinnedObject());
      gch.Free();
      gch2.Free();
      return n;
    }
    public int getContinuousGesture(out int[] gid, out float[] progress){
      gid = new int[bodyCount * getGidMapSize()];
      progress = new float[bodyCount * getGidMapSize()];
      GCHandle gch = GCHandle.Alloc(gid,GCHandleType.Pinned);
      GCHandle gch2 = GCHandle.Alloc(progress,GCHandleType.Pinned);
      int n = getContinuousGesture(kinect,gch.AddrOfPinnedObject(),gch2.AddrOfPinnedObject());
      gch.Free();
      gch2.Free();
      return n;
    }
  }
  public class RigBone {
    public GameObject gameObject;
    public HumanBodyBones bone;
    public bool isValid;
    public Transform transform {
      get { return animator.GetBoneTransform(bone); }
    }
    Animator animator;
    Quaternion savedValue;
    public RigBone(GameObject g, HumanBodyBones b) {
      gameObject = g;
      bone = b;
      isValid = false;
      animator = gameObject.GetComponent<Animator>();
      if (animator == null) {
	Debug.Log("no Animator Component");
	return;
      }
      Avatar avatar = animator.avatar;
      if (avatar == null || !avatar.isHuman || !avatar.isValid) {
	Debug.Log("Avatar is not Humanoid or it is not valid");
	return;
      }
      isValid = true;
      savedValue = animator.GetBoneTransform(bone).localRotation;
    }
    public void set(float a, float x, float y, float z) {
      set(Quaternion.AngleAxis(a, new Vector3(x,y,z)));
    }
    public void set(Quaternion q) {
      animator.GetBoneTransform(bone).localRotation = q;
      savedValue = q;
    }
    public void mul(float a, float x, float y, float z) {
      mul(Quaternion.AngleAxis(a, new Vector3(x,y,z)));
    }
    public void mul(Quaternion q) {
      Transform tr = animator.GetBoneTransform(bone);
      tr.localRotation = q * tr.localRotation;
    }
    public void offset(float a, float x, float y, float z) {
      offset(Quaternion.AngleAxis(a, new Vector3(x,y,z)));
    }
    public void offset(Quaternion q) {
      animator.GetBoneTransform(bone).localRotation = q * savedValue;
    }
    public void changeBone(HumanBodyBones b) {
      bone = b;
      savedValue = animator.GetBoneTransform(bone).localRotation;
    }
  }

  class HumanoidSkeleton {
    private static int[] jointSegment = new int[] {
      Kinect.JointType_SpineBase, Kinect.JointType_SpineMid,             // Spine
      Kinect.JointType_Neck, Kinect.JointType_Head,                      // Neck
      // left
      Kinect.JointType_ShoulderLeft, Kinect.JointType_ElbowLeft,         // LeftUpperArm
      Kinect.JointType_ElbowLeft, Kinect.JointType_WristLeft,            // LeftLowerArm
      Kinect.JointType_WristLeft, Kinect.JointType_HandLeft,             // LeftHand
      Kinect.JointType_HipLeft, Kinect.JointType_KneeLeft,               // LeftUpperLeg
      Kinect.JointType_KneeLeft, Kinect.JointType_AnkleLeft,             // LeftLowerLeg6
      Kinect.JointType_AnkleLeft, Kinect.JointType_FootLeft,             // LeftFoot
      // right
      Kinect.JointType_ShoulderRight, Kinect.JointType_ElbowRight,       // RightUpperArm
      Kinect.JointType_ElbowRight, Kinect.JointType_WristRight,          // RightLowerArm
      Kinect.JointType_WristRight, Kinect.JointType_HandRight,           // RightHand
      Kinect.JointType_HipRight, Kinect.JointType_KneeRight,             // RightUpperLeg
      Kinect.JointType_KneeRight, Kinect.JointType_AnkleRight,           // RightLowerLeg
      Kinect.JointType_AnkleRight, Kinect.JointType_FootRight,           // RightFoot
    };
    public Vector3[] joint = new Vector3[Kinect.jointCount];
    public int[] jointState = new int[Kinect.jointCount];

    Dictionary<HumanBodyBones,Vector3> trackingSegment = null;
    Dictionary<HumanBodyBones, int>  trackingState = null;

    private static HumanBodyBones[] humanBone = new HumanBodyBones[] {
      HumanBodyBones.Hips,
      HumanBodyBones.Spine,
      HumanBodyBones.UpperChest,
      HumanBodyBones.Neck,
      HumanBodyBones.Head,
      HumanBodyBones.LeftUpperArm,
      HumanBodyBones.LeftLowerArm,
      HumanBodyBones.LeftHand,
      HumanBodyBones.LeftUpperLeg,
      HumanBodyBones.LeftLowerLeg,
      HumanBodyBones.LeftFoot,
      HumanBodyBones.RightUpperArm,
      HumanBodyBones.RightLowerArm,
      HumanBodyBones.RightHand,
      HumanBodyBones.RightUpperLeg,
      HumanBodyBones.RightLowerLeg,
      HumanBodyBones.RightFoot,
    };

    private static HumanBodyBones[] targetBone = new HumanBodyBones[] {
      HumanBodyBones.Spine,
      HumanBodyBones.Neck,
      HumanBodyBones.LeftUpperArm,
      HumanBodyBones.LeftLowerArm,
      HumanBodyBones.LeftHand,
      HumanBodyBones.LeftUpperLeg,
      HumanBodyBones.LeftLowerLeg,
      HumanBodyBones.LeftFoot,
      HumanBodyBones.RightUpperArm,
      HumanBodyBones.RightLowerArm,
      HumanBodyBones.RightHand,
      HumanBodyBones.RightUpperLeg,
      HumanBodyBones.RightLowerLeg,
      HumanBodyBones.RightFoot,
    };

    public GameObject humanoid;
    private Dictionary<HumanBodyBones, RigBone> rigBone = null;
    private bool isSavedPosition = false;
    private Vector3 savedPosition;
    private Quaternion savedHumanoidRotation;

    public HumanoidSkeleton(GameObject h) {
      humanoid = h;
      rigBone = new Dictionary<HumanBodyBones, RigBone>();
      foreach (HumanBodyBones bone in humanBone) {
	rigBone[bone] = new RigBone(humanoid,bone);
      }
      savedHumanoidRotation = humanoid.transform.rotation;
      trackingSegment = new Dictionary<HumanBodyBones,Vector3>(targetBone.Length);
      trackingState = new Dictionary<HumanBodyBones, int>(targetBone.Length);
    }
    private void swapJoint(int a, int b) {
      Vector3 tmp = joint[a]; joint[a] = joint[b]; joint[b] = tmp;
      int t = jointState[a]; jointState[a] = jointState[b]; jointState[b] = t;
    }
    public void set(float[] jt, int[] st, int offset, bool mirrored, bool move) {
      if (isSavedPosition == false && jointState[Kinect.JointType_SpineBase] != Kinect.TrackingState_NotTracked) {
	isSavedPosition = true;
	int j = offset * Kinect.jointCount + Kinect.JointType_SpineBase;
	savedPosition = new Vector3(jt[j*3],jt[j*3+1],jt[j*3+2]);
      }
      for (int i=0; i<Kinect.jointCount; i++) {
	int j = offset * Kinect.jointCount + i;
	if (mirrored) {
	  joint[i] = new Vector3(-jt[j*3], jt[j*3+1], -jt[j*3+2]);
	} else {
	  joint[i] = new Vector3(jt[j*3], jt[j*3+1], savedPosition.z*2-jt[j*3+2]);
	}
	jointState[i] = st[j];
      }
      if (mirrored) {
	swapJoint(Kinect.JointType_ShoulderLeft, Kinect.JointType_ShoulderRight);
	swapJoint(Kinect.JointType_ElbowLeft, Kinect.JointType_ElbowRight);
	swapJoint(Kinect.JointType_WristLeft, Kinect.JointType_WristRight);
	swapJoint(Kinect.JointType_HandLeft, Kinect.JointType_HandRight);
	swapJoint(Kinect.JointType_HipLeft, Kinect.JointType_HipRight);
	swapJoint(Kinect.JointType_KneeLeft, Kinect.JointType_KneeRight);
	swapJoint(Kinect.JointType_AnkleLeft, Kinect.JointType_AnkleRight);
	swapJoint(Kinect.JointType_FootLeft, Kinect.JointType_FootRight);
	swapJoint(Kinect.JointType_HandTipLeft, Kinect.JointType_HandTipRight);
	swapJoint(Kinect.JointType_ThumbLeft, Kinect.JointType_ThumbRight);
      }
      for (int i=0; i<targetBone.Length; i++) {
	int s = jointSegment[2*i], e = jointSegment[2*i+1];
	trackingSegment[targetBone[i]] = joint[e] - joint[s];
	trackingState[targetBone[i]] = System.Math.Min(jointState[e],jointState[s]);
      }

      Vector3 waist = joint[Kinect.JointType_HipRight] - joint[Kinect.JointType_HipLeft];
      waist = new Vector3(waist.x, 0, waist.z);
      Quaternion rot = Quaternion.FromToRotation(Vector3.right,waist);
      Quaternion rotInv = Quaternion.Inverse(rot);
 
      Vector3 shoulder = joint[Kinect.JointType_ShoulderRight] - joint[Kinect.JointType_ShoulderLeft];
      shoulder = new Vector3(shoulder.x, 0, shoulder.z);
      Quaternion srot = Quaternion.FromToRotation(Vector3.right,shoulder);
      Quaternion srotInv = Quaternion.Inverse(srot);

      humanoid.transform.rotation = Quaternion.identity;
      foreach (HumanBodyBones bone in targetBone) {
	rigBone[bone].transform.rotation = rotInv * Quaternion.FromToRotation(Vector3.up,trackingSegment[bone]);
      }
      rigBone[HumanBodyBones.UpperChest].offset(srot);
      Quaternion bodyRot = rot;
      if (mirrored) {
	bodyRot = Quaternion.AngleAxis(180,Vector3.up) * bodyRot;
      }
      humanoid.transform.rotation = bodyRot;
      if (move == true) {
	Vector3 m = joint[Kinect.JointType_SpineBase];
	if (mirrored) m = new Vector3(-m.x, m.y, -m.z);
	humanoid.transform.position = m;
      }
    }
  }
}