package us.ihmc.avatar.reachabilityMap.voxelPrimitiveShapes;

import us.ihmc.commons.MathTools;
import us.ihmc.euclid.referenceFrame.FrameQuaternion;
import us.ihmc.euclid.referenceFrame.ReferenceFrame;
import us.ihmc.euclid.referenceFrame.interfaces.FramePose3DBasics;
import us.ihmc.euclid.tuple3D.Point3D;
import us.ihmc.euclid.tuple3D.Vector3D;
import us.ihmc.euclid.tuple4D.Quaternion;
import us.ihmc.robotics.geometry.SpiralBasedAlgorithm;

/**
 * SphereVoxelShape creates N points uniformly distributed on the surface of a sphere. For each
 * point, a ray, which goes from the point on the surface of the sphere to its origin, is generated.
 * For each ray M rotations are generated by computing the orientation aligning the x-axis to the
 * ray and transforming this orientation by M rotations around the ray. This class is meant to help
 * discretizing the 3D space of orientations and also to simulate the different possibilities for
 * grasping a spherical object.
 */
public class SphereVoxelShape
{
   public enum SphereVoxelType
   {
      graspOrigin, graspAroundSphere
   };

   private final Quaternion[][] rotations;
   private final Point3D[] pointsOnSphere;
   /** Origin of the sphere in the current voxel coordinate. Should probably always be set to zero. */
   private final Point3D sphereOrigin = new Point3D();

   private final int numberOfRays;
   private final int numberOfRotationsAroundRay;

   private final SphereVoxelType type;
   private ReferenceFrame referenceFrame = ReferenceFrame.getWorldFrame();

   public SphereVoxelShape(double voxelSize, int numberOfRays, int numberOfRotationsAroundRay, SphereVoxelType type)
   {
      this.type = type;
      this.numberOfRays = numberOfRays;
      this.numberOfRotationsAroundRay = numberOfRotationsAroundRay;

      pointsOnSphere = SpiralBasedAlgorithm.generatePointsOnSphere(sphereOrigin, voxelSize, numberOfRays);
      rotations = SpiralBasedAlgorithm.generateOrientations(numberOfRays, numberOfRotationsAroundRay);
   }

   public void setReferenceFrame(ReferenceFrame referenceFrame)
   {
      this.referenceFrame = referenceFrame;
   }

   public int getNumberOfRays()
   {
      return numberOfRays;
   }

   public int getNumberOfRotationsAroundRay()
   {
      return numberOfRotationsAroundRay;
   }

   public void getRay(Vector3D rayToPack, int rayIndex)
   {
      MathTools.checkIntervalContains(rayIndex, 0, numberOfRays - 1);

      rayToPack.sub(sphereOrigin, pointsOnSphere[rayIndex]);
      rayToPack.normalize();
   }

   public void getOrientation(FrameQuaternion orientation, int rayIndex, int rotationAroundRayIndex)
   {
      MathTools.checkIntervalContains(rayIndex, 0, numberOfRays - 1);
      MathTools.checkIntervalContains(rotationAroundRayIndex, 0, numberOfRotationsAroundRay - 1);

      orientation.setIncludingFrame(referenceFrame, rotations[rayIndex][rotationAroundRayIndex]);
   }

   public void getPose(FramePose3DBasics pose, int rayIndex, int rotationAroundRayIndex)
   {
      MathTools.checkIntervalContains(rayIndex, 0, numberOfRays - 1);
      MathTools.checkIntervalContains(rotationAroundRayIndex, 0, numberOfRotationsAroundRay - 1);

      if (type == SphereVoxelType.graspAroundSphere)
      {
         pose.setIncludingFrame(referenceFrame, pointsOnSphere[rayIndex], rotations[rayIndex][rotationAroundRayIndex]);
      }
      else
      {
         pose.setToZero(referenceFrame);
         pose.getOrientation().set(rotations[rayIndex][rotationAroundRayIndex]);
      }
   }

   public Point3D[] getPointsOnSphere()
   {
      return pointsOnSphere;
   }
}
