import pygame
import keras
from keras.layers import Conv2D, MaxPooling2D, Input, Dense, Flatten, Dropout, LSTM
from keras.models import Sequential
import cv2
import numpy as np
import pickle
import os
from keras import backend as K
from keras.models import load_model
import random
user_train = False
model_train = False
if user_train:
    pygame.init()

    screen = pygame.display.set_mode([455,256])

    imgs = list(["output/"+fl for fl in os.listdir("output")])
    for i in imgs:
        if not i.endswith(".png"):
            imgs.remove(i)
    img_obs = list([pygame.image.load(img) for img in imgs])
    img_index = 0
    dataset = {}
    hold = None
    run_mainloop = True
    while True:
        m_pos = pygame.mouse.get_pos()
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                pygame.quit()
                raise SystemExit
            if event.type == pygame.MOUSEBUTTONDOWN:
                dataset[imgs[img_index]] = m_pos
                img_index+=1
                clicks = []
                print(dataset)
                if img_index == len(imgs):
                    pygame.quit()
                    run_mainloop = False

        if not run_mainloop:
            break

        screen.blit(img_obs[img_index],[0,0])
        pygame.display.flip()

    with open("dataset_raw.txt","wb") as fl:
        pickle.dump(dataset,fl)
else:
    with open("dataset_raw.txt","rb") as fl:
        dataset = pickle.load(fl)

train_x = []
train_y = []
i = 0
for imgfile in dataset:
    img_src = cv2.imread(imgfile,0).flatten()/255

    train_x.append(img_src)
    pos = dataset[imgfile]
    train_y.append(np.array([pos[0]/455,pos[1]/256]))
if model_train:
    train_x = np.array(train_x)
    train_y = np.array(train_y)
    print(train_x)
    print(train_y)

    model = Sequential()
    model.add(Dense(units=512,input_dim=256*455,activation='relu'))
    model.add(Dense(units=512,activation='relu'))
    model.add(Dense(units=512,activation='relu'))
    model.add(Dense(units=256,activation='relu'))
    model.add(Dense(units=256,activation='relu'))
    model.add(Dense(units=2,activation='tanh'))
    model.compile(loss="mean_squared_error",
                  optimizer='sgd',
                  metrics=['accuracy'])
    last_loss = 99999999999999 # Don't question it
    while 1:
        model.fit(train_x,train_y,epochs=5)
        loss = model.evaluate(train_x, train_y, batch_size=128,verbose=False)[0]

        model.save("desklamp.h5")
        #cv2.imwrite("tmp.png",model.predict(train_x[0].reshape(1,256,455,1))[0])
else:
    model = load_model("desklamp.h5")
cam = cv2.VideoCapture(0)
while True:
    ret_val, img = cam.read()
    img = cv2.flip(img,1)
    arr = cv2.cvtColor(cv2.resize(img, (455, 256), interpolation = cv2.INTER_CUBIC), cv2.COLOR_BGR2GRAY).flatten()/255
    output = model.predict(arr.reshape((1,455*256)))[0]
    pos = (int(output[0]*1280),int(output[1]*720))
    cv2.circle(img,pos,16,[255,0,0],-1)
    cv2.imshow("hackathon",img)
    key = cv2.waitKey(1)
    if key == 27:
        cv2.destroyAllWindows()
        break