import math
import random
import time
import datetime
import json
from copy import copy

activation = math.tanh
deactivation = math.tanh

class NeuralNetwork:
    def __init__(self, neurons_size:list, weights:list=None):
        self.neuron_size = neurons_size
        for i in range(len(neurons_size)-1):
            self.neuron_size[i] += 1

        if weights is None:
            self.weights = []
            for i in range(len(neurons_size)-1):
                self.weights.append([[random.uniform(-1, 1) for x in range(neurons_size[i+1])] for y in range(neurons_size[i])])
        else:
            self.weights = weights
        
    def out(self, inp):
        out=[inp + [1]]
        for i in range(1, len(self.weights)+1):
            a = []
            for j in range(self.neuron_size[i]):
                s = sum([out[i-1][k] * self.weights[i-1][k][j] for k in range(self.neuron_size[i-1])])
                a.append(activation(s))
            if i != len(self.weights):
                a += [1]
            out.append(a)
        return out
            
    def correct(self, inp, answer, learning_rate=0.1):
        out = self.out(inp)
        errors = [[answer[i] - out[-1][i] for i in range(len(out[-1]))]]

        for i in range(len(self.weights) - 1, 0, -1):
            a = []
            for j in range(self.neuron_size[i]):
                s = sum([errors[0][k] * self.weights[i][j][k] for k in range(self.neuron_size[i + 1])])
                a.append((1 - out[i][j] ** 2) * s)
            errors.insert(0, a)

        for i in range(len(self.weights)):
            for j in range(self.neuron_size[i]):
                for k in range(self.neuron_size[i + 1]):
                    self.weights[i][j][k] += learning_rate * errors[i][k] * out[i][j]

        error_count = sum([sum(abs(en) for en in el) for el in errors])
        return out, error_count

    def save(self, name):
        with open(name, 'w') as f:
            neuron_size = copy(self.neuron_size)
            for i in range(len(neuron_size) - 1):
                neuron_size[i] -= 1
            f.write(json.dumps({'shape': neuron_size, 'weights': self.weights}))

    def open(name):
        with open(name, 'r') as f:
            data = json.loads(f.read())
            return NeuralNetwork(data['shape'], data['weights'])

    def show(self, name):
        import matplotlib.pyplot as plt
        import networkx as nx
        from networkx.drawing.nx_agraph import graphviz_layout

        G = nx.DiGraph()

        for layer in range(len(self.neuron_size)):
            for neuron in range(self.neuron_size[layer]):
                G.add_node((layer, neuron))

        for layer in range(len(self.neuron_size) - 1):
            for from_neuron in range(self.neuron_size[layer]):
                for to_neuron in range(self.neuron_size[layer + 1]):
                    weight = self.weights[layer][from_neuron][to_neuron]
                    G.add_edge((layer, from_neuron), (layer + 1, to_neuron), weight=weight)

        pos = graphviz_layout(G, prog='dot', args="-Grankdir=LR")

        edge_widths = [2 + abs(G.edges[edge]['weight']) for edge in G.edges]
        edge_alpha = [abs(activation(G.edges[edge]['weight']))/2 for edge in G.edges]
        edge_colors = ['green' if G.edges[edge]['weight'] >= 0 else 'red' for edge in G.edges]

        nx.draw_networkx_nodes(G, pos, node_size=300, node_color='skyblue', alpha=0.8)

        nx.draw_networkx_edges(G, pos, width=edge_widths, alpha=edge_alpha, edge_color=edge_colors, arrows=True)

        layer_labels = {}
        for layer in range(len(self.neuron_size)):
            for neuron in range(self.neuron_size[layer]):
                    layer_labels[(layer, neuron)] = f"{neuron}"
        nx.draw_networkx_labels(G, pos, labels=layer_labels, font_size=12, font_color='r')

        plt.axis('off')
        plt.title(f"Neural Network Weights")
        plt.savefig(f'{name}.png')

from neuro import NeuralNetwork
from tkinter import *
import matplotlib.pyplot as plt
import matplotlib.animation as animation
import random
import asyncio

mind = NeuralNetwork([25, 4])
canvas_size = 1280
realsize = 32

pix = canvas_size / realsize

