#!/bin/python
#
#
#py-segway

import sys, os, random
import pygame 
from pygame.locals import *
from cgtypes import *
from math import *
from OpenGL.GL import *
from OpenGL.GLU import *
from OpenGL.GLUT import *
import ode

# prepare_GL
def prepare_GL():
    """Prepare drawing.
    """
    # Viewport
    glViewport(0,0,640,480)

    # Initialize
    glClearColor(0.2,0.8,1,0)
    glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT);
    glEnable(GL_DEPTH_TEST)
    glDisable(GL_LIGHTING)
    glEnable(GL_LIGHTING)
    glEnable(GL_NORMALIZE)
    glShadeModel(GL_FLAT)

    # Projection
    glMatrixMode(GL_PROJECTION)
    glLoadIdentity()
    P = mat4(1).perspective(45,1.3333,0.2,20)
    glMultMatrixd(P.toList())

    # Initialize ModelView matrix
    glMatrixMode(GL_MODELVIEW)
    glLoadIdentity()

    # Light source
    glLightfv(GL_LIGHT0,GL_POSITION,[0,0,1,0])
    glLightfv(GL_LIGHT0,GL_DIFFUSE,[1,1,1,1])
    glLightfv(GL_LIGHT0,GL_SPECULAR,[1,1,1,1])
    glEnable(GL_LIGHT0)

    # View transformation
    V = mat4(1).lookAt(1.2*vec3(4,1,6),(0.5,0.5,0), up=(0,1,0))
    V.rotate(pi,vec3(0,1,0))  
    V = V.inverse()
    glMultMatrixd(V.toList())


# draw_body
def draw_body(body):
    """Draw an ODE body.
    """
    
    x,y,z = body.getPosition()
    R = body.getRotation()
    T = mat4()
    T[0,0] = R[0]
    T[0,1] = R[1]
    T[0,2] = R[2]
    T[1,0] = R[3]
    T[1,1] = R[4]
    T[1,2] = R[5]
    T[2,0] = R[6]
    T[2,1] = R[7]
    T[2,2] = R[8]
    T[3] = (x,y,z,1.0)

    glPushMatrix()
    glMultMatrixd(T.toList())
    if body.shape=="box":
        sx,sy,sz = body.boxsize
        glScale(sx, sy, sz)
        glutSolidCube(1)
    if body.shape=="sphere":
        sx,sy,sz = body.boxsize
        glScale(sx, sx, sx)
        glutSolidSphere(1,20,3)
    if body.shape=="floor":
        sx,sy,sz = body.boxsize
        glScale(sx, sx, sx)
        glutSolidCube(1)
    glPopMatrix()


# create_box
def create_box(world, space, density, lx, ly, lz):
    """Create a box body and its corresponding geom."""

    # Create body
    body = ode.Body(world)
    M = ode.Mass()
    M.setBox(density, lx, ly, lz)
    body.setMass(M)

    # Set parameters for drawing the body
    body.shape = "box"
    body.boxsize = (lx, ly, lz)

    # Create a box geom for collision detection
    geom = ode.GeomBox(space, lengths=body.boxsize)
    geom.setBody(body)

    return body


# create_cylinder
def create_cylinder(world, space, density, radius, length):
    """Create a box body and its corresponding geom."""

    # Create body
    body = ode.Body(world)
    M = ode.Mass()
    M.setCylinder(density, 1, radius, length)
    body.setMass(M)

    # Set parameters for drawing the body
    body.shape = "sphere"
    body.boxsize = (radius*2, radius*2,length)

    # Create a box geom for collision detection
    geom = ode.GeomCCylinder(space, radius,length)
    geom.setBody(body)

    return body


def wind():
    """Pull the objects back to the origin.

    Every object will be pulled back to the origin.
    Every couple of frames there'll be a thrust upwards so that
    the objects won't stick to the ground all the time.
    """
    global fps
    windforce=0.001
    
    
    global bodies
    for b in bodies:
        p = vec3([0,0,1])
        b.addForce((-random.random()*windforce+windforce/2)*p.normalize())
        

    
# Collision callback
def near_callback(args, geom1, geom2):
    """Callback function for the collide() method.

    This function checks if the given geoms do collide and
    creates contact joints if they do.
    """

    # Check if the objects do collide
    contacts = ode.collide(geom1, geom2)

    # Create contact joints
    world,contactgroup = args
    for c in contacts:
        c.setBounce(0.2)
        c.setMu(5000)
        j = ode.ContactJoint(world, contactgroup, c)
        j.attach(geom1.getBody(), geom2.getBody())


######################################################################

# Initialize pygame
passed, failed = pygame.init()

# Open a window
srf = pygame.display.set_mode((640,480), OPENGL | DOUBLEBUF|OPENGLBLIT)
pygame.font.init()
print "Font init",pygame.font.get_init()


# Create a world object
world = ode.World()
world.setGravity( (0,-9.81,0) )
world.setERP(0.8)

world.setCFM(1E-5)

# Create a space object
space = ode.Space()

# Create a plane geom which prevent the objects from falling forever
floor = ode.GeomPlane(space, (0,1,0), 0)

# A list with ODE bodies
bodies = []

# A joint group for the contact joints that are generated whenever
# two bodies collide
contactgroup = ode.JointGroup()

# Some variables used inside the simulation loop
fps = 100
dt = 1.0/fps
running = True
state = 0
counter = 0
objcount = 0
clk = pygame.time.Clock()


#First we create the blingway's body.....
body = create_box(world, space, 380, 0.1,.6,0.1)
body.setPosition( (0,0.8,0.00) )
bodies.append(body)
objcount+=1

