###################################################################################################
#  Deployment specific global definitions - executed after startup.py
###################################################################################################
import traceback

energy.setBlockingWrite(True)


#The persistence strategy is overriden to adjust file names
import ch.psi.pshell.data.LayoutTable
class DataLayout( ch.psi.pshell.data.LayoutTable):
    def getLogFileName(self):
        return time.strftime('%Y%m%d_%H%M') + '_' + get_context().name + '_logs'
    
    def getDatasetName(self, scan):
        print get_context().count
        data_file =  time.strftime('%Y%m%d_%H%M') + '_' + get_context().name  + '_' + str(get_context().count).zfill(4)        
        print "Opened data file: " + get_context().path + "/" + data_file
        return data_file
controller.dataManager.setLayout(DataLayout())


#Reading Energy Scan configuration file
def getPars(element):
    f = open(controller.setup.getConfigPath() + '/energy_scan.properties')
    try:
        for line in f:
            tokens = line.split("=")
            if tokens[0] == str(element):
                tokens = tokens[1].split(";")
                for i in range(len(tokens)):
                    tokens[i] = float(tokens[i].strip())

                return tokens
    finally:
        f.close()
    raise Exception ("Invalid element: " + str(element))


###################################################################################################
# PSeudo-devices
###################################################################################################
sim_energy = None
class SimulatedEnergy(Writable):
    def write(self, value):
        self.put(value)

    def put(self, value, timeout = None):
        global sim_energy
        sim_energy = value

    def close(self):
        pass


class SimulatedEnergyReadback(Readable):
    def read(self):
        global sim_energy
        return sim_energy;
    
    def get(self):
        return self.read()

    def close(self):
        pass

sim_energy = SimulatedEnergy()
sim_energy_readback = SimulatedEnergyReadback()



class TeyNorm(ReadonlyRegisterBase):
    def read(self):
        return signal_tey.take() / signal_i0.take();
    
class TransNorm(ReadonlyRegisterBase):
    def read(self):
        return signal_trans.take() / signal_i0.take();

add_device(TeyNorm("tey_norm"), True)
add_device(TransNorm("trans_norm"), True)


###################################################################################################
# Utilities
###################################################################################################

run ("Accumulator")

def get_next_fid(folder, prefix):
    try:
        import glob   
        files = glob.glob(folder + prefix + '*_*.txt')
        last = max(files)
        index = int (last[last.rfind('_')+1 : last.rfind('.')]) + 1
        return index
    except:
        return 0

def wait_channel(name, value, timeout=None, type='s'):
    print "Waiting " + str(name) + " = " + str(value)
    cawait(name, value, timeout = timeout, type=type)
    print "Done waiting"

def wait_device(dev, value, timeout=-1):
    timeout = int(timeout *1000) if timeout>0 else timeout
    print "Waiting " + dev.getName() + " = " + str(value)
    dev.waitValue(value,timeout)
    print "Done waiting"
        
def convert_file(input_file_name, output_file_name, pol = None):   
    print "Converting data file: " + input_file_name + " to " + output_file_name
    sep = "\t"
    line_sep = "\n"
    if pol is None:
        pol = pol_angle.read() if (MODE == "LINEAR") else pol_mode.readback.read()      
    with open(input_file_name) as inp:
        lines = inp.readlines()
        with open(output_file_name, "wb") as out:            
            out.write("Ecrbk" + sep + "CADC1" + sep + "CADC2" + sep + "CADC3" + sep + "CADC4" + sep + "CADC5" + sep + "MCurr" + sep + "Time" + sep + "FieldX" + sep + "FieldZ" + sep + "Pol" + sep + "Temperature" + sep + "NORMtey" + sep + "NORMdiode" + line_sep)
            s = sep
            for line in lines[1:]:
                line = line.strip()
                if line=="": break
                try:              
                    (Ecrbk, CADC1, CADC2, CADC3, CADC4, CADC5, MCurr, cffrbk, IDErbk, NORM, time, MAGX, MAGZ, EXPT) = line.split(" ")                   
                    normtey=repr( float(CADC1)/float(CADC2))
                    normdiode=repr(float(CADC3)/float(CADC2))
                    out.write(Ecrbk + s + CADC1 + s + CADC2 + s + CADC3 + s + CADC4 + s + CADC5 + s + MCurr + s + time + s +  MAGX + s +  MAGZ + s + str(pol) + s + EXPT + s + normtey + s + normdiode + line_sep)
                except:
                    traceback.print_exc()  
                    
def plot_file(file_name, title = None):
    """
    """
    table = Table.load(file_name, "\t", '#')
    plots = plot(table, title = title)
                        
    

###################################################################################################
# Controller state listener
###################################################################################################

