import sys
sys.path.append('../')

from europilot.screen import stream_local_game_screen
from europilot.screen import Box
from europilot.joystick import LinuxVirtualJoystick

"""
Explanation: Test the pretrained model
Get the game screen as the input, and use the virtual joystick to steer the wheel.
End of explanation
"""


%matplotlib inline
import matplotlib.pyplot as plt

import os
import numpy as np
from PIL import Image
from time import time

import keras
from keras.models import load_model
from keras.preprocessing import image

parent_path = os.path.dirname(os.getcwd())
model_path = os.path.join(parent_path, 'model')

# multiply by constant to undo normalization
OUTPUT_NORMALIZATION = 655.35

# limit GPU memory usage
import tensorflow as tf
from keras.backend.tensorflow_backend import set_session
config = tf.ConfigProto()
config.gpu_options.per_process_gpu_memory_fraction = 0.5
set_session(tf.Session(config=config))

"""
Explanation: Load pretrained model, and check performance for each forward pass.
The performance needs to be checked, because it determines the fps of our program.
End of explanation
"""


%ls $model_path

model = load_model(os.path.join(model_path, 'v5-PilotNet_v2-009-0.67221.h5'))

"""
Explanation: load pretrained model
End of explanation
"""


%ls ../sample/img/raw

sample_img = Image.open(os.path.join('../sample/img/raw', '9d0c3c2b_2017_07_27_14_55_08_16.jpg')).convert('RGB')

front_coord = (289,167, 851, 508)
sample_img_front = sample_img.crop(front_coord)

plt.imshow(sample_img_front)

sample_arr = image.img_to_array(sample_img_front)
sample_arr = np.reshape(sample_arr, (1,) + sample_arr.shape)

model.predict(sample_arr, batch_size = 1)
start = time()
for i in range(100):
    model.predict(sample_arr, batch_size = 1)
end = time()

fps = 100. / (end - start)
print("fps: %f" % fps)

"""
Explanation: Get sample image, to find out the performance of our model
End of explanation
"""


x1, y1 = (68, 54)
x2, y2 = (x1 + 1024, y1 + 768)
box = Box(x1, y1, x2, y2)

joy = LinuxVirtualJoystick()

"""
Explanation: Be aware that the performance of the model can affect the performance of the overall program, since it may not react quickly enough to the changing environment.
Setup virtual joystick.
Make sure to setup the virtual joystick inside ETS2 inside the controller settings.
Get the coordinates for the game screen. The values may be different per game setting.
End of explanation
"""


!cat /proc/bus/input/devices

"""
Explanation: As seen below, a virtual device can now be detected.
End of explanation
"""


# values computed from dataset sample.
def normalize(img):
    img[:,:,0] -= 89.5761
    img[:,:,0] /= 58.4214

    img[:,:,1] -= 97.5966
    img[:,:,1] /= 61.7917

    img[:,:,2] -= 88.3135
    img[:,:,2] /= 68.2043
    
    return img

def get_angle(predict):
    angle = predict[0][0]
    angle *= OUTPUT_NORMALIZATION
    
    return int(angle)

streamer = stream_local_game_screen(box=box, default_fps=60)
while True:
    image_data = next(streamer)
    im = Image.fromarray(image_data)
    img_front = im.crop(front_coord)
    
    arr = image.img_to_array(img_front)
    arr = normalize(arr)
    arr = np.reshape(arr, (1,) + arr.shape)
    
    angle = get_angle(model.predict(arr, batch_size = 1))
    
    joy.emit(angle)

"""
Explanation: Set up game loop
End of explanation
"""


