from PyQt4.QtGui import *
from PyQt4.QtCore import *
from PyQt4.QtOpenGL import QGLWidget
from OpenGL.GL import *
from OpenGL.GLU import gluPerspective
import sys
from random import random
from math import pi, cos, sin, fabs, sqrt
from numpy import mat, array, ones, zeros, eye
from numpy.linalg import norm
import numpy as np
import time
import scipy.integrate
import pyopencl

"""
  Test script that lets a dice bounce.
  Converted from 20-sim equations into python code.

  20-sim website:
    http://www.20sim.com

"""
def drawCube(w):
      glBegin(GL_QUADS)			# Start Drawing The Cube
      glColor3f(0.0,1.0,0.0)			# Set The Color To Blue
      glVertex3f( w, w,-w)		# Top Right Of The Quad (Top)
      glVertex3f(-w, w,-w)		# Top Left Of The Quad (Top)
      glVertex3f(-w, w, w)		# Bottom Left Of The Quad (Top)
      glVertex3f( w, w, w)		# Bottom Right Of The Quad (Top)

      glColor3f(1.0,0.5,0.0)			# Set The Color To Orange
      glVertex3f( w,-w, w)		# Top Right Of The Quad (Bottom)
      glVertex3f(-w,-w, w)		# Top Left Of The Quad (Bottom)
      glVertex3f(-w,-w,-w)		# Bottom Left Of The Quad (Bottom)
      glVertex3f( w,-w,-w)		# Bottom Right Of The Quad (Bottom)

      glColor3f(1.0,0.0,0.0)			# Set The Color To Red
      glVertex3f( w, w, w)		# Top Right Of The Quad (Front)
      glVertex3f(-w, w, w)		# Top Left Of The Quad (Front)
      glVertex3f(-w,-w, w)		# Bottom Left Of The Quad (Front)
      glVertex3f( w,-w, w)		# Bottom Right Of The Quad (Front)

      glColor3f(1.0,1.0,0.0)			# Set The Color To Yellow
      glVertex3f( w,-w,-w)		# Bottom Left Of The Quad (Back)
      glVertex3f(-w,-w,-w)		# Bottom Right Of The Quad (Back)
      glVertex3f(-w, w,-w)		# Top Right Of The Quad (Back)
      glVertex3f( w, w,-w)		# Top Left Of The Quad (Back)

      glColor3f(0.0,0.0,1.0)			# Set The Color To Blue
      glVertex3f(-w, w, w)		# Top Right Of The Quad (Left)
      glVertex3f(-w, w,-w)		# Top Left Of The Quad (Left)
      glVertex3f(-w,-w,-w)		# Bottom Left Of The Quad (Left)
      glVertex3f(-w,-w, w)		# Bottom Right Of The Quad (Left)

      glColor3f(1.0,0.0,1.0)			# Set The Color To Violet
      glVertex3f( w, w,-w)		# Top Right Of The Quad (Right)
      glVertex3f( w, w, w)		# Top Left Of The Quad (Right)
      glVertex3f( w,-w, w)		# Bottom Left Of The Quad (Right)
      glVertex3f( w,-w,-w)		# Bottom Right Of The Quad (Right)
      glEnd()				# Done Drawing The Quad

def drawFloor(w, h):
      glBegin(GL_QUADS)			# Start Drawing The Cube

      glColor3f(1.0,0.5,0.0) # Set The Color To Orange
      glVertex3f( w,-w,h)# Top Right Of The Quad (Bottom)
      glVertex3f(-w,-w,h)# Top Left Of The Quad (Bottom)
      glVertex3f(-w,w,h)# Bottom Left Of The Quad (Bottom)
      glVertex3f( w,w,h)# Bottom Right Of The Quad (Bottom)
      glEnd()				# Done Drawing The Quad

def drawAxis():
  glLineWidth(0.5)
  glBegin(GL_LINES)
  glColor3f(1.0, 0.0, 0.0)
  glVertex3f(0,0,0)
  glVertex3f(1,0,0)
  glColor3f(0.0, 1.0, 0.0)
  glVertex3f(0,0,0)
  glVertex3f(0,1,0)
  glColor3f(0.0, 0.0, 1.0)
  glVertex3f(0,0,0)
  glVertex3f(0,0,1)
  glEnd()