#class Listener(ControllerListener):
#    def onControllerStateChanged(self, state, former):
#        global FILENAME
#        if former == State.Busy:
#            FILENAME = None
#controller.addListener(Listener())    #Old listeners are cleande in startup


###################################################################################################
#Definitions for importing text batch files
###################################################################################################
#TODO: Should set devices? K10, k11, k24 seem not to be defined....
keithleys = {
    #name:[Setpoint,                 range,                      Readback] #TODO
    "k1": [None ,                     'X07MA-PC-K428:1:setGain', 'X07MA-PC-K428:1:getGain'],
    "k2": [None ,                     'X07MA-PC-K428:2:setGain', 'X07MA-PC-K428:2:getGain'],
    "k3": [None ,                     'X07MA-PC-K428:3:setGain', 'X07MA-PC-K428:3:getGain'],
    "k10":['X07MA-KEI10:SETVOLTAGE' , 'X07MA-KEI10:RANGE',       None],
    "k11":['X07MA-KEI11:SETVOLTAGE' , 'X07MA-KEI11:RANGE',       None],
    "k24":['X07MA-KEI2400:setVoltAO', None,                      None]
}

def set_hx(field, timeout = -1):
    """
    """
    timeout = int(timeout *1000) if timeout>0 else timeout
    field_x.write(float(field))
    field_x_done.waitValue(1,timeout)
#Only test x today
def setHz(field, timeout = -1):
    """
    """
    timeout = int(timeout *1000) if timeout>0 else timeout
    field_z.write(float(field))
    field_z_done.waitValue(1,timeout)

def set_pol_cplus(offset = None, timeout = -1):
    """
    """
    timeout = int(timeout *1000) if timeout>0 else timeout
    pol_mode.write("CIRC +")
    if offset is not None:
        pol_offset.write(float(offset))
    pol_done.waitValue("DONE",timeout)    

def set_pol_cminus(offset = None, timeout = -1):
    """
    """
    timeout = int(timeout *1000) if timeout>0 else timeout
    pol_mode.write("CIRC -")
    if offset is not None:
        pol_offset.write(float(offset))
    pol_done.waitValue("DONE",timeout)    

def set_pol_lin(angle, offset = None, timeout = -1):
    """
    """
    timeout = int(timeout *1000) if timeout>0 else timeout
    pol_mode.write("LINEAR")
    pol_angle.write(float(angle))
    if offset is not None:
        pol_offset.write(float(offset))
    pol_done.waitValue("DONE",timeout)    

def set_temp(value, delay = 30.0):    
    """
    """
    temperature.write(float(value))
    time.sleep(float(delay))
    
def open_valve(delay = 5.0):    
    """
    """
    valve_try_open.write(1)    #TODO: CAPUT
    time.sleep(float(delay))

def close_valve(delay = 0.0):    #TODO: Check default delay
    """
    """
    valve_try_open.write(0)
    time.sleep(fset_pol_cplusloat(delay))

def set_energy (value, delay=0.5):
    """
    """
    energy.write(float(value))
    energy_done.waitValue(1, -1)
    time.sleep(float(delay))

def set_cff (value):
    """
    """
    cff.write(float(value))
    caput("X07MA-PGM:setE.PROC",1)
    energy_done.waitValue(1, -1)

def set_slit(value):
    """
    """
    exit_slit.write(float(value))
    
#1 or 3
def set_har(value):
    """
    """
    harmonic.write(float(value))

#Not connected
def set_volt(keithley, value, delay=0.1):
    """
    """
    #keithley.write(fl'10^'+str(8)oat(value))
    if not keithltey in keithleys.keys(): raise Exception("Invalid keithley: " + keithley)
    caput(keithleys[keithley][0], float(value))
    time.sharmonicleep(float(delay))

#value is int from 1 to 10
def set_range(keithley, value):
    """
    """
    if not keithley in keithleys.keys(): raise Exception("Invalid keithley: " + keithley)
    v='10^'+str(value)    
    while True:
        caput(keithleys[keithley][1], v)
        time.sleep(1.0)
        if caget(keithleys[keithley][2],'s') == v: 
            break

def set_fe(opening):
    """
    """
    opening = int(opening*1000)
    if opening==0:
        aperture.write("closed")
    elif opening==100:
        aperture.write("0.1x0.1 mm")
    elif opening==250:
        aperture.write("0.25x0.25 mm")
    elif opening==500:
        aperture.write("0.5x0.5 mm")
    elif opening==1000:
        aperture.write("1x1 mm")
    elif opening==1250:
        aperture.write("1.25x1.25 mm")
    elif opening==2000:
        aperture.write("2x2 mm")
    else:
        raise Exception("Invalid aperture opening: " + str(opening))

