import traceback

is_embedded = (get_exec_pars().getCommand(False).parent != None)
has_args = is_embedded or (get_exec_pars().source != CommandSource.ui)

STREAM_CREATION_RETRIES = 3
MAX_RANGE_STEP = 3

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

#LayoutSF
DATA_GROUP_PREFIX = "data/"  
DATASET_SUFIX = "/value"

set_exec_pars(layout="default")
#LayoutDefault
DATA_GROUP_PREFIX = "" 
DATASET_SUFIX = ""

#Paramter parsing
prefix     = args[0] if has_args else "SARCL01-DWSC160" # "SINDI01-DWSC090" # "S10DI01-DWSC010" #"S10CB07-DWSC440"  #"SINDI01-DWSC090" 
scan_type  = args[1] if has_args else WireScanner.WireX1
scan_range = args[2] if has_args else []
cycles     = args[3] if has_args else 2
#velocity   = args[4] if has_args else 200
n_shot     = args[4] if has_args else 200
bpms       = args[5] if has_args else [] #get_wire_scanners_bpms(prefix)
blms       = args[6] if has_args else get_wire_scanners_blms(prefix)
bkgrd      = args[7] if has_args else 10
plt        = args[8] if has_args else plot(None, title = "Wire Scan")[0]
save_raw   = args[9] if has_args else False
bunch      = args[10] if has_args else 1
adaptive   = args[11] if has_args else 1 #0=Off, 1=Gain, 2=Gain+range
filter_beam_ok = (args[12] if (has_args and len(args)>12) else True)
do_elog    = True if (has_args and (not is_embedded) and (plt is not None)) else False
print  has_args, is_embedded, do_elog

print "WireScan parameters: ", prefix, scan_type, scan_range, cycles, n_shot, bpms, blms, bkgrd, bunch, adaptive

#Adaptive mode parameters
MIN_GAIN, MAX_GAIN = 0.5, 1.1
INCREMENT_FACTOR = 0.5
SCAN_RANGE_FACTOR = 6
MAX_RANGE_STEP = 300

SET_BLM_WS_MODE = True
SET_BLM_WS_SETTLING_TIME = 2.0
SET_BLM_WS_BS_READBACK_TIMEOUT = 10000 #ms

ADD_CHANNELS = [] 
#ADD_CHANNELS = ["SARFE10-PBPG050:HAMP-INTENSITY-CAL", "SARFE10-PBIG050-EVR0:CALCI"]

#TODO: configure biggerCAJ buffer size and set back to 10000
SAMPLE_CHANNEL_SIZE = 2049 #10000 
SAMPLE_CHANNEL_EMPTY_VALUE = 0.0

BPM_SENSORS = [("x","X"+str(bunch)), ("y","Y"+str(bunch)), ("q","Q"+str(bunch))] #(logic name suffix, channel suffix)

#Plot setup
if plt is not None:
    plt.clear()
    plt.removeMarker(None)
    plt.getAxis(plt.AxisId.X).setLabel("Position");
    plt.getAxis(plt.AxisId.Y).setLabel("");
    plt.getAxis(plt.AxisId.Y2).setLabel("");
    plt.setLegendVisible(True);       
snapshots = []

if scan_range is None or len(scan_range) !=4:
    if scan_type in  [WireScanner.WireX2, WireScanner.WireY2, WireScanner.Set2]:
        scan_range = [ caget(prefix+":W2X_START_SP", 'd'), \
                       caget(prefix+":W2X_END_SP", 'd'), \
                       caget(prefix+":W2Y_START_SP", 'd'), \
                       caget(prefix+":W2Y_END_SP", 'd') ]
    else:
        scan_range = [ caget(prefix+":W1X_START_SP", 'd'), \
                       caget(prefix+":W1X_END_SP", 'd'), \
                       caget(prefix+":W1Y_START_SP", 'd'), \
                       caget(prefix+":W1Y_END_SP", 'd') ]          


rr = get_repetition_rate(bunch)
velocity_x = abs(scan_range[1]-scan_range[0])*rr/n_shot
velocity_y = abs(scan_range[3]-scan_range[2])*rr/n_shot

#Creating WireScanner object
print "Creating scanner..."
if prefix not in get_wire_scanners():
    raise Exception("Invalid wire scan: " + prefix)    
