from OpenGL.GL import *
from OpenGL.GLUT import *
from OpenGL.GLU import *
import sys
from math import pi, sqrt
from time import perf_counter
import random

from vector import multivector
from polybar import polybar
from polybox import polybox

dimensions = int(sys.argv[1]) if len(sys.argv)>1 else 3
bar_length = 1
tick_length = 20

width = 800
height = 800

factor = sqrt(dimensions)

global bar
last = perf_counter()

force_pos = None
force = None

first_dim = None

mode = sys.argv[2] if len(sys.argv)>2 else 'edge'
shape = sys.argv[3] if len(sys.argv)>3 else 'box'

def random_vector():
	return multivector.vectorize(
		[random.gauss(0,1) for _ in range(dimensions)]
	)

def generate_new_force():
	global force_pos, force, bar
	force_pos = bar.random_point()
	force = random_vector()
	bar.impulse(force, force_pos)

def update_force(dir):
	global first_dim, force_pos, force, bar
	
	if first_dim is None:
		first_dim = dir
	else:
	#	force_pos = multivector.vectorize([1]*dir + [0]*(dimensions-dir))
	#	force_pos = bar.local_to_global(force_pos)
		force = multivector.vectorize([0]*(dir-1) + [1] + [0]*(dimensions-dir))
		force_pos = multivector.vectorize([0]*(first_dim-1) + [1] + [0]*(dimensions-first_dim))
		bar.impulse(force, force_pos)
		first_dim = None

def tick():
	update()
	display()

def update():
	global last
	global bar
	
	new = perf_counter()
	bar.update(new - last)
	last = new

def to_color(a):
	return (a+factor)/(2*factor)

def colors(v):
	if len(v) < 3: # In 2 dimensions, don't bother doing color stuff
		return [1, 1, 1]
	elif len(v) == 3: # In 3 dimensions, use bright-dark to indicate 3d
		a = 1 - to_color(v[2])
		return [a, a, a]
	elif len(v) == 4: # In 4 dimensions, use ??? for 3d and red-blue for 4d
		a = 1# - to_color(v[2])
		b = to_color(v[3])
		c = 1 - b
		return [a*b, a*min(b,c), a*c]
	elif len(v) == 5: # In 5 dimensions, use bright-dark for 3d, red for 4d, green for 5d
		a = 1 - to_color(v[2])
		b = to_color(v[3])
		c = to_color(v[4])
		return [a*b, a*c, a*min(b,c)]
	else: # In higher dimensions, use bright-dark, red, green, blue, and ignore the rest
		a = 1 - to_color(v[2])
		b = to_color(v[3])
		c = to_color(v[4])
		d = to_color(v[5])
		return [a*b, a*c, a*d]
#	a = to_color(v[2]) if v.dim > 2 else 0 # Reddishness
#	b = to_color(v[3]) if v.dim > 3 else 1-a # Greenishness
#	c = to_color(v[4]) if v.dim > 4 else 1-b # Bluishness
#	return [a, b, c]

arrow = [multivector.vectorize(i+[0]*(dimensions-2)) for i in [ # Pad with zeroes to make dimensions line up, then convert to multivectors
	[0, 0],
	[1, 0],
	[0.75, 0.25]
]]

def draw_force(pos, vec):
	if pos is None or vec is None:
		return
	delta = (vec @ arrow[1]) # End @ start gives the rotor between them
	for point in arrow:
		place_vertex(pos + (delta @ point)/2) # Geometric product then applies the rotation (divide by 2 to scale the arrow down)

def draw_rot_axis():
	plane = bar.velocity[2]
	axis = plane.hodge() # Hodge dual
	
	glBegin(GL_LINE_STRIP)
	place_vertex(-axis)
	
	glColor3f(1,0,0)
	glVertex3f(0,0,0)
	
	place_vertex(axis)
	glEnd()

def draw_details():
	glColor3f(0,1,0)
	glRasterPos2f(-2, 1.9)
	for c in str(bar.velocity):
		glutBitmapCharacter(GLUT_BITMAP_9_BY_15, ord(c))
	
	if first_dim is not None:
		glRasterPos2f(-2,-1.9)
		for c in str(first_dim):
			glutBitmapCharacter(GLUT_BITMAP_9_BY_15, ord(c))

def place_vertex(v):
	glColor3fv(colors(v.vector()))
	glVertex3fv(v.projection(3))

def display():
	glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT)
	glMatrixMode(GL_MODELVIEW)
	glLoadIdentity()
	
	v = bar.rotated_points()
	if mode == 'edge':
		glBegin(GL_LINES)
		for (a,b) in bar.get_edges():
			place_vertex(v[a])
			place_vertex(v[b])
		glEnd()
	elif mode == 'face':
		glBegin(GL_QUADS)
		for (a,b,c,d) in bar.get_faces():
			place_vertex(v[a])
			place_vertex(v[b])
			place_vertex(v[c])
			place_vertex(v[d])
		glEnd()
	
	glBegin(GL_LINE_STRIP)
	draw_force(force_pos, force)
	glEnd()
	
	if dimensions == 3: # Axis of rotation only makes sense in 3D
		draw_rot_axis()
	
	draw_details()
	
	glutSwapBuffers()

def setup():
	global bar
	if shape == 'line':
		bar = polybar(dimensions, constrain_position=True)
	elif shape == 'box':
		bar = polybox(dimensions, constrain_position=True)

def timerEvent(_):
	glutPostRedisplay()
	glutTimerFunc(tick_length, timerEvent, 1)

def keypress(key, x, y):
	if key == b' ':
		generate_new_force()
	elif key >= b'0' and key <= b'9':
		dir = int(key)
		update_force(dir)

def glmain(argv):
	glutInitDisplayMode(GLUT_RGB | GLUT_DEPTH | GLUT_DOUBLE)
	glutInitWindowSize(width, height)
	glutInit(argv)
	glutCreateWindow(b'Polybar')
	glutDisplayFunc(tick)
	glutIdleFunc(tick)
	glutKeyboardFunc(keypress)
	glEnable(GL_DEPTH_TEST)
	glMatrixMode(GL_PROJECTION)
	glLoadIdentity()
	glOrtho(-2.0, 2.0, -2.0, 2.0, -2.0, 2.0)
#	glutTimerFunc(tick_length, timerEvent, 1)
	glutMainLoop()

def main():
	setup()
	glmain(sys.argv)

main()