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

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

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

"""

# OpenGL drawing madness:
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()


# Math helper functions:

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):
    """ Create a H matrix from rotation and position """
    H = mat(eye(4))
    H[0:3, 0:3] = R
    H[0:3, 3] = p
    return H


class Simulation:
    def simulate(self, endtime, dt):
        """ Evaluate the state of the system over a given time """
        state = self.initial()
        times = np.arange(0.0, endtime, dt)
        y0 = state.T.A1
        def fWrapper(y, t):
           y = mat(y).T
           dy = self.rateOfChange(y, t)
           return dy.T.A1
        res = scipy.integrate.odeint(fWrapper, y0, times)
        states = []
        res = res.T
        r,c =  res.shape
        for ci in range(0,c):
            states.append(mat(res[:,ci]).T)
        self.result = states
        return states
  

class BouncingCube(Simulation):
    """ Simulation of a bouncing cube """
    def __init__(self):
        # Parameters:
        self.gravity = mat( array([0,0,-9.81]) ).T
        self.massI = mat(eye(6)) * 1.01 # Mass matrix

    def initial(self):
        """ Return initial pose """
        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)

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

    def rateOfChange(self, states, thetime):
        """ Calculates rate of change given the current states. """
        quat = states[0:4,  0] # Orientation (4)
        pos  = states[4:7,  0] # Position (3)
        P    = states[7:13, 0] # Momentum (6)
        massI, gravity = self.massI, self.gravity
        # 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 True:
                # Contact forces:
                Fx = 0 #np.exp(-5.0*vx)
                Fy = 0 # np.exp(-5.0*vy)
                Fz = np.exp(-z*10.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 draw(self, state):
        Dicequat = state[0:4, 0]
        Dicepos = state[4:7, 0]

        x, y, z = Dicepos.A1
        R = quatToR(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)


class RollingPendulum(Simulation):
    """ A 2D pendulum with a rolling contact. """
    def __init__(self):
        self.R = 0.01
        self.L = 1
        self.M = 1
        self.g = 9.81

    def initial(self):
        i = mat(zeros((2,1)))
        i[0, 0] = 
        return i

    def rateOfChange(self, states, thetime):
        phi, omega = states[0:2, 0]
        rates = mat(zeros((2,1)))
        return rates

    def draw(self, state):
        pass
        #drawCube(0.6)


class WobbelingWheel(Simulation):
    """ A wheel pushed sideways over a surface """
    def __init__(self):
        self.r1 = 0.5
        self.r2 = 0.05

    def initial(self):
        i = mat(zeros((5,1)))
        return i

    def rateOfChange(self, states, thetime):
        x, y = states[0:2, 0] 
        rates = mat(zeros((5,1)))
        return rates

    def draw(self, state):
        pass
        #drawCube(0.6)


class W(QGLWidget):
    time = 0.0
    index = 0
    def __init__(self, sim):
        super().__init__()
        self.firstRun = True
        self.savePNGS = False
        self.index = 0
        self.sim = sim
        self.resize(500,500)
        self.UP()
        t = QTimer(self)
        t.timeout.connect(self.UP)
        t.start(40)

    def UP(self):
        if self.sim.result:
            if self.index + 1 < len(self.sim.result):
                self.index += 1
            else:
                self.index = 0
        # 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
      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))
      if self.sim.result:
          self.sim.draw(self.sim.result[self.index])


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