#!/usr/bin/env python3

from math import sqrt
from gmpy2 import is_prime
from PIL import Image
from pyglet.gl import *
from time import time

import pyglet

class TheEnd(pyglet.window.Window):
    def __init__(self, i_max = 512, x_min = 0, y_min = 0, x_max = 64, y_max = 64):
        config = Config(double_buffer=True)
        super().__init__(1920, 1080, resizable=True, config=config)

        self.sensitivity = 0.2
        self.x = 0
        self.y = 0
        self.z = -100
        self.rx = 0
        self.ry = 0
        self.rz = 0

        glClearColor(0, 0, 0, 0)
        glColor3f(1, 0, 0)

        print("Creating the End...")

        theend = {}

        for i in range(i_max):
             for j in range(i):
                  a = i - j
                  b = i + j
                  c = a * b
                  d = int(sqrt(c))
                  e = int(c - (d * d))
                  f = int(e - ((2 * d) + 1))
                  n = int(i - d)
                  x = int(d - a)

                  if e not in theend:
                      theend[e] = {}
                  if n not in theend[e]:
                      theend[e][n] = []
                  if f not in theend:
                      theend[f] = {}
                  if n - 1 not in theend[f]:
                      theend[f][n - 1] = [];

                  theend[e][n].append([e, n, d, x, a, b, c])
                  theend[f][n - 1].append([f, n - 1, d + 1, x + 1, a, b, c])

        print("Generating points...")

        points = []
        colors = []
        for e in theend:
            for n in theend[e]:
                for cell in theend[e][n]:
                    _, _, d, x, a, b, c = cell
                    if not c % 2 == 1:
                        continue
                    if is_prime(a):
                        colors.extend([255, 0, 0])
                    else:
                        colors.extend([255, 255, 255])
                    points.extend([e, n, a])

        self.batch = pyglet.graphics.Batch()
        self.batch.add(len(points)//3, GL_POINTS, None, ("v3i", points), ("c3B", (colors)))

    def on_resize(self, width, height):
        glViewport(0, 0, width, height)
        glMatrixMode(GL_PROJECTION)
        glLoadIdentity()
        gluPerspective(60., width / float(height), .1, 1000.)
        glMatrixMode(GL_MODELVIEW)
        return pyglet.event.EVENT_HANDLED

    def on_draw(self):
        glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT)
        glLoadIdentity()
        glTranslatef(0, 0, self.z)
        glTranslatef(self.x, self.y, 0)
        glRotatef(self.rz, 0, 0, 1)
        glRotatef(self.ry, 0, 1, 0)
        glRotatef(self.rx, 1, 0, 0)
        self.draw()
        return pyglet.event.EVENT_HANDLED

    def on_mouse_drag(self, x, y, dx, dy, buttons, modifiers):
        if buttons == pyglet.window.mouse.LEFT:
            self.x += dx * self.sensitivity
            self.y += dy * self.sensitivity
        elif buttons == pyglet.window.mouse.RIGHT:
            self.rx += dy * self.sensitivity
            self.ry += dx * self.sensitivity
        return pyglet.event.EVENT_HANDLED

    def on_mouse_scroll(self, x, y, scroll_x, scroll_y):
        self.z += scroll_y * 10
        return pyglet.event.EVENT_HANDLED

    def on_key_press(self, symbol, modifier):
        if symbol == pyglet.window.key.SPACE:
            buf = (GLubyte * (self.width * self.height * 3))(0)
            glReadPixels(0, 0, self.width, self.height, gl.GL_RGB, gl.GL_UNSIGNED_BYTE, buf)
            img = Image.frombytes(mode="RGB", size=(self.width, self.height), data=buf)
            img = img.transpose(Image.FLIP_TOP_BOTTOM)
            epoch = int(time())
            name = "end_{}.png".format(epoch)
            img.save(name)

            print("Saved {}".format(name))
        return pyglet.event.EVENT_HANDLED

    def draw(self):
        self.batch.draw()



end = TheEnd()
pyglet.app.run()