import plotutils
from mathutils import fit_gaussian, Gaussian

SINGLE_PASS = False
if SINGLE_PASS:
    STEP_SIZE = 0.5
else:
    STEP_SIZE = 1.0
RANGE = [-5.0, 5.0]
LATENCY = 0.005

robot.enable()
move_to_laser()

step_x = STEP_SIZE
step_z = STEP_SIZE
range_x = [RANGE[0], RANGE[1]]
range_z = [RANGE[0], RANGE[1]]

robot.set_motors_enabled(True)
current_x = robot_x.getPosition()
current_z = robot_z.getPosition()

print "Current pos x,z" , current_x, ",", current_z
ret = ascan([robot_x, robot_z], ue.readable, [range_x[0], range_z[0]], [range_x[1], range_z[1]], [step_x,step_z], latency = LATENCY, relative = True , zigzag=False)
data = ret.getData(0)[0]
#plot(Convert.transpose(data), title="Data")

integ = []
for x in data: integ.append(sum(   [ (0.0 if (math.isnan(y)) else y)  for y  in x]))

xdata= frange(range_x[0], range_x[1], step_x , False, True)
p = plot(integ, title = "Fit", xdata=xdata)[0]


max_x_index = integ.index(max(integ))
max_x = xdata[max_x_index]
(normalization, mean_val, sigma) = fit_gaussian(integ, xdata)
gaussian = Gaussian(normalization, mean_val, sigma)
plot_function(p, gaussian, "Fit", xdata, show_points = False, show_lines = True, color = Color.BLUE)

#So 
if abs(mean_val - max_x) > 1.0:
    raise Exception("Invalid X detection")
x_offset = mean_val    
center_x = current_x + x_offset

print "X offset = ", x_offset


if SINGLE_PASS:
    z_scan_data = data[max_x_index]        
else:
    robot_x.move(center_x)
    step_z = step_z/4.0
    ret2 = lscan(robot_z, ue.readable, range_z[0], range_z[1], step_z, latency = LATENCY, relative = True , zigzag=False)
    z_scan_data = ret2.getData(0)[0]
    
max_z_index= z_scan_data.index(max(z_scan_data))    
last_z_index = 0
for i in range(len(z_scan_data)):
    if not math.isnan(z_scan_data[i]):
        last_z_index = i    
#Shape is cone: z is inceraseing. For proper detection last Z must be furthest    
if abs(max_z_index - last_z_index) * step_z > 1.0: 
    raise Exception("Invalid Z detection")

if SINGLE_PASS:
    max_z = ret.getPositions(1)[len(data[0]) * max_x_index + last_z_index]    
else:    
    max_z = ret2.getPositions(0)[last_z_index]    

z_offset = max_z-current_z

print "Z offset = ", z_offset


#Updating tool:
#update_tool(None, x_offset=x_offset, z_offset=z_offset)