from Structures import*
from SearchAlgorithms import*
from numpy import*
import pygame

#-----------
# DISPLAY A STRUCTURE USING PYGAME
#-----------
def show(s):
	screen = pygame.display.set_mode((640, 480))
	clock = pygame.time.Clock()    
	k = 0; dk = 1000; maxk = 10000
	xscale = yscale = 40
	xshift = 20; yshift = 100
	running = 1; count = -1;
	while running:
		# Pygame stuff:
		pygame.time.wait(10)	
		event = pygame.event.poll()
		if event.type == pygame.QUIT:
			running = 0
		screen.fill((255, 255, 255))	
		# Set the load:	
		load = k
		s.F[7] = k
		# Find the displacements
		u = s.analyze()
		# Draw the structure:	
		for i, pt in enumerate(connectivity):			
			x1 = xscale * (geometry[pt[0],0] + u[3*(pt[0]+1)-3]) + xshift
			y1 = yscale*(geometry[pt[0],1] + u[3*(pt[0]+1)-2])+yshift
			x2 = xscale*(geometry[pt[1],0] + u[3*(pt[1]+1)-3])+xshift
			y2 = yscale*(geometry[pt[1],1] + u[3*(pt[1]+1)-2])+yshift
			pygame.draw.line(screen, (0,0,255), (x1, y1), (x2, y2), 3)
			pygame.draw.circle(screen, (255,255,0), (int(x1),int(y1)), 4, 0)
			pygame.draw.circle(screen, (255,255,0), (int(x2),int(y2)), 4, 0)			
		if k >= maxk or k <= -maxk:
			dk *= -1
			count += 1
		k += dk
		pygame.display.flip()
		if(count == -1):
			pygame.time.wait(5000)
			count += 1		
		if count>10:
			running = 0

#-----------
# INPUT DATA FOR STRUCTURE
#-----------
# The geometry of the structural nodes (10-bar truss)
geometry = array([[0., 0. ], [ 5., 0. ], [ 10., 0. ], [ 10., 5. ], [ 5., 5. ], [ 0., 5. ]])
# A connectivity matrix defines how nodes are connected with beams. a row contains ( beam - node - node )
connectivity = array([[ 0, 1 ], [ 1, 2 ], [ 2, 3 ] , [ 3, 4 ], [ 4, 5 ], [ 5, 1 ], [ 0, 4 ], [ 4, 2 ], [ 1, 3 ], [ 4, 1 ] ], dtype = int)
# Restrained degrees of freedom
fixity = array([15,16,0,1])
load = zeros((3*geometry.shape[0])); 
load[7] = -10
s = structure(geometry, connectivity, fixity, load)
show(s)

#-----------
# OPTIMIZATION STEP
#-----------
opt_geom = s.geometry[1:5,1] # to be  optimized - all y coordinates
optimizer = NelderMead(opt_geom.shape[0], s, opt_geom, 1e-5) 
optimizer.Run()
show(s)