import pygame
from math import sin, cos, pi
from numpy import array, dot
from random import randint

def random_color():
    return (randint(0, 255), randint(0, 255), randint(0, 255))

win_size = (800, 600)
screen = pygame.display.set_mode(win_size)
clock = pygame.time.Clock()

class Shape:
    def __init__(self, shape, size, colors, pos = (0, 0)):
        self.shape = shape
        self.colors = colors
        self.pos = pos
        self.rotated = [0, 0, 0]
        half = 'int' in str(type(size)) and [size / 2.0] * 3 or [size[0] / 2.0, size[1] / 2.0, size[2] / 2.0]
        vertices = {
            'cube': [
                (-half[0], -half[1], -half[2], 0),
                (-half[0], half[1], -half[2], 0),
                (half[0], half[1], -half[2], 0),
                (half[0], -half[1], -half[2], 0),
                (-half[0], -half[1], half[2], 0),
                (-half[0], half[1], half[2], 0),
                (half[0], half[1], half[2], 0),
                (half[0], -half[1], half[2], 0)],
            'pyramid': [
                (0, -half[1], half[2], 0),
                (-half[0], half[1], half[2], 0),
                (half[0], half[1], half[2], 0),
                (0, 0, -half[2], 0)],
            'diamond': [
                (0, 0, half[2] * 3, 0),
                (-half[0], -half[1], half[2], 0),
                (-half[0], half[1], half[2], 0),
                (half[0], half[1], half[2], 0),
                (half[0], -half[1], half[2], 0),
                (0, 0, -half[2], 0)]}
        self.vertices = vertices[shape]

    def rotate(self, axis, angle):
        radians = angle * pi / 180
        rotation = [
            array([1, 0, 0, 0, 0, cos(radians), -sin(radians), 0, 0, sin(radians), cos(radians), 0, 0, 0, 0, 1]).reshape(4, 4),
            array([cos(radians), 0, sin(radians), 0, 0, 1, 0, 0, -sin(radians), 0, cos(radians), 0, 0, 0, 0, 1]).reshape(4, 4),
            array([cos(radians), -sin(radians), 0, 0, sin(radians), cos(radians), 0, 0, 0, 0, 1, 0, 0, 0, 0, 1]).reshape(4, 4)]
        for i in range(len(self.vertices)):
            self.vertices[i] = list(dot(list(self.vertices[i]), rotation[axis]))
        self.rotated[axis] = (self.rotated[axis] + angle) % 360
        return self

    def draw(self, surface = screen, fov = 256, distance = 300):
        faces = {
            'cube': [(0,1,2,3), (4,5,6,7), (0,3,7,4), (1,2,6,5), (0,4,5,1), (3,7,6,2)],
            'pyramid': [(0,1,2), (0,1,3), (0,2,3), (1,2,3)],
            'diamond': [(5,1,2), (5,2,3), (5,3,4), (5,1,4), (0,1,2), (0,2,3), (0,3,4), (0,1,4)]}
        self.faces = faces[self.shape]
        z_index = {}
        for i in range(len(self.faces)):
            z_index[i] = 0
            for j in range(len(self.faces[i])):
                z_index[i] += self.vertices[self.faces[i][j]][2]
            z_index[i] /= len(self.faces)
        z_index = sorted(z_index, key = z_index.get, reverse = True)
        projected = []
        for v in self.vertices:
            factor = fov / (distance + float(v[2]))
            projected.append((v[0] * factor + self.pos[0], v[1] * factor + self.pos[1]))
        for f in z_index:
            points = tuple([projected[e] for e in self.faces[f]])
            pygame.draw.polygon(surface, self.colors[f], points)

opts = ['cube', 'pyramid', 'diamond']
shapes = []

for i in range(40):
    rand_colors = [random_color() for e in range(8)]
    shapes.append(Shape(opts[randint(0, 2)], [randint(20, 80) for e in range(3)], rand_colors, (randint(0, win_size[0] * .9), randint(0, win_size[1] * .9))))

done = False
keys_down = {}

step = 10
rotate_keys = {273: (0, step), 274: (0, -step), 275: (1, step), 276: (1, -step), 122: (2, step), 120: (2, -step)}

while not done:
    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            done = True
        if event.type == pygame.KEYDOWN:
            keys_down[event.key] = True
        if event.type == pygame.KEYUP:
            del(keys_down[event.key])
    for i in keys_down:
        if i in rotate_keys:
            for s in shapes:
                s.rotate(rotate_keys[i][0], rotate_keys[i][1])
    screen.fill((0, 0, 0))
    for s in shapes:
        s.draw()
    pygame.display.update()
    clock.tick(600)