canvas = [[0.1 for x in range(realsize)] for y in range(realsize)]
cellx, celly = 15, 5


def cellvision(vis):
	global cellx
	global celly
	global canvas
	inp = []

	if vis != -1:
		for j in range(vis):
			for i in range(vis):
				inp.append(canvas[int(cellx - vis // 2 + i) % realsize][int(celly - vis // 2 + j) % realsize])
		return inp

	inp.append(canvas[int(cellx + 0) % realsize][int(celly - 1) % realsize])
	inp.append(canvas[int(cellx + -1) % realsize][int(celly + 0) % realsize])
	inp.append(canvas[int(cellx + 1) % realsize][int(celly + 0) % realsize])
	inp.append(canvas[int(cellx + 0) % realsize][int(celly + 1) % realsize])

	return inp


def move(out):
	global cellx
	global celly
	if out == 0:
		celly -= 1
	if out == 1:
		cellx -= 1
	if out == 2:
		cellx += 1
	if out == 3:
		celly += 1
	if cellx == realsize:
		cellx = 0
	if cellx == -1:
		cellx = realsize - 1
	if celly == realsize:
		celly = 0
	if celly == -1:
		celly = realsize - 1
	cell(cellx, celly)
	return


def goodpoint(x, y):
	color = "#476042"
	x, y = x * pix, y * pix
	x1, y1 = (x - pix / 2), (y - pix / 2)
	x2, y2 = (x + pix / 2), (y + pix / 2)
	w.create_oval(x1, y1, x2, y2, outline=color, fill=color)


def badpoint(x, y):
	color = "#ff0000"
	x, y = x * pix, y * pix
	x1, y1 = (x - pix / 2), (y - pix / 2)
	x2, y2 = (x + pix / 2), (y + pix / 2)
	w.create_oval(x1, y1, x2, y2, outline=color, fill=color)


def cell(x, y):
	color = "#ffffff"
	x, y = x * pix, y * pix
	x1, y1 = (x - pix / 2), (y - pix / 2)
	x2, y2 = (x + pix / 2), (y + pix / 2)
	w.create_oval(x1, y1, x2, y2, outline=color, fill=color)


def canvas_print():
	global canvas
	w.delete("all")
	ans = ''
	for y in range(realsize):
		for x in range(realsize):
			ans += str(canvas[x][y]) + " "
			if canvas[x][y] == 1:
				goodpoint(x, y)
			if canvas[x][y] == -1:
				badpoint(x, y)
			if canvas[x][y] == 0:
				cell(x, y)
		ans += "\n"


def usergoodpoint(event):
	x, y = int(event.x / pix), int(event.y / pix)
	canvas[x][y] = 1


def userbadpoint(event):
	x, y = int(event.x / pix), int(event.y / pix)
	canvas[x][y] = -1


master = Tk()
master.title("Среда обучения")
w = Canvas(master, bg="black",
		   width=canvas_size,
		   height=canvas_size)
w.pack(expand=YES, fill=BOTH)
w.bind("<B1-Motion>", usergoodpoint)
w.bind("<B3-Motion>", userbadpoint)

iterat = -1
allg = 0
graphic = []
while True:
	iterat += 1
	if iterat % 200 == 0:
		plt.plot(graphic)
		plt.pause(0.0000001)
	good = 0
	if iterat % 10000 == 0:
		plt.close()
		mind.show()

	canvas[random.randint(0, realsize - 1)][random.randint(0, realsize - 1)] = 1
	canvas[random.randint(0, realsize - 1)][random.randint(0, realsize - 1)] = -1

	canvas_print()

	visn = cellvision(5)
	visnn = cellvision(-1)

	out = mind.out(visn)
	move_ = out[-1].index(max(out[-1]))
	mind.correct(visn, visnn, 0.1)

	answer = [0]*4
	answer[move_] = 1
	move(move_)

	if canvas[cellx][celly] == 1:
		good += 50
		canvas[cellx][celly] = 0.1

	elif canvas[cellx][celly] == -1:
		good -= 50
		canvas[cellx][celly] = 0.1

	# print(input())
	allg += good
	graphic.append(allg)

	master.title("Среда обучения: " + " i:" + str(iterat) + " good:" + str(good))
	master.update()

plt.show()
master.mainloop()