scanner = WireScanner(prefix, scan_range, cycles, None, True)

#List of stream channels
channels = [("m_pos", scanner.motor_bs_readback.get_channel_name()), 
            ("cur_cycle", scanner.curr_cycl.get_channel_name()),
            ("scanning", scanner.status_channels[0].get_channel_name())]  


for i in range (len(blms)):
    channels.append (("blm" + str(i+1), blms[i] + ":B" + str(bunch) + "_LOSS_RAW"))
    #channels.append (("blm" + str(i+1), blms[i] + ":B" + str(bunch) + "_LOSS"))
    if plt is not None:
        series = LinePlotSeries(blms[i], None, min(i+1, 2))
        plt.addSeries(series)     
        series.setLinesVisible(False)
        series.setPointSize(2)
    if save_raw:    
        channels.append (("blm" + str(i+1) + "_raw" , blms[i] + ":LOSS_SIGNAL_RAW"))
for i in range (len(bpms)):
    for sensor in BPM_SENSORS:
        channels.append (("bpm" + str(i+1) + "_" + sensor[0], bpms[i] + ":" + sensor[1]))

if SET_BLM_WS_MODE and (len(blms)>0):
    channels.append(("blm1_ws_mode", blms[0] + ":WS_RUNNING"))   
channels.append(("beam_ok", get_beam_ok_channel(bunch)))       

for ch in ADD_CHANNELS:
    channels.append((ch, ch))       
    
#Metadata
set_attribute("/", "Wire Scanner", prefix)    
set_attribute("/", "Scan Type", scan_type)    
set_attribute("/", "Range", scan_range)    
set_attribute("/", "Cycles", cycles)    
set_attribute("/", "Scan Points", n_shot)  
set_attribute("/", "Motor Velocity X", velocity_x*math.sqrt(2))  
set_attribute("/", "Wire Velocity X", velocity_x)
set_attribute("/", "Motor Velocity Y", velocity_y*math.sqrt(2))  
set_attribute("/", "Wire Velocity Y", velocity_y)
set_attribute("/", "Background Measures", bkgrd)
set_attribute("/", "BPMs", bpms)
set_attribute("/", "BLMs", blms)
set_attribute("/", "Bunch", bunch)

filename = get_exec_pars().path



#Stream creation
for retry in range(STREAM_CREATION_RETRIES):
    try:
        print "Starting stream..."
        st = Stream("pulse_id", dispatcher)
        for c in channels:
            if c[1].endswith("LOSS_SIGNAL_RAW"):
                st.addWaveform(c[0], c[1], int(100.0 / rr), 0)   
            else:
                st.addScalar(c[0], c[1], int(100.0 / rr), 0)    
        st.initialize()
        st.start()
        add_device(st, True)
        st.waitCacheChange(10000)    #Wait stream be running before starting scan
        break
    except:
        print sys.exc_info()[1]
        if retry >= (STREAM_CREATION_RETRIES-1):
            channels_names =  [dev.name for dev in st1.children]
            msg = "Error creating stream with channels:\n" + "\n".join(channels_names)
            show_message(msg)
            print(msg)
            raise
        else:
            print "Error creating stream, retrying..."
           
class Timestamp(Readable):
    def read(self):        
        return st.getTimestamp()

        

   
#Pseudo-device returning the wire position
class w_pos(Readable):
    def read(self):
        return scanner.get_sel_wire_pos(st.getChild("m_pos").take())

#End of  scan checking
scan_complete, cur_cycle, wire = None, None, None
rec =None

def check_end_scan(record, scan):
    global scan_complete,cur_cycle
    global rec
    if (rec is None) and (record is not None):
        print "Started receiving from stream"
    rec = record
    if record["scanning"]<1:
        print "Data aquisition completed"
        scan_complete=True
        scan.abort()
        record.cancel() #So it won't be saved
    else:
        position = record["w_pos"]
        if record["cur_cycle"] != cur_cycle:
            cur_cycle = record["cur_cycle"]
            get_context().dataManager.splitScanData(scan)
            #if plt is not None: for s in plt.getAllSeries(): s.clear()
        if plt is not None:    
            for i in range (len(blms)):
                try:
                    val = record["blm" + str(i+1)]
                    if val is None:
                        val = float("NaN")
                        print "Null value for blm ", ("blm" + str(i+1)) 
                    plt.getSeries(i).appendData(position, val)                
                except:
                    print "Error plotting position " , position, record["blm" + str(i+1)]

