import pygame
import numpy as np
import random
import matplotlib.pyplot as plt
from examples import Point
from examples import Examples
from neural_network import NeuralNetwork

arm_length = 75.0

def translate(center, angle):
    return Point(center.x + arm_length * np.sin(angle), center.y - arm_length * np.cos(angle))

def alpha_point(alpha, beta, center):
    first_point = translate(center, alpha)
    second_joint = translate(first_point, np.pi - beta + alpha)
    return first_point, second_joint

def unstandarise(ang):
    #ang_x= np.array(ang)/np.pi*0.8 - 0.1
    ang_x= np.array(ang)*np.pi
    #ang_x= (np.array(ang)+0.1)*np.pi/0.8
    return ang_x

def main():
    
    ex = Examples(arm_length)
    e = ex.generate(10000)

    #fig, ax = plt.subplots()
    #ax.axis('equal')
    #for (x, y) in e[0]:
    #    plt.scatter(x, y, marker='o')
    #plt.show()

    np.max(e[0]), np.max(e[1])
    x_train = (np.array(e[0]) + arm_length*2.0)/(arm_length*4.0)*0.8 + 0.1
    y_train = np.array(e[1])/np.pi*0.8 + 0.1
    NN = NeuralNetwork(layers=(2,10,2))
    for i in range(2000):
        NN.train(x_train, y_train)
    
    err = NN.errors
    plt.plot(range(len(err)), err)
    plt.show()

    pygame.init()
    pygame.font.init()
    myfont = pygame.font.SysFont('Open Sans', 24)
    screen = pygame.display.set_mode([500, 500])


    ex2 = Examples(arm_length)
    e2 = ex2.generate(10)
    xx = (np.array(e2[0]) + arm_length*2.0)/(arm_length*4.0)*0.8 + 0.1
    yy = np.array(e2[1])/np.pi*0.8 + 0.1



    screen.fill((255, 255, 255))
    image = pygame.image.load('robot.png')
    screen.blit(image, (0, 0))
    pygame.display.flip()
    running = True

    while running:
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                running = False
            elif event.type == pygame.MOUSEBUTTONDOWN:
                pos = pygame.mouse.get_pos()
                if pos[0] > 250 or True:
                    position = np.array(pos)
                    position[0] =  position[0] - 250
                    position[1] =  250 - position[1]
                    position = (np.array(position) + arm_length*2.0)/(arm_length*4.0)*0.8 + 0.1
                    ang = NN.forward( (position[0],position[1]) )
                    alpha, beta = unstandarise( ang )
                    points = alpha_point( np.pi - alpha, -beta,Point(250,250))
                    
                    print(points[0].x,points[0].y)
                    screen.fill((255, 255, 255))
                    image = pygame.image.load('robot.png')
                    screen.blit(image, (0, 0))
                    pygame.draw.line(screen, (255, 128, 0),
                                     (250, 250), (points[0].x, points[0].y), 5)
                    pygame.draw.line(
                        screen, (255, 128, 0), (points[0].x ,  points[0].y), (points[1].x , points[1].y ), 5)
                    pygame.display.flip()
                
    pygame.quit()

if __name__ == "__main__":
    main()