def cross(A, B):
   a = A.A1
   b = B.A1
   return mat(np.cross(a, b)).T

def skew(X):
  Y = mat(zeros( (3, 3) ))
  a,b,c = X.A1
  Y[0,1] = -c
  Y[0,2] = b
  Y[1,0] = c
  Y[1,2] = -a
  Y[2,0] = -b
  Y[2,1] = a
  return Y

def adjoint(T):
  W = T[0:3, 0]
  V = T[3:6, 0]
  a = mat(zeros( (6,6) ) )
  a[0:3, 0:3] = skew(W)
  a[3:6, 0:3] = skew(V)
  a[3:6, 3:6] = skew(W)
  return a

def Adjoint(H):
  R = H[0:3, 0:3]
  P = H[0:3, 3]
  a = mat(zeros( (6,6) ) )
  a[0:3, 0:3] = R
  a[3:6, 3:6] = R
  a[3:6, 0:3] = skew(P) * R
  return a

def quatToR(q):
  x, y, z, w = q.A1
  r = mat(eye(3))
  r[0,0] = 1 - (2*y**2+2*z**2)
  r[0,1] = 2*x*y+2*z*w
  r[0,2] = 2*x*z - 2*y*w
  r[1,0] = 2*x*y-2*z*w
  r[1,1] = 1 - (2*x**2 + 2*z**2)
  r[1,2] = 2*y*z + 2*x*w
  r[2,0] = 2*x*z+2*y*w
  r[2,1] = 2*y*z - 2*x*w
  r[2,2] = 1 - (2*x**2+2*y**2)
  return r

def rotateAbout(axis, angle):
   ax, ay, az = (axis/norm(axis)).A1
   qx = ax*sin(angle/2.0)
   qy = ay*sin(angle/2.0)
   qz = az*sin(angle/2.0)
   qw = cos(angle/2.0)
   q = mat(array([qx,qy,qz,qw])).T
   return q

def normalizeQuaternion(quat):
  x,y,z,w = quat.A1
  magnitude = sqrt(x*x + y*y + z*z + w*w)
  x = x / magnitude
  y = y / magnitude
  z = z / magnitude
  w = w / magnitude
  quat[0, 0] = x
  quat[1, 0] = y
  quat[2, 0] = z
  quat[3, 0] = w
  return quat

def VTo4x4(V):
    v1, v2, v3 = V.A1
    return mat(array( \
      [[0.0, -v3,  v2, -v1], \
      [ v3,  0.0, -v1, -v2], \
      [-v2, v1, 0.0, -v3], \
      [v1, v2, v3, 0.0] ]))

def homogeneous(R,p):
    H = mat(eye(4))
    H[0:3, 0:3] = R
    H[0:3, 3] = p
    return H