if scanner.take() != "At start":
    print "Not at start: parking scanner"
    scanner.park(wait=True)        
#Process background
def do_background():
    #Store Background
    if bkgrd>0:
        #scanner.park(wait=True)        
        set_exec_pars(group = "background")
        r = mscan (st, st.getReadables()[4:], bkgrd)    
        for i in range(len(r.getReadables())):
            d = r.getReadable(i)              
            try:
                path = get_exec_pars().group + "/" + DATA_GROUP_PREFIX + r.getReadables()[i].name
                set_attribute(path, "Mean", mean(d))            
                set_attribute(path, "Sigma", stdev(d) )   
            except: 
                pass                             

def set_blm_gain(scan_type, index):
    if SET_BLM_WS_MODE and len(blms)>0:    
         if scan_type == WireScanner.Set1:
            scan_type = (WireScanner.WireX1 if index == 0  else WireScanner.WireY1)
         if scan_type == WireScanner.Set2:
            scan_type = (WireScanner.WireX2 if index == 0  else WireScanner.WireY2)
         cfg_gain = read_ws_gain (blms[0], scan_type)
         print "cfg_gain = " , cfg_gain
         if cfg_gain is not None:
             set_blm_ws_gain(blms[0],cfg_gain)    
             print "Set = " , blms[0], cfg_gain
     


#Scan   
def do_scan(index):    
    global scan_complete, cur_cycle, wire, rec
    rec =None
    wire = "y" if (index==1) or (scan_type in [WireScanner.WireY1, WireScanner.WireY2]) else "x"
    if wire == "x":        
        scanner.set_velocity(velocity_x)
    else:
        scanner.set_velocity(velocity_y)
    
    set_exec_pars(group=wire+"_{count}", reset=True)               
    scanner.set_selection(get_scan_selection(scan_type, index))    
    if plt is not None:
        if wire == "x":        
            plt.getAxis(plt.AxisId.X).setRange(scan_range[0], scan_range[1])        
        else:
            plt.getAxis(plt.AxisId.X).setRange(scan_range[2], scan_range[3])    

    if adaptive>0:
        set_blm_gain(scan_type, index)
            
    scanner.init(wait=True)
    scanner.curr_cycl.write(0)
    scan_complete=False
    cur_cycle = 1.0

    if plt is not None:
        for s in plt.getAllSeries():
            s.clear()
        plt.removeMarker(None)
    try:     
        scanner.scan()        #scanner.waitState(State.Busy, 60000)  Not needed as stream  filter will make the wait
        st.getChild("scanning").waitValue(1.0, 10000)    
        #print st.getValues()     

        #TODO: Check what the problem is
        #mscan (st, [w_pos(),] + st.getReadables() + [Timestamp(),], -1, -1, take_initial = True, after_read = check_end_scan)                           

        l=[w_pos()] ; l.extend(st.getReadables()); l.append(Timestamp())
        print "Start scan"
        mscan (st, l, -1, -1, take_initial = True, after_read = check_end_scan)                           
        print "End scan"
        #tscan([w_pos()] + st.getReadables() + [Timestamp(),], 10, 0.5)
    except:        
        print sys.exc_info()[1]
        if not scanner.isReady():
            print "Aborting  scan"
            scanner.abort()
        if not scan_complete:
            raise
    finally:        
        #Combining data of multiple series
        #s=plt.getSeries(0)
        #indexes = sorted(range(len(s.x)),key=lambda x:s.x[x])
        #x,y = [s.x[x] for x in indexes], [s.y[x] for x in indexes]
        #plot(y, xdata = x)    
        if rec is None:      
            raise Exception("Didn't receive from stream")
        if not scan_complete:
            raise Exception("Scan didn't conmplete")
        print "Calculating"
        rec
        calculate()        
        print "Ok"
        img_file = os.path.abspath(filename + "_" + get_exec_pars().group[0:1] + ".png")
        time.sleep(0.1) #Give some time to plot finish (async)
        if plt is not None:
            plt.saveSnapshot(img_file, "png")
            snapshots.append(img_file)     
        print "Finished"

