import plotutils
from mathutils import fit_gaussian, Gaussian

SINGLE_PASS = False
if SINGLE_PASS:
    STEP_SIZE = 0.2
else:
    STEP_SIZE = 1.0
STEP_Z_FINAL = 0.1
    
RANGE = [-5.0, 5.0]
LATENCY = 0.05

Z_FINAL_OFFSET = 0.0


SINGLE_PASS = True
STEP_SIZE = 0.1


robot.enable()
move_to_laser()

step_y = STEP_SIZE
step_z = STEP_SIZE
range_y = [RANGE[0], RANGE[1]]
range_z = [RANGE[0], RANGE[1]]

robot.set_motors_enabled(True)
current_y = robot_y.getPosition()
current_z = robot_z.getPosition()

print "Current pos y,z" , current_y, ",", current_z
ret = ascan([robot_y, robot_z], ue.readable, [range_y[0], range_z[0]], [range_y[1], range_z[1]], [step_y,step_z], latency = LATENCY, relative = True , zigzag=False, title = "Scan XY")
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_y[0], range_y[1], step_y , False, True)
p = plot(integ, title = "Fit", xdata=xdata)[0]


max_x_index = integ.index(max(integ))
max_x = xdata[max_x_index]
try:
    (normalization, mean_val, sigma) = fit_gaussian(integ, xdata)
except:
    raise Exception("Invalid Fit")    
gaussian = Gaussian(normalization, mean_val, sigma)
xdata= frange(range_y[0], range_y[1], step_y/100.0 , False, True)
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 Y detection")
y_offset = mean_val    
center_y = current_y + y_offset

print "Y offset = ", y_offset

robot_y.move(center_y)
if SINGLE_PASS:
    z_scan_data = data[max_x_index]        
else:    
    step_z = STEP_Z_FINAL
    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 + Z_FINAL_OFFSET

robot_z.move(max_z + Z_FINAL_OFFSET)

print "Z offset = ", z_offset


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