#!/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, x, y, i_max):
        config = Config(double_buffer=True)
        super().__init__(x, y, resizable=True, config=config, vsync=False)

        self.rotate = True
        self.sensitivity = 0.2
        self.show_fps = True

        self.x = 0
        self.y = 0
        self.z = -500

        self.rx = 0
        self.ry = 0
        self.rz = 0

        glClearColor(0, 0, 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)
                  t = int((x + 1) / 2) if e % 1 else int((x + 2) / 2)

                  #if not is_prime(a) and not is_prime(b):
                  #     continue

                  if e not in theend:
                       theend[e] = {}
                  if n not in theend[e]:
                      theend[e][n] = {}
                  if t not in theend[e][n]:
                      theend[e][n][t] = []

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

                  theend[e][n][t] = [e, n, d, x, a, b, c]
                  theend[f][n - 1][t] = [f, n - 1, d + 1, x + 1, a, b, c]

        print("Generating vertexes...")

        points = []
        points_rgb = []
        lines = []
        lines_rgb = []

        for e in theend:
            for n in theend[e]:
                for t in theend[e][n]:
                    points.extend([e, n, d])
                    points_rgb.extend([128, 128, 128])

                    _, _, _, _, _, _, C = theend[e][n][t]

                    _e = e
                    _t = 1

                    if 1 not in theend[e] or _t not in theend[e][1]:
                        break
                    _, _, d, x, a, b, c = theend[e][1][_t]
                    for i in range(100):
                        _t += 1
                        x += 2
                        a = b
                        b = a + x * 2 + n * 2
                        d = a + x
                        c = a * b

                        if C % a == 0:
                            break

                    p = a
                    i = 1
                    _t = 1
                    while True:
                        _t += p
#                        _t = (i * p + 1) - _t

                        _e, _n, _d, _x, _a, _b, _c = self.ent(e, 1, _t)

                        if _c == C:
                            lines.extend([e, n, d])
                            lines.extend([_e, _n, _d])
                            lines_rgb.extend([255, 0, 0])
                            lines_rgb.extend([255, 165, 0])
                            break

                        if _c > C:
                            break

                        i += 1

        print("Done")

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

        self.fps = pyglet.clock.ClockDisplay()
        pyglet.clock.schedule(self.update)

    def ent(self, e, n, t):
        if e == 0:
            x1 = 2
            d1 = 0
            d2 = t+1
            d3 = t
        elif e % 2 == 0:
            x1 = 0
            d1 = int(e/2)
            d2 = t
            d3 = t-1
        else:
            x1 = 1
            d1 = int((e+3)/2)
            d2 = t+1
            d3 = t-1
        dt = 2 * d3 * d2 + d1
        xt = 2 * (t - 1) + x1
        ct = dt * dt + e;
        at = dt - xt;
        if at == 0:
            return None
        bt = int(ct / at);
        return [e, n, dt, xt, at, bt, ct]

    def on_resize(self, width, height):
        glViewport(0, 0, width, height)
        return pyglet.event.EVENT_HANDLED

    def on_draw(self):
        glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT)

        glMatrixMode(GL_PROJECTION)
        glLoadIdentity()
        gluPerspective(60.0, float(self.width / self.height), 0.1, 10000.0)
        glMatrixMode(GL_MODELVIEW)
        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()

        if self.show_fps:
            glMatrixMode(GL_PROJECTION)
            glLoadIdentity()
            gluOrtho2D(0, self.width, 0, self.height)
            glMatrixMode(GL_MODELVIEW)
            glLoadIdentity()
            self.fps.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:
            epoch = int(time())
            name = "end_{}.png".format(epoch)
            self.save(name)
        if symbol == pyglet.window.key.F:
            self.show_fps = not self.show_fps
        if symbol == pyglet.window.key.R:
            self.rotate = not self.rotate
        return pyglet.event.EVENT_HANDLED

    def save(self, name):
        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)
        img.save(name)

        print("Saved {}".format(name))

    def update(self, dt):
        if not self.rotate:
            return
        self.rx += 0.3
        if self.rx > 360:
            self.rx = 0
        self.ry += 0.6
        if self.ry > 360:
            self.ry = 0
        self.rz += 0.9
        if self.rz > 360:
            self.rz = 0

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



end = TheEnd(x=1280, y=720, i_max=512)
pyglet.app.run()