pars = []
msg = ""
ret = []
def calculate():
    global msg
    stats, samples_val, samples_pos = [], [],[]
    for i in range(len(blms)):
        msg += "Wire " + wire + "  -  BLM " + str(i+1) + ":\n"
        try:          
            blm_back_path = "background/" + DATA_GROUP_PREFIX + "blm" + str(i+1)
            print "BLM path:  ", blm_back_path
            blm_back_attrs = get_attributes(blm_back_path)
            print "BLM attrs: ", blm_back_attrs
            bg = blm_back_attrs["Mean"] if bkgrd>0 else 0.0
            samples = [[], [], [], [], [], []]   
            for cycle in range (cycles):  
                pos_path = wire+"_" + ("%04d" % (cycle+1)) + "/" + DATA_GROUP_PREFIX+ "w_pos"  
                print "Loading: ", pos_path
                pos = load_data(pos_path +  DATASET_SUFIX)
                path = wire+"_" + ("%04d" % (cycle+1)) + "/" + DATA_GROUP_PREFIX + "blm" + str(i+1)
                print "Loading ", path
                data = load_data(path + DATASET_SUFIX)         
                print "OK"
                sp = data #blm_remove_spikes(data)
                sig = sp if bg is None else [v-bg for v in sp]   
                                    
                #print [com, rms]
                [off, amp, com, sigma] = profile_gauss_stats(pos, sig, off=None, amp=None, com=None, sigma=None)
                set_attribute(path, "Gauss COM", float("nan") if (com is None) else com)            
                set_attribute(path, "Gauss Sigma", float("nan") if (sigma is None) else sigma)                                            

                #[rms_com, rms_sigma] = profile_rms_stats(pos, sig,noise_std=0, n_sigma=3.5)
                [rms_com, rms_sigma] =  profile_rms_stats_with_estimate(pos, sig, com_estimate = com, window_size = sigma * 3.5)
                set_attribute(path, "RMS COM", float("nan") if (rms_com is None) else rms_com)            
                set_attribute(path, "RMS Sigma", float("nan") if (rms_sigma is None) else  rms_sigma)            


                if i==0: #Write channels for BLM 1
                     samples_val.extend(sig)
                     samples_pos.extend(pos)
        
                samples[0].append(rms_com);samples[1].append(rms_sigma);
                samples[2].append(com);samples[3].append(sigma);samples[4].append(amp);samples[5].append(off)
                #print [off, amp, com, sigma]
        
                #from mathutils import Gaussian
                #g = Gaussian(amp, com, sigma)
                #gauss = [g.value(v)+off for v in pos]
                #plot([data, sp, sig, gauss], ["data", "sp", "signal", "gauss", ], xdata = pos, title="Fit blm" + str(i+1) + " - " + str(cycle+1))         
                ret.extend([rms_com, rms_sigma, com, sigma, filename + "|"+ pos_path + DATASET_SUFIX,  filename + "|"+ path + DATASET_SUFIX])  
        
            stats.append([])
            for sample in samples:
                sample =  [v for v in sample if v is not None]
                stats[i].append( (mean(sample), stdev(sample)) if len(sample)>0 else (float("nan"), float("nan")) )
            if plt is not None:    
                plt.addMarker(stats[i][2][0], None, "Gcom=" + "%.2f" % stats[i][2][0], plt.getSeries(i).color)
                plt.addMarker(stats[i][0][0], None, "Rcom=" + "%.2f" % stats[i][0][0], plt.getSeries(i).color.brighter())
                
            msg += "    RMS COM: " + "%.4f" % stats[i][0][0]  +  " +- " +"%.4f" % stats[i][0][1] + "\n"  #unichr(0x03C3) + "="
            msg += "    RMS Sigma: " + "%.4f" % stats[i][1][0]  +  " +- " + "%.4f" % stats[i][1][1] + "\n"
            msg += "    Gauss COM: " + "%.4f" % stats[i][2][0]  + " +- " + "%.4f" % stats[i][2][1] + "\n"
            msg += "    Gauss Sigma: " + "%.4f" % stats[i][3][0]  +  " +- " + "%.4f" % stats[i][3][1] + "\n"

            if i==0: #Write channels for BLM 1
                scanner.set_out_com(bunch,wire,stats[i][0][0])
                scanner.set_out_rms(bunch,wire,stats[i][1][0])
                scanner.set_out_mean(bunch,wire,stats[i][2][0])
                scanner.set_out_sigma(bunch,wire,stats[i][3][0] )                
                scanner.set_out_amp(bunch,wire,stats[i][4][0])                
                scanner.set_out_off(bunch,wire,stats[i][5][0])
                if len(samples_pos) > SAMPLE_CHANNEL_SIZE:
                    samples_pos = samples_pos[0:SAMPLE_CHANNEL_SIZE]
                else:
                    samples_pos += [SAMPLE_CHANNEL_EMPTY_VALUE] * (SAMPLE_CHANNEL_SIZE - len(samples_pos))
                if len(samples_val) > SAMPLE_CHANNEL_SIZE:
                    samples_val = samples_pos[0:SAMPLE_CHANNEL_SIZE]
                else:
                    samples_val += [SAMPLE_CHANNEL_EMPTY_VALUE] * (SAMPLE_CHANNEL_SIZE - len(samples_val))                
                scanner.set_out_pos(bunch,wire,samples_pos)
                scanner.set_out_samples(bunch,wire,samples_val)                      

                try:
                    gauss = Gaussian(stats[i][4][0], stats[i][2][0], stats[i][3][0])
                    gauss = [gauss.value(v)+stats[i][5][0] for v in samples_pos]                    
                except:      
                    traceback.print_exc()              
                    gauss = [0.0] * SAMPLE_CHANNEL_SIZE
                scanner.set_out_gauss(bunch,wire,gauss)    

                #com = stats[i][2][0]
                #sigma = stats[i][3][0]
                #saturration_loss = get_blm_saturation(blms[i])
                #current_gain = get_blm_ws_gain(blms[i]) #read_ws_gain (blms[i], wire)
                #current_loss = get_blm_loss(blms[i])
                #global pars
                #pars = [blms, com, sigma, saturration_loss, current_gain, current_loss]
                com = stats[i][0][0]
                rms = stats[i][1][0]
                mn = stats[i][2][0]
                sigma = stats[i][3][0]
                (cur_x_range_min, cur_x_range_max) = (scan_range[0], scan_range[1])  if wire == "x" else    (scan_range[2], scan_range[3])     
                x_range_min, x_range_max  = mn - SCAN_RANGE_FACTOR * sigma,  mn + SCAN_RANGE_FACTOR * sigma                         
                off_min, off_max = (x_range_min - cur_x_range_min),  (x_range_max - cur_x_range_max)
                valid_range = (abs(off_min) < MAX_RANGE_STEP) and (abs(off_max) < MAX_RANGE_STEP)            
                valid_fitting =  (com!=float('NaN')) and (mn!=float('NaN')) and \
                                 (com>cur_x_range_min) and (com<cur_x_range_max) and \
                                 (mn>cur_x_range_min) and (mn<cur_x_range_max) 
                print "Scan range: ", (cur_x_range_min, cur_x_range_max) 
                print "Valid fitting: ", valid_fitting
                print "Valid range: ", valid_range
                if adaptive > 0  and \
                    valid_fitting and \
                    (scan_type not in [WireScanner.Set1, WireScanner.Set2]):                    
                     
                     new_gain, new_voltage = None, None
                     #desired_loss= 800 
                     desired_loss= get_blm_saturation(blms[i])                    
                     current_voltage = get_blm_ws_gain(blms[i]) #read_ws_gain (blms[i], scan_type)
                     #current_loss = float(get_blm_loss(blms[i]))
                     current_loss = max(sig)
                     print "current_loss=", current_loss
                     current_gain = get_gain_from_voltage(current_voltage)
                                          
                     dg =  current_gain * (desired_loss/ current_loss -1)
                     new_gain = current_gain + dg * INCREMENT_FACTOR
                     new_voltage = get_voltage_from_gain(new_gain)
                     new_voltage = max(min(new_voltage, MAX_GAIN),MIN_GAIN)
                     print "dg=", dg, " ng=", new_gain   
                     
                     if (new_voltage is not None) and (new_voltage != current_voltage):
                       logstr= "Adapting " + str(blms[i]) + " - " +  str(scan_type) + "gain: " +  str(new_voltage)
                       print logstr
                       log(logstr)
                       write_ws_gain(blms[i], scan_type, new_voltage)
                       set_blm_ws_gain(blms[i], new_voltage) 
                                         
                     if adaptive > 1:    
                         if valid_range:
                             new_min, new_max = (cur_x_range_min + INCREMENT_FACTOR*off_min), (cur_x_range_max + INCREMENT_FACTOR*off_max)
                             logstr= "Adapting range: " , new_min , " to " , new_max 
                             set_wire_scan_range(prefix, scan_type, new_min, new_max)                        
                             print logstr
                             log(logstr)                                      
                                                      
        except Exception, e:
            print >> sys.stderr, traceback.format_exc()
            msg += str(e)+ "\n"
                        