def rateOfChange(states, thetime, parameters):
    quat = states[0:4,  0] # Orientation (4)
    pos  = states[4:7,  0] # Position (3)
    P    = states[7:13, 0] # Momentum (6)
    massI, gravity = parameters
    # Rigid body parts:
    # Forward Kinematic chain:
    H = homogeneous(quatToR(quat), pos) # Forward kinematics

    AdjX2 = mat(eye(6)) # The connectionpoint in the real world
    adjAdjX2_1 = adjoint(AdjX2[0:6,0])
    adjAdjX2_2 = adjoint(AdjX2[0:6,1])
    adjAdjX2_3 = adjoint(AdjX2[0:6,2])
    adjAdjX2_4 = adjoint(AdjX2[0:6,3])
    adjAdjX2_5 = adjoint(AdjX2[0:6,4])
    adjAdjX2_6 = adjoint(AdjX2[0:6,5])
    AdjInv2 = Adjoint(H.I)
    M2 = AdjInv2.T * (massI * AdjInv2) # Transfor mass to base frame
    MassMatrix = M2

    wrenchGrav2 = mat( zeros((1,6)) )
    wrenchGrav2[0, 0:3] = -cross(gravity, pos).T
    wrenchGrav2[0, 3:6] = gravity.T

    Bk = mat( zeros( (6,6) ))
    Bk[0:3, 0:3] = skew(P[0:3, 0])
    Bk[3:6, 0:3] = skew(P[3:6, 0])
    Bk[0:3, 3:6] = skew(P[3:6, 0])

    # TODO: do this a cholesky:
    v = np.linalg.solve(MassMatrix, P) # Matrix inverse like thingy !

    T2_00 = v # Calculate the relative twist!
    TM2 = T2_00.T * M2
    twistExternal = T2_00 # Twist van het blokje
    TMSum2 = TM2

    PDotBodies = mat( zeros( (6,1)) )
    PDotBodies[0,0] = TMSum2 * (adjAdjX2_1 * T2_00)
    PDotBodies[1,0] = TMSum2 * (adjAdjX2_2 * T2_00)
    PDotBodies[2,0] = TMSum2 * (adjAdjX2_3 * T2_00)
    PDotBodies[3,0] = TMSum2 * (adjAdjX2_4 * T2_00)
    PDotBodies[4,0] = TMSum2 * (adjAdjX2_5 * T2_00)
    PDotBodies[5,0] = TMSum2 * (adjAdjX2_6 * T2_00)
    
    PDot = -PDotBodies - Bk * v
    PDot += wrenchGrav2.T
    
    ##### Contact wrench part:
    HB_W = H # Is H-matrix van het blokje
    WrenchB = mat(zeros( (1,6) ))
    for px in [-0.5, 0.5]:
      for py in [-0.5, 0.5]:
        for pz in [-0.5, 0.5]:
          HB1_B = homogeneous(mat(eye(3)), mat([px,py,pz]).T)
          HB1_W = HB_W * HB1_B
          HW1_W = homogeneous(mat(eye(3)), HB1_W[0:3,3])
          HW_W1 = HW1_W.I
          HB_W1 = HW_W1 * HB_W
    
          AdjHB_W1 = Adjoint(HB_W1)
          TB_W1_W1 = AdjHB_W1 * twistExternal
          z = HB1_W[2, 3]
          vx, vy, vz = TB_W1_W1[3:6, 0].A1
          if z < 0:
            # Contact forces:
            Fx = -50.0*vx
            Fy = -50.0*vy
            Fz = -z*50000.0 
          else:
            Fx = 0.0
            Fy = 0.0
            Fz = 0.0
          # TODO: reflect impulse
          WrenchW1 = mat([0,0,0,0,0,Fz])
          # Transform it back:
          WrenchB += (AdjHB_W1.T * WrenchW1.T).T
    ##### End of contact wrench

    PDot += (WrenchB * AdjInv2).T

    # Position and orientation rates:
    QOmega = VTo4x4(v[0:3, 0])
    quatDot = 0.5 * QOmega * quat
    vel = v[3:6, 0]
    posDot = skew(v[0:3]) * pos + vel
    # The rate vector:
    rates = mat(zeros( (13,1) ))
    rates[0:4, 0] = quatDot
    rates[4:7, 0] = posDot
    rates[7:13, 0] = PDot
    return rates

def fWrapper(y, t, parameters):
   y = mat(y).T
   dy = rateOfChange(y, t, parameters)
   return dy.T.A1
  
def SCIPY(endtime, dt, state, parameters):
    times = np.arange(0.0, endtime, dt)
    y0 = state.T.A1
    res = scipy.integrate.odeint(fWrapper, y0, times, args=(parameters,))
    states = []
    res = res.T
    r,c =  res.shape
    for ci in range(0,c):
      states.append(mat(res[:,ci]).T)
    return states
  