#Next we create the roller
roller = create_cylinder(world,space,300,.125,.125)
roller.setPosition((0,.125,0.000))
print roller.getRotation()
m = mat4().rotation(1, (1,0,0))
#m.rotate (pi, (0,1,0))
roller.setRotation([0,0,1,1,0,0,0,1,0])
bodies.append(roller)
objcount+=1

# Connect body2 with body1
j1 = ode.HingeJoint(world)
j1.attach(body,roller)
j1.setAnchor( (0,.125,0) )
j1.setAxis( (1,0,0) )
j1.setParam(ode.ParamVel, 0)
j1.setParam(ode.ParamFMax,1)

error_total=0;
pos_desired=0;
speed_desired=0;
lastxpos=0
curxpos=0
lastypos=0
curypos=0
error_current=0
xvelocity=0
xacceleration=0
lastxvelocity=0
yvelocity=0
yacceleration=0
lastyvelocity=0
poserror=0
angle=0
rate=0
lastrate=0
lastangle=0
ac_x=0
ac_y=0
rateavg=0
speedavg=0
rate_wanted=0
v_wanted=0
torque=0
lasttorque=0.5
torque_mux=10
k_rate=0
torqueavg=0
forcedir=0
xposavg=0
paused=0

while running:
    events = pygame.event.get()
    for e in events:
        if e.type==QUIT:
            running=False
        elif e.type==KEYDOWN:
        	if paused==1: paused=0
		else: paused=1

        elif e.type==MOUSEBUTTONDOWN:
        	if paused==1: paused=0
		else: paused=1
	elif e.type==MOUSEMOTION:
		pass
    
    if pygame.key.get_pressed()[K_ESCAPE]:
    
	running=False
    if paused!=1:
    
	# Draw the scene
	prepare_GL()
	for b in bodies:
	    draw_body(b)
    
	font = pygame.font.Font(None, 36)
	text = font.render("Inverted pendulum bot", 0, (0, 0, 0))
	srf.blit(text, (200,200)) 
    
    
	pygame.display.flip()
    
	# Simulate
	n = 8
    
	#Do the sim of the actual motor and sensor...
	lastxpos=curxpos
	curxpos=bodies[1].getPosition()[2]
	lastypos=curypos
	curypos=bodies[1].getPosition()[1]
	
	xposavg+=curxpos
	xposavg=xposavg*fps/(fps+1)
	
	lastxvelocity=xvelocity
	xvelocity=curxpos-lastxpos
	xacceleration=(xvelocity-lastxvelocity)*fps
	
	lastyvelocity=yvelocity
	yvelocity=curxpos-lastypos
	yacceleration=(yvelocity-lastyvelocity)*fps
	
	ac_x=bodies[0].getRotation()[4]*xacceleration
	ac_y=bodies[0].getRotation()[5]*xacceleration
	
	ac_x+=bodies[0].getRotation()[5]*yacceleration
	ac_y+=bodies[0].getRotation()[4]*yacceleration   
	
	ac_x+=bodies[0].getRotation()[5]*-9.81
	ac_y+=bodies[0].getRotation()[4]*-9.81
	
	lastangle=angle
	
	angle=asin(ac_x/sqrt(ac_x**2+ac_y**2))
	
	lastrate=rateavg/5
	
	rate=(angle-lastangle)*fps
	rateavg+=rate
	rateavg-=rateavg/5
    
	last_v=speedavg
	speedavg=speedavg+(xvelocity)
	speedavg=speedavg-speedavg/fps
	
	
    
	#WE OPTIMIZE...


	if abs(torqueavg)<=0.1:
		torqueavg+=(random.random()-0.4)/4
	
	if abs(rateavg)>=0.05:
		torqueavg+=(abs((torqueavg/torque_mux)*(rate_wanted/(rate))))
		print "TA:",torqueavg/torque_mux

	torqueavg=abs(torqueavg*torque_mux/(torque_mux+1))
	
	if torqueavg>1:torqueavg=1
	if torqueavg<-1:torqueavg=-1
	
	
	angle_wanted=0
	speed_wanted=0
	
	if (xposavg/fps>.3):
		speed_wanted= -.002
	elif (xposavg/fps<-.3):
		speed_wanted= .002


	if ((speedavg/fps-speed_wanted)> .001):
		print "want -angle" 
		angle_wanted= -0.05
	elif ((speedavg/fps-speed_wanted)< -.001):
		print "want +angle" 
		angle_wanted= 0.05
		
	if abs(angle-angle_wanted)>pi/45:
		rate_wanted=-20*(angle-angle_wanted)
	else:
		rate_wanted=0
	
	
	if (rate_wanted!=lastrate):
		forcedir=(rate_wanted-lastrate)/abs(rate_wanted-lastrate)
	else:
		forcedir=0
		
	j1.setParam(ode.ParamVel, 60*forcedir)
	j1.setParam(ode.ParamFMax, 5*abs(torqueavg))


	

	print 'T:%03.3f LR:%03.3f RW:%03.3f A:%03.3f X:%03.3f Y:%03.3f V:%03.3f SW:%3.3f XP:%2.2f'%(torqueavg,lastrate,rate_wanted,angle,ac_x,ac_y,speedavg/fps,speed_wanted,xposavg/fps)
    
	for i in range(n):
	    wind() 
	   # Detect collisions and create contact joints
	    space.collide((world,contactgroup), near_callback)
    
	    # Simulation step
	    world.step(dt/n)
    
	    # Remove all contact joints
	    contactgroup.empty()
      
	clk.tick(fps)