import traceback

is_panel = get_exec_pars().source != CommandSource.ui  #Must be checked before callin "run"

run("Devices/Elements")
run("Devices/WireScanner")
run("Diagnostics/sig_process_wrapper")

###################################################################################################
# Arguments and constants
###################################################################################################

ws_prefix         = args[0] if is_panel else "SINDI01-DWSC090" #"S10DI01-DWSC010" #"S10CB07-DWSC440"  #"SINDI01-DWSC090" \\
plt               = args[1] if is_panel else plot(None, title = "Wire Scan Calibration")[0]
ws_wire           = args[2] if is_panel else WireScanner.WireX1
range_start       = args[3] if is_panel else -2000
range_end         = args[4] if is_panel else 2000
n_shot            = args[5] if is_panel else 200
saturation        = args[6] if is_panel else 19950 / 10        #S10DI01-DBLM113:AL1-WS-PMT-GAIN
scan_range_factor = args[7] if is_panel else  6.0
initial_gain      = args[8] if is_panel else  0.6

ws_blm       = get_wire_scanners_blms(ws_prefix )[0]

MIN_GAIN, MAX_GAIN = 0.5, 1.1
OPT_STEP = 0.02


print "WireScanCalibration parameters: ", ws_prefix, ws_wire, range_start, range_end, n_shot, saturation, scan_range_factor, initial_gain


###################################################################################################
# Utilities
###################################################################################################
def write_ws_gain(val):
    set_setting(ws_blm + "GainWs" + ws_wire, val)

def read_ws_gain():
    return get_setting(ws_blm + "GainWs" + ws_wire)

def get_gain():
    return get_blm_ws_gain(ws_blm)
    
def set_gain(val):
    set_blm_ws_gain(ws_blm,val)    

def set_wire_scan_range(wire, start, end):  
    sel = {'X1':"W1X" , 'Y1': "W1Y", 'X2':"W2X", 'Y2' : "W2Y"}
    start = min (max(start, -2000), 2000.0)
    end = min (max(end, -2000), 2000.0)
    caput((ws_prefix + ":" + sel[wire] +"_START_SP"), start)
    caput((ws_prefix + ":" + sel[wire] +"_END_SP"), end)    



###################################################################################################
# Find COM
###################################################################################################
print "---------------          Find COM          --------------- "
set_status("Performing wire scan to find initial COM...")

set_gain(initial_gain)
print "Setting gain=", initial_gain

#Calculate speed
x_min, x_max = range_start, range_end
#rr = get_repetition_rate()
#ws_speed = (x_max- x_min)*rr/n_shot
args = [ ws_prefix , ws_wire, [x_min, x_max, x_min, x_max], 1, n_shot, [], [ws_blm], 10, plt, False,1]
ret = run("Diagnostics/WireScan", args)
[rms_com, rms_sigma, com, sigma, pos_path,  path] = ret



###################################################################################################
# Optimize gain
###################################################################################################
print "---------------         Optimize gain         --------------- "



#Need 30s for 25Hz at full gain range. Have a 2x factor.
start_blm_ws(ws_blm, 2*7.5*(100/get_repetition_rate()))
ws_info = WireScanInfo("ws_info", ws_prefix )
motor_pos= ws_info.get_motor_pos(com, ws_wire)

set_status("Setting motor position to COM...")
caput(ws_prefix+":MOTOR_1.VELO", caget(ws_prefix + ":TRAVEL_VELO_SP"))
caput(ws_prefix+":MOTOR_1.VAL", motor_pos) #DVAL?





print "Starting stream..."
set_status("Creating stream for gain search..." )

st = Stream("blm_stream", dispatcher)
ch = ws_blm + ":B1_LOSS"
st.addScalar(ch, ch, int(100.0 / get_repetition_rate()), 0)    
st.addScalar("blm1_ws_mode", ws_blm + ":WS_RUNNING", int(100.0 / get_repetition_rate()), 0)    
st.initialize()
st.start()
st.waitCacheChange(10000)    #Wait stream be running before starting scan
    

def change_blm_ws_gain(gain):
    set_gain(gain)
    #stop_blm_ws(ws_blm)                
    #st.getChild("blm1_ws_mode").waitValue(0, SET_BLM_WS_BS_READBACK_TIMEOUT)    
    #start_blm_ws(ws_blm, 600.0)                
    #st.getChild("blm1_ws_mode").waitValue(1, SET_BLM_WS_BS_READBACK_TIMEOUT)        
    time.sleep(0.25)

def get_loss():   
    global ch
    samples = []
    for i in range(10):
         st.waitCacheChange(-1)
         val = st.getValue(ch)
         samples.append(val)
    samples.remove(max(samples))#Remove max value
    return max(samples)

change_blm_ws_gain(MIN_GAIN)
time.sleep(2.0)

start_gain = get_gain()
pos = start_val = get_loss()
loss = get_loss()
target = saturation * 0.8

print "Start Gain = ", start_gain
print "Start Loss = ", start_val 
print "Target = ", target 
set_status("Searching gain to match peak losses of " + str(target) + "...")




try:
    #Search loop
    for pos in frange(MIN_GAIN, MAX_GAIN, OPT_STEP, True):
        change_blm_ws_gain(pos)   
        loss = get_loss()
        print "Pos = ", pos, " Loss = ", loss
        if loss>=target:
            break           
        #stop_blm_ws(ws_blm)                
finally:
    stop_blm_ws(ws_blm)                                


print "Final Gain: ", pos
print "Final Loss: ", loss 


result = "Loss value for final gain: " + str(loss) + " - Target value: " + str(target)  + "\n"

if loss<target:
    result = result + "Cannot reach target value. Setting gain to " + str(pos)  + "\n"
else:
    result = result + "Optimized gain: " + str(pos)  + "\n" 

st.close()
write_ws_gain(pos)
time.sleep(1.0)

###################################################################################################
# Optimize scan range
###################################################################################################

set_exec_pars(reset=True, defaults=True)

print "---------------       Optimize scan range       --------------- "
set_status("Performing scan with optimal gain to define optimal range..." )
#caget(ws_blm+":SAT_RAW_SUM")

args = [ ws_prefix , ws_wire, [x_min, x_max, x_min, x_max], 1, n_shot, [], [ws_blm], 10, plt, False,1]
[rms_com, rms_sigma, com, sigma, pos_path,  path] = run("Diagnostics/WireScan", args)

x_range_min, x_range_max  = com - scan_range_factor * sigma,  com + scan_range_factor * sigma
print "Optimized range: " , x_range_min , " to " , x_range_max 
set_wire_scan_range(ws_wire, x_range_min, x_range_max)


result = result + "Optimized range: " + str(x_range_min) +  " to " + str(x_range_max)

set_return(result)