def get_scan_time():
    global bkgrd, scan_type, scan_range, velocity_x, velocity_y
    ret = 0
    if not scan_type in [WireScanner.WireY1, WireScanner.WireY2]:    
        ret +=  abs(scan_range[0] - scan_range[1])/velocity_x    #X SCAN
    if not scan_type in [WireScanner.WireX1, WireScanner.WireX2]:
        ret +=  abs(scan_range[2] - scan_range[3])/velocity_y    #Y SCAN
    ret += 2.0                                        #ACC/DEC
    ret *= cycles
    ret += bkgrd * (1.0/rr)        #BACK   
    print "Scan time = " + str(ret), 
    ret = ret * 2.0 + 10.0 #Tolernces
    print "; with tolerance = " + str(ret), 
    return ret
    
print "Starting scan..."         
try:     
    if SET_BLM_WS_MODE and len(blms)>0:
        for i in range(len(blms)):
            start_blm_ws(blms[i], get_scan_time())                
        #TODO: Wait for stream variable indicate bllm is in ws mode
        print "Waiting for WS mode..."         
        #time.sleep(SET_BLM_WS_SETTLING_TIME)
        st.getChild("blm1_ws_mode").waitValue(1, SET_BLM_WS_BS_READBACK_TIMEOUT)
    print "Reading background..."         
    do_background()       
    stream_filter = scanner.curr_cycl.get_channel_name() + ">0"
    if filter_beam_ok:
        stream_filter = stream_filter + " AND " + get_beam_ok_channel(bunch) + " == 1"
    st.setFilter(stream_filter)
    print "Executing scan 1..."         
    do_scan(0)
    if scan_type in [WireScanner.Set1, WireScanner.Set2]:        
        print "Executing scan 2..."         
        do_scan(1)                             
except:
    print >> sys.stderr, traceback.format_exc()
    raise
finally:
    try:
        scanner.park(wait=False)      
        pass
    except:
        pass
    if SET_BLM_WS_MODE and len(blms)>0:
        for i in range(len(blms)):
            stop_blm_ws(blms[i])        
    print "Closing scanner"
    scanner.close()
    print "Closing stream"    
    st.close()       
           
print msg

# save the entry in the logbook
if do_elog:
    if get_option("Generated data file:\n" + filename +"\n\n" + msg + "\n\n" + "Save to ELOG?", "YesNo") == "Yes":    
        log_msg = "Data file: " + filename
        log_msg = log_msg + "\nWire Scanner: " + prefix
        log_msg = log_msg + "\nScan Type: " + str(scan_type)
        log_msg = log_msg + "\nRange: " + str(scan_range)
        log_msg = log_msg + "\nCycles: " + str(cycles)
        log_msg = log_msg + "\nScan Points: " + str(n_shot)
        log_msg = log_msg + "\nWire Velocity X: " + str(velocity_x)
        log_msg = log_msg + "\nWire Velocity Y: " + str(velocity_y)
        log_msg = log_msg + "\nBackground Measures: " + str(bkgrd)
        log_msg = log_msg + "\nBPMs: " + str(bpms)
        log_msg = log_msg + "\nBLMs: " + str(blms)
        log_msg = log_msg + "\nBunch: " + str(bunch)
        
        log_msg = log_msg + "\n" + msg
        elog("Wire Scan", log_msg, snapshots)

set_exec_pars(open=False)
print ret
set_return(ret)