def refill_1k():
    """
    """
    run("refill_1kpot")
    
def wait_temp():
    """
    """
    if (temperature.read() < 2.0):
        t=ct=temperature_readback.read()
        #hl=caget("X07MA-PC-HE:LEVELB", 'd') # TODO: not used	 
        if (t > 3.7):
            refill_1k()

def rampdown():
    """
    """
    field_x.write(0.0)
    field_z.write(0.0)
    
def shutdown():
    """
    """
    set_fe(0)
    rampdown()


def has_beam():
    """
    """
    return beam_status.readback.read() !="Machine Down"

def wait_beam():
    """
    """
    print "Waiting for beam..."
    while not has_beam():
        sleep(0.1)
    print "Beam OK"

def wait_inj (value, delay=0.5):
    """
    """
    wait_channel('X07MA-OP2-EVR:TOPUP-STAT', 'TOPUP-ON')
    wait_channel('X07MA-OP2-EVR:TOPUP-STAT', 'TOPUP-OFF') 
    time.sleep(float(delay))
        

def set_file(file_name):
    """
    """
    set_context(name = file_name)

class pol_mod(Readable):
    def read(self):  
        mode = pol_mode.read()
        if mode == "LINEAR": return 0.0
        if mode == "CIRC +": return 1.0
        if mode == "CIRC -": return 2.0
        return -1.0
polmod = pol_mod()
        
    
def otf(start, end, time, delay=0.0, mode = None, offset = None, alpha = None, name = None):
    """
    """    
    if name is None:
        name = get_context().name
    folder = controller.setup.expandPath("{year}_{month}/{date}");    
    run("EnergyScan", {"E1":start, "E2":end, "TIME":time, "DELAY":float(delay), "MODE":mode, "OFFSET":(offset), "FOLDER":folder, "FILE":name, "ALPHA":alpha})    

def hyst_cont(field, init_field, final_field, ramp_speed, energies):
    """
    """
    run("HystScan",{"FIELD":field, "START_FIELD":init_field, "END_FIELD":final_field, "ENERGIES":energies, "RAMP_RATE":ramp_speed, "ENERGY_CHANGE_SLEEP":0.5, "MODE":None, "OFFSET":None})

def hyst_cont_mult(field, ranges, energies):
    """
    """
    run("HystScanMult",{"FIELD":field, "RANGES":ranges, "ENERGIES":energies, "ENERGY_CHANGE_SLEEP":0.5, "MODE":None, "OFFSET":None})

def hyst_step(forward, field,  init_field, final_field, step, energies, energy_change_sleep = 0.5, field_change_sleep = 22.0, mode = None, offset=None):
    """
    """
    run("HystScanStep",{"FIELD":field, "RANGES":[(init_field, final_field, step),], "ENERGIES":energies, "ENERGY_CHANGE_SLEEP":energy_change_sleep, "FIELD_CHANGE_SLEEP":field_change_sleep, "MODE":mode, "OFFSET":offset})

def hyst_step_mult(forward, ranges, ramp_speed, energies, energy_change_sleep = 0.5, field_change_sleep = 22.0, mode = None, offset=None):
    """
    """
    run("HystScanStep",{"FIELD":field, "RANGES":ranges, "ENERGIES":energies, "ENERGY_CHANGE_SLEEP":energy_change_sleep, "FIELD_CHANGE_SLEEP":field_change_sleep, "MODE":mode, "OFFSET":offset})

def scan_e(start, end, step, settling_time = 0, accumulation_time = None, name = None):
    """
    """
    if name is not None:
        set_context(name = name)
    #wait_beam()    
    acc = Accumulator([signal_tey, signal_i0, signal_trans, tey_norm, trans_norm], accumulation_time)        
    detectors = acc.getSensors() + [polmod, pol_angle, temperature, current]
    set_preference(Preference.ENABLED_PLOTS, acc.getSensors())
    lscan(energy, detectors, float(start), float(end), float(step), latency = settling_time)

def scan_e_mult(ranges, settling_time = 0, accumulation_time = None, name = None):
    """
    """
    if name is not None:
        set_context(name = name)        
    acc = Accumulator([signal_tey, signal_i0, signal_trans, tey_norm, trans_norm], accumulation_time)
    detectors = acc.getSensors() + [polmod, pol_angle, temperature, current]
    set_preference(Preference.ENABLED_PLOTS, acc.getSensors())
    rscan(energy, detectors, ranges, latency = settling_time)
    
#not connected
def scan_v(keithley,start, end, step):
    """
    """
    setpoint = Channel (keithleys[keithley][0], 'd')
    readback = Channel (keithleys[keithley][2], 'd')
    lscan(setpoint, readback, float(start), float(end), float(step), latency = 0.1)