# import os
# from keras.preprocessing.image import ImageDataGenerator
# from PIL import Image
# from sklearn.metrics import accuracy_score
import gym
import time
import random
from collections import deque

from keras.models import Sequential
from keras.layers import Conv2D
from keras.layers import MaxPooling2D
from keras.layers import AveragePooling2D
from keras.layers import GlobalMaxPooling2D
from keras.layers import Dense
from keras.models import model_from_json

import win32gui
from PIL import ImageGrab
import cv2
import numpy as np


def conv_net():
    cnn = Sequential()
    cnn.add(Conv2D(32, (3, 3), input_shape=(128, 128, 3), padding='same', activation='relu'))
    cnn.add(Conv2D(32, (3, 3), input_shape=(128, 128, 3), padding='same', activation='relu'))
    cnn.add(MaxPooling2D(pool_size=(3, 3), padding='same', strides=(2, 2)))
    cnn.add(Conv2D(64, (5, 5), padding='same', activation='relu', strides=(2, 2)))
    cnn.add(Conv2D(64, (5, 5), padding='same', activation='relu', strides=(2, 2)))
    cnn.add(MaxPooling2D(pool_size=(2, 2), padding='same', strides=(2, 2)))
    cnn.add(Conv2D(128, (5, 5), padding='same', activation='relu'))
    cnn.add(Conv2D(128, (5, 5), padding='same', activation='relu'))
    cnn.add(Conv2D(256, (3, 3), padding='same', activation='relu'))
    cnn.add(Conv2D(256, (3, 3), padding='same', activation='relu'))
    cnn.add(AveragePooling2D(pool_size=(3, 3), padding='same'))
    cnn.add(GlobalMaxPooling2D())
    cnn.add(Dense(128, activation='relu'))
    cnn.add(Dense(64, activation='relu'))
    cnn.add(Dense(16, activation='relu'))
    cnn.add(Dense(2, activation='linear'))
    cnn.compile(optimizer='Adam', loss='mse')
    return cnn


def load_model():
    try:
        raise FileNotFoundError
        json_file = open('CNNmodel.json', 'r')
        loaded_model_json = json_file.read()
        json_file.close()
        cnn = model_from_json(loaded_model_json)
        cnn.load_weights("CNNmodel.h5")
        cnn.compile(optimizer='Adam', loss='mse')
        print(cnn.summary())
    except FileNotFoundError:
        cnn = conv_net()
        print(cnn.summary())
    return cnn


def save_model(cnn):
    model_json = cnn.to_json()
    with open("CNNmodel.json", "w") as json_file:
        json_file.write(model_json)
    cnn.save_weights("CNNmodel.h5")


def grab_frame():
    while True:
        try:
            hwnd = (win32gui.GetForegroundWindow())
            bbox = win32gui.GetWindowRect(hwnd)
            bbox = (bbox[0] + 10, bbox[1] + 50, bbox[2] - 5, bbox[3] - 5)
            img = np.asarray(ImageGrab.grab(bbox))
            return img
        except:
            pass



def reshape_frame(image):
    return cv2.resize(image, (128, 128))


def remember(state, action, reward, next_state, done):
    global memory
    memory.append((state, action, reward, next_state, done))


def act(state):
    global epsilon, cnn
    if np.random.rand() <= epsilon:
        return random.randrange(action_size)
    act_values = cnn.predict(state.reshape(1, 128, 128, 3))
    return np.argmax(act_values[0])


def replay(batch_size):
    global epsilon, epsilon_min, epsilon_decay, cnn, memory, gamma
    minibatch = random.sample(memory, batch_size)
    targets = []
    states = []
    for state, action, reward, next_state, done in minibatch:
        target = reward
        if not done:
            target = reward + gamma * np.amax(cnn.predict(next_state.reshape(1, 128, 128, 3))[0])
        target_f = cnn.predict(state.reshape(1, 128, 128, 3))
        target_f[0][action] = target
        targets.append(target_f)
        states.append(state)
    states = np.asarray(states)
    targets = np.asarray(targets)
    cnn.fit(states.reshape(batch_size, 128, 128, 3), targets.reshape(batch_size, 2), epochs=1, verbose=0)
    if epsilon > epsilon_min:
        epsilon *= epsilon_decay


memory = deque(maxlen=1000)
epsilon = 1.0
epsilon_decay = 0.995
epsilon_min = 0.01
gamma = 0.90
action_size = 2

cnn = load_model()

env = gym.make("CartPole-v0")
with open("CNN_Log.txt", 'w') as f:
    for games in range(1, 10000):
        env.reset()
        time.sleep(0.1)
        frame = grab_frame()
        s = reshape_frame(frame)
        steps = 0
        for time_t in range(500):
            env.render()
            a = act(s)
            _, r, d, _ = env.step(a)
            print(str(a), end=' ', file=f)
            frame = grab_frame()
            ns = reshape_frame(frame)
            time.sleep(0.1)
            remember(np.asarray(s).reshape(128, 128, 3), a, r, np.asarray(ns).reshape(128, 128, 3), d)
            s = ns
            if d:
                print("Episode: {}/{}, Score: {}".format(games, 10000, time_t), end='\n', file=f, flush=True)
                print("Episode: {}/{}, Score: {}".format(games, 10000, time_t))
                break
            steps += 1
        replay(10)
        if games % 100 == 0:
            save_model(cnn)

save_model(cnn)