def RK4(endtime, dt, state, parameters):
    t = 0.0
    states = []
    while t < endtime:
      newstate = mat (zeros( state.shape )) # Create a new object

      #### Runge Kutta integration:
      k1 = rateOfChange(state,             t, parameters)
      k2 = rateOfChange(state + 0.5*dt*k1, t, parameters)
      k3 = rateOfChange(state + 0.5*dt*k1, t, parameters)
      k4 = rateOfChange(state + dt*k3,     t, parameters)
      newstate = state + (dt/6.0)*(k1+2*k2+2*k3+k4)

      # Normalize quat:
      newstate[0:4, 0] = normalizeQuaternion(newstate[0:4, 0])
      states.append(newstate)

      state = newstate

      t += dt
      print state[6,0], t, ' ', (t/endtime)*100.0, '%'
    return states
  

def simulate(endtime, dt):
    PInitial = mat( zeros((6,1)) )
    posInitial = mat(array([-1.2, -1.3, 2.8])).T
    quatInitial = rotateAbout(mat(array([0.2, 1.0, 0.4])).T, 0.5)
    # Parameters:
    gravity = mat( array([0,0,-9.81]) ).T
    massI = mat(eye(6)) * 1.01 # Mass matrix
    parameters = (massI, gravity)

    # The state vector!
    state = mat(zeros( (13,1) ))
    state[0:4, 0] = quatInitial
    state[4:7, 0] = posInitial
    state[7:13, 0] = PInitial

    return SCIPY(endtime, dt, state, parameters)

class W(QGLWidget):
  time = 0.0
  index = 0
  def __init__(self, states, dt, parent=None):
    super(W, self).__init__(parent)
    self.firstRun = True
    self.savePNGS = False
    self.dt = dt
    self.resize(500,500)
    self.states = states
    self.UP()
    t = QTimer(self)
    t.timeout.connect(self.UP)
    t.start(self.dt*1000.0)

  def UP(self):
    self.time += self.dt
    if self.states:
      state = self.states[self.index]
      self.Dicequat = state[0:4, 0]
      self.Dicepos = state[4:7, 0]
      self.index += 1
      if self.index == len(self.states):
        self.index = 0
        self.firstRun = False
    # Paint:
    self.update()
    if self.firstRun:
       # Create png images for the movie:
       if self.savePNGS:
          pm = self.renderPixmap()
          pm.save('image'+str(self.index)+'.png')

  def initializeGL(self):
      glClearColor(0.0, 0.5, 0.0, 1.0)
      glEnable(GL_DEPTH_TEST)
      glDepthFunc(GL_LESS)
      glShadeModel(GL_SMOOTH)
  def resizeGL(self,w,h):
      glViewport(0, 0, w, h)
      glMatrixMode(GL_PROJECTION)
      glLoadIdentity()
      gluPerspective(45.0, float(w)/float(h), 0.1, 100.0)
      glMatrixMode(GL_MODELVIEW)
  def paintGL(self):
      glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT) # Clear buffers
      glLoadIdentity() # Reset The View

      glLoadIdentity()
      glTranslatef(0.0,-2.0,-10.0) # Move Left And Into The Screen
      glRotatef(-90.0, 1.0, 0.0, 0.0)
      drawFloor(2.0, 0.0)
      drawAxis()
      self.renderText(1.0, 0.0, 0.0, 'X')
      self.renderText(0.0, 1.0, 0.0, 'Y')
      self.renderText(0.0, 0.0, 1.0, 'Z')

      self.renderText(0.0,0.0,1.2,str(self.time))

      x,y,z = self.Dicepos.A1
      R = quatToR(self.Dicequat)

      glTranslatef(x, y, z)
      # Trick to rotate the openGL matrix:
      r = R.A1
      rotR = (r[0], r[3], r[6], 0.0,
              r[1], r[4], r[7], 0.0,
              r[2], r[5], r[8], 0.0,
              0.0,  0.0,  0.0,  1.0)
      glMultMatrixd(rotR)
      
      drawCube(0.6)

et = 20.0
dt = 0.04
print 'starting integration... endtime =', et, ' stepsize =', dt
t0 = time.time()
states = simulate(et, dt)
t1 = time.time()
print 'That was heavy, it took me ', t1-t0, ' seconds!'
app = QApplication(sys.argv)
w = W(states, dt)
w.show()
sys.exit(app.exec_())