#!/usr/bin/python
# -*- coding: utf-8 -*-
"""Simple Pygame test program for automatic differentiation test.

This currently displays the four corners of what is supposed to be a
3×4 rectangle after re-solving it from scratch after each frame.  We
can see that the optimization algorithm is kind of working.

If you click, it resets point 1 to where you clicked.

"""
from math import ceil

import pygame
from pygame import display, Surface, event, draw

import revdiff


if __name__ == '__main__':
    scr = display.set_mode((640, 480), 0)
    p1, p2, p3, p4, total = revdiff.dumb_geometric_example()

    inputs = None

    xmin = ymin = 0
    xmax = ymax = 1

    v = None

    #total = total * total  # square the error to make gradient descent work better
    while True:
        while True:
            ev = event.poll()
            if ev.type in (pygame.QUIT, pygame.NOEVENT):
                break
            elif ev.type == pygame.MOUSEBUTTONDOWN:
                if inputs is None:
                    continue
                x, y = ev.pos
                inputs[p1[0]], inputs[p1[1]] = x0 + x / xscale, y0 + y / yscale
            elif ev.type != pygame.MOUSEMOTION:
                print(ev)

        if ev.type in (pygame.QUIT, pygame.MOUSEBUTTONDOWN):
            break

        result = err, inputs, graph, v = revdiff.gdm_forward(total,
                                                             initial_learning=.0001,
                                                             final_learning=.0001,
                                                             iters=8,
                                                             width=1,  # i.e. no random restarts
                                                             start=inputs,
                                                             beta=0.9,
                                                             v=v)

        print(err)
        points = [(inputs[x], inputs[y]) for x, y in p1, p2, p3, p4]

        if err > 1000:
            inputs = None  # random restart!
            continue

        xmin = min(xmin, min(x for x, y in points))
        xmax = max(xmax, max(x for x, y in points))
        ymin = min(ymin, min(y for x, y in points))
        ymax = max(ymax, max(y for x, y in points))
        xdist = xmax - xmin
        ydist = ymax - ymin
        x0 = xmin
        y0 = ymin
        xscale = yscale = min(640 / (xmax - x0),
                              480 / (ymax - y0))

        scr.fill(0)
        for x, y in points:
            draw.ellipse(scr, 1234567, ((x - x0) * xscale - 5,
                                        (y - y0) * yscale - 5,
                                        10, 10))

        for x in range(int(x0), int(ceil(x0 + 640 / xscale))):
            for y in range(int(y0), int(ceil(y0 + 480 / yscale))):
                draw.ellipse(scr, 1234567, ((x - x0) * xscale - 1,
                                            (y - y0) * yscale - 1,
                                            2, 2))

        display.flip()