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
bunch             = 1

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


###################################################################################################
# Find COM
###################################################################################################
print "---------------          Find COM          --------------- "
set_status("Performing wire scan to find initial COM...")
write_ws_gain(ws_blm,ws_wire, initial_gain)
set_blm_ws_gain(ws_blm, initial_gain)
print "Setting gain=", initial_gain

#Calculate speed
x_min, x_max = range_start, range_end
rr = get_repetition_rate(bunch)
#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,bunch,0, True]
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/rr))
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_int = ws_blm + ":B1_LOSS"
ch_raw = ws_blm + ":LOSS_SIGNAL_RAW"
st.addScalar(ch_int, ch_int, int(100.0 / rr), 0)    
st.addWaveform(ch_raw, ch_raw, int(100.0 / rr), 0)   
st.addScalar("blm1_ws_mode", ws_blm + ":WS_RUNNING", int(100.0 / rr), 0)    
st.initialize()
st.start()
st.waitCacheChange(10000)    #Wait stream be running before starting scan
    
baseline=get_blm_baseline(ws_blm)

def change_blm_ws_gain(gain):
    set_blm_ws_gain(ws_blm, gain)      
    time.sleep(0.25)

def get_loss_int():   
    samples = []
    for i in range(10):
         st.waitCacheChange(-1)
         val = st.getValue(ch_int)
         samples.append(val)
    samples.remove(max(samples))#Remove max value
    return max(samples)
    
def get_loss_range():
    st.waitCacheChange(-1)
    raw = st.getValue(ch_raw)
    m = min(raw)   
    ret = (baseline-m)/baseline
    ret = min(max(ret, 0.0), 1.0)
    return ret 

change_blm_ws_gain(MIN_GAIN)
time.sleep(2.0)


pos = read_ws_gain(ws_blm, ws_wire)
loss =  get_loss_int()
print "Current Gain = ", pos, "\nStart Loss = ", loss 

"""
target = saturation * 0.8
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_int()
        print "Pos = ", pos, " Loss = ", loss
        if loss>=target:
            break                         
finally:
    stop_blm_ws(ws_blm)                                

print "Final Gain: ", pos, "\nFinal 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" 
"""

target = 0.9
loss_range = 0.0
blm_saturation = None
set_status("Searching loss range of " + str(target) + "...")
try:
    #Search loop
    for pos in frange(MIN_GAIN, MAX_GAIN, OPT_STEP, True):
        change_blm_ws_gain(pos)   
        time.sleep(0.1)
        loss_range = get_loss_range()
        print "Pos = ", pos, " Loss Range= ", loss_range
        if loss_range>=target:
            print "Done..."
            blm_saturation =  get_loss_int()
            set_blm_saturation(ws_blm,blm_saturation)
            break                         
finally:
    stop_blm_ws(ws_blm)                                

print "Final Gain: ", pos, "\nFinal Loss Range: ", loss_range , "\nBLM satiuration: ", blm_saturation 
result = "Loss range for final gain: " + str(loss_range) + " - Target value: " + str(target)  + "\n"
if blm_saturation is None:
    result = result + "Cannot reach target value. Setting gain to " + str(pos)  + "\n"
else:
    result = result + "Optimized gain: " + str(pos)  + " - BLM satiuration: " +str(blm_saturation) + "\n"  
    
st.close()
write_ws_gain(ws_blm, ws_wire, 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,0]
[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_prefix, ws_wire, x_range_min, x_range_max)


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

set_return(result)