#import ch.psi.pshell.epics.ChannelDouble as ChannelDouble
import ch.psi.pshell.bs.Scalar as Scalar
import ch.psi.pshell.bs.Waveform as Waveform
import ch.psi.pshell.epics.DiscretePositioner as DiscretePositioner
import ch.psi.pshell.epics.ChannelDouble as ChannelDouble


#Arguments
SAMPLES = 100
GAINS = ["SINDI02-DBLM084:M06-1-CH03-V-MM","SINDI02-DBLM084:M06-2-CH03-V-MM","S10DI01-DBLM113:M06-1-CH03-V-MM"] #, "SINDI02-DBLM084:M06-0-CH03-V-MM"]
BLMS = ["SINDI02-DBLM025", "SINDI02-DBLM085", "S10DI01-DBLM045"]
LLM = ["SINDI01-DLLM105"]
BPMS = ["SINDI01-DBPM060:Q1", "SINDI02-DBPM010:Q1"]
ATTENUATORS = ["SINDI02-DBLM084:M06-1-ATT2-VAL", "SINDI02-DBLM084:M06-2-ATT2-VAL", "S10DI01-DBLM113:M06-1-ATT2-VAL"] #,"SINDI02-DBLM084:M06-0-ATT2-VAL"]
ATT_SP = 15.0 #3.0,6.0,9.0,12.0,15.0
RANGE = [0.5, 1.1]
STEP_SIZE = 0.05
SETTLING_TIME = 3
SIMULATION = False
do_elog = True


gain_positioners = []
for i in range(len(GAINS)):
   gain_positioners.append( DummyPositioner("gain " + str(i+1)) if SIMULATION else Channel(GAINS[i], alias = "gain " + str(i+1)))
    
attenuators = []
for i in range(len(ATTENUATORS)):
    att = ChannelDouble("Att"+str(i+1), ATTENUATORS[i])
    att.initialize()
    attenuators.append(att)

#Channel-based
#blm1 = ChannelDouble("blm1", "SINDI02-DBLM025:B1_LOSS"); blm1.setMonitored(True); blm1.initialize()
#blm2 = ChannelDouble("blm2", "SINDI02-DBLM085:B1_LOSS"); blm2.setMonitored(True); blm2.initialize()

#Stream creation
for att in attenuators:
     att.write(ATT_SP)
     
sensors = []
line_plots = []
st = Stream("pulse_id", dispatcher)
st.setFilter("SIN-CVME-TIFGUN-EVR0:BEAMOK == 1")
for i in range(len(BLMS)):
    blm = Scalar("blm" + str(i+1), st, BLMS[i] + ":B1_LOSS", 1, 0)
    av = create_averager(blm, SAMPLES, interval = -1)
    av.setMonitored(i>0)
    sensors.append(av)  
    sensors.append(av.stdev)  
    sensors.append(av.samples)  
    line_plots.append(av.samples)  
    blmw = Waveform("blmw" + str(i+1), st, BLMS[i] + ":LOSS_SIGNAL_RAW", 1, 0)
    sensors.append(blmw)

#for i in range(len(LLMS)):
#    llm = Scalar("llm" + str(i+1), st, LLMS[i] + ":XX_LOSS", 1, 0)
#    avl = create_averager(llm, SAMPLES, interval = -1)
#    avl.setMonitored(i>0)
#    sensors.append(avl)  
#    sensors.append(avl.stdev)  
#    sensors.append(avl.samples)  
#    line_plots.append(avl.samples)  

    
for i in range(len(BPMS)):
    bpm = Scalar("bpm" + str(i+1), st, BPMS[i], 1, 0)
    av1 = create_averager(bpm, SAMPLES, interval = -1)
    av1.setMonitored((i>0) or (len(BLMS)>0))
    sensors.append(av1)  
    sensors.append(av1.stdev)  
    sensors.append(av1.samples)  
    line_plots.append(av1.samples)      

#Scalar("beam_ok" , st, "SIN-CVME-TIFGUN-EVR0:BEAMOK" , 1, 0)
st.initialize()
st.start()
st.waitCacheChange(5000)    #Wait stream be running before starting scan


"""
#Averaging
ablm1 = create_averager(blm1, SAMPLES, interval = -1)
ablm2 = create_averager(blm2, SAMPLES, interval = -1)
ablm3 = create_averager(blm3, SAMPLES, interval = -1)
ablm2.setMonitored(True)
ablm3.setMonitored(True)
"""

#Plot setup
setup_plotting( line_plots = line_plots)



#Metadata
set_attribute("/", "Samples" , SAMPLES)
set_attribute("/", "StepSize" , STEP_SIZE)
set_attribute("/", "SamplingTime" , SETTLING_TIME)
set_attribute("/", "Range" , RANGE)
set_attribute("/", "BLM" , BLMS)
#set_attribute("/", "BLMW" , BLMW)
set_attribute("/", "Gain" , GAINS)
set_attribute("/", "Attenuators" , ATTENUATORS)
set_attribute("/", "BPM" , BPMS)
for att in attenuators:
    set_attribute("/", att.channelName, att.read())

try:
    r=lscan(gain_positioners, sensors, [RANGE[0],] * len(gain_positioners), [RANGE[1],] * len(gain_positioners), [STEP_SIZE,] * len(gain_positioners), latency = SETTLING_TIME)
     #r=bscan(st, 10)
finally:
    st.close()    


# save the entry in the logbook
if do_elog:
    if get_option("Generated data file:\n" + get_exec_pars().path + "\n\n" + "Save to ELOG?", "YesNo") == "Yes":    
        gsa_log_msg = "Data file: " + get_exec_pars().path
        gsa_log_msg = gsa_log_msg + "\nSamples = " + str(SAMPLES)
        gsa_log_msg = gsa_log_msg + "\nStepSize = " + str(STEP_SIZE)
        gsa_log_msg = gsa_log_msg + "\nSamplingTime = " + str(SETTLING_TIME)
        gsa_log_msg = gsa_log_msg + "\nRange = " + str(RANGE)
        gsa_log_msg = gsa_log_msg + "\nBLM = " + str(BLMS)        
        #gsa_log_msg = gsa_log_msg + "\nBLMW = " + str(BLMW)  
        gsa_log_msg = gsa_log_msg + "\nBPM = " + str(BPMS)
        gsa_log_msg = gsa_log_msg + "\nGain = " + str(GAINS)
        gsa_log_msg = gsa_log_msg + "\nAttenuators:" # " + str(ATTENUATORS)
        for att in attenuators:
            gsa_log_msg = gsa_log_msg + "\n    " + att.setpoint.channelName + " = " + str(att.read())

        plots = get_plot_snapshots()
        save=[]
        for i in range(len(BLMS)):
            save.append(plots[3*i])
        for i in range(len(BPMS)):
            save.append(plots[3*i])
        elog("BLM Gain scan", gsa_log_msg, save)