import numpy as np
from EV3Robot import *

robot = Robot()
robot.connect_motor( 'left' )
robot.connect_motor( 'right' )
robot.connect_sensor( 'color' )

#syn0, syn1=train()

#syn0 = [[-0.05111396, -0.65152816],
 #[-0.60031302, -0.42462132],
 #[-0.61138226, -0.94931736],
 #[-0.54869344, -1.05984886],
 #[ 0.01826786,  0.24129794],
 #[ 0.14427619, -0.29771331],
 #[-0.74305135,  0.076575  ],
 #[-0.08188871, -0.31810276],
 #[-0.26736444,  0.20061916],
 #[-1.58011504, -1.44969247],
 #[-0.47182455,  2.43261639],
 #[ 2.17045716, -1.52101072]],

#syn1 = [[ -3.47021005, -61.89030242,  12.78322924],
 #[ -2.8108547,  22.64381972, -65.90059129]]

syn0 = [[-0.22387899,  0.0272135 ],
 [-0.53776054, -0.36374634],
 [-0.61603734, -0.48518088],
 [-0.57052126, -0.56987098],
 [ 0.2442093,   0.23051107],
 [-0.17470505,  0.42684878],
 [-0.53413718,  0.49524244],
 [-0.49455046,  0.49381262],
 [-0.17848447,  0.17378944],
 [-0.84917023, -1.20753232],
 [ 2.3610183,   0.48650416],
 [-1.5081981,   1.38669447]]

syn1 = [[ -1.50414787,  17.11258305, -19.87962655],
 [ -1.63558919, -16.87799893,   3.60433587]]

print("get train data",syn0)
print(syn1)


def nonlin(x, deriv=False):
    if (deriv == True):
        return x * (1 - x)

    return 1 / (1 + np.exp(-x))

def calculate_output(syn0, syn1, l0):
    l1 = nonlin(np.dot(l0, syn0))
    l2 = nonlin(np.dot(l1, syn1))
    return l2

def run1(speed):
    previous = []

    while(1):
        #check the state
        state_1 = get_state()
        #action taken according to maximum value of q table in color column
        if(previous==[]):
            action=1

        #get action from
       # action = np.argmax(q,axis=0)[state_1]
        else:
            vector=toBinary(previous)
            current=[0,0,0]
            current[int(state_1)]=1
            input=np.append(vector, current)

            l2 = calculate_output(syn0, syn1, [input])


            action=np.argmax(l2)
        # do the action
        do_action(action, speed)

        state_2=get_state()
        previous=[state_1, state_2, action]

def toBinary(set):
    vector = np.zeros((9,), dtype=np.int)
    vector[int(set[0])] = vector[int(set[1]) + 3] = vector[int(set[2]) + 6] = 1;
    return vector


def bgw(isee, follow_color=50, grey_zone=25):
    if isee < follow_color - grey_zone: # BLACK
        return 0
    elif isee > follow_color + grey_zone: # WHITE
        return 1
    else: # move forward if in the grey zone
        return 2

def get_state():
    isee = robot.color_sensor_measure('reflected_light_intensity')
    color = bgw(isee)
    return color

def do_action(action, speed):
    if action == 0:
        robot.move(0,speed) #left
    elif action == 1:
        robot.move(speed, 0) #right
    elif action == 2:
        robot.move(speed, speed) #forward

run1(40)