"""
Arguments:

"""

MOTORS = [Eph]
POSITIONS = [600.0, 610.0]
LATENCY = 0.0

REGION1 = {'name': 'Peak 1', 'elo': 243.0, 'ehi': 253.0, 'estep': 0.1, 'epass': 50, 'tstep': 0.1, 'iter': 1, 'cis': True}
REGION2 = {'name': 'Peak 2', 'elo': 214.0, 'ehi': 219.0, 'estep': 0.1, 'epass': 50, 'tstep': 0.1, 'iter': 2, 'cis': False}
REGIONS = [REGION1, REGION2]

CLOSE_SHUTTER_AT_END = True

Scienta.setAcquisitionMode(ch.psi.pshell.epics.Scienta.AcquisitionMode.Swept)

class SpectrumReader(ReadonlyRegisterBase, ReadonlyRegisterArray):
    def initialize(self):
        #super(SpectrumReader, self).initialize()
        self.scan_index = -1

    def create_datasets(self):
        path = get_exec_pars().scanPath + self.region_name + "/"

        self.channel_begin_dataset_name = path + "ScientaChannelBegin"
        self.channel_end_dataset_name = path + "ScientaChannelEnd"
        self.pass_energy_dataset_name = path + "ScientaPassEnergy"
        self.step_energy_dataset_name = path + "ScientaStepEnergy"
        self.step_time_dataset_name = path + "ScientaStepTime"
        self.iterations_dataset_name = path + "ScientaIterations"
        
        create_dataset(self.channel_begin_dataset_name, 'd')
        create_dataset(self.channel_end_dataset_name, 'd')
        create_dataset(self.pass_energy_dataset_name, 'd')
        create_dataset(self.step_energy_dataset_name, 'd')
        create_dataset(self.step_time_dataset_name, 'd')
        create_dataset(self.iterations_dataset_name, 'd')
        
    def setup(self):
        if self.scan_index != get_exec_pars().index:
            self.scan_index = get_exec_pars().index
            self.create_datasets()
            if self.region_index == 0:
                print "scan {0}".format(self.scan_index)
            
        ephot = Eph.read()
        try:
            if self.region['cis']:
                edelta = ephot - self.ephot_start
            else:
                edelta = 0.0
        except AttributeError:
            self.ephot_start = ephot
            edelta = 0.0
            
        elo = self.region['elo'] + edelta
        ehi = self.region['ehi'] + edelta

        Scienta.lowEnergy.write(elo)
        Scienta.highEnergy.write(ehi)
        Scienta.stepSize.write(self.region['estep'])
        Scienta.setPassEnergy(self.region['epass'])
        Scienta.stepTime.write(self.region['tstep'])
        Scienta.setIterations(self.region['iter'])
        Scienta.update()

        append_dataset(self.channel_begin_dataset_name, elo)
        append_dataset(self.channel_end_dataset_name, ehi)
        append_dataset(self.pass_energy_dataset_name, self.region['epass'])
        append_dataset(self.step_energy_dataset_name, self.region['estep'])
        append_dataset(self.step_time_dataset_name, self.region['tstep'])
        append_dataset(self.iterations_dataset_name, self.region['iter'])

    def read(self):
        global current_region_index
        current_region_index = self.region_index
        self.setup()
        print "Photon energy {ephot} eV. Acquiring spectrum {elo}-{ehi} eV.".format(ephot=Eph.read(), elo=Scienta.lowEnergy.value, ehi=Scienta.highEnergy.value)
        trig_scienta()
        time.sleep(0.5)
        sp = Scienta.getSpectrum().read()
        return sp

    def getSize(self):
        nx = int((self.region['ehi'] - self.region['elo']) / self.region['estep']) + 1
        return nx


class ImageReader(ReadonlyRegisterBase, ReadonlyRegisterMatrix):   
    def read(self):
        return Scienta.getDataMatrix().read()

    def getWidth(self):
        nx = int((self.region['ehi'] - self.region['elo']) / self.region['estep']) + 1
        return nx

    def getHeight(self):
        ny = Scienta.slices.read()
        return ny

def setup_live_plots(regions):
    global live_plots
    global current_region_index
    names = [region['name'] for region in regions]
    live_plots = plot(None, names, title="Live Spectra")
    current_region_index = 0

def update_live_plots():
    global live_plots
    global current_region_index
    try:
        while get_context().state.running:
            y = Scienta.spectrum.take(100)
            x = Scienta.spectrumX
            try:
                series = live_plots[current_region_index].getSeries(0)
                series.setData(x, y)        
            except IndexError:
                pass
            time.sleep(1.0)
    finally:
        print "Stopping live spectra"

def do_scan(motors, positions, regions, latency):
    global SENSORS
    
    SENSORS = []
    
    for (index, region) in enumerate(regions):
        reader = SpectrumReader()
        reader.region_index = index
        reader.region_name = "region{0}".format(index + 1)
        reader.region = region
        reader.initialize()
        set_device_alias(reader, reader.region_name + "/ScientaSpectrum")
        SENSORS.append(reader)

        image = ImageReader()
        image.region_index = index
        image.region = region
        image.initialize()
        set_device_alias(image, reader.region_name + "/ScientaImage")
        SENSORS.append(image)

    SENSORS.append(SampleCurrent)
    SENSORS.append(RefCurrent)
    adjust_sensors()
    set_adc_averaging()

    vscan(motors, SENSORS, positions, True, latency,False, before_read=wait_beam, after_read = after_readout)

    for (index, region) in enumerate(regions):
        set_attribute(get_exec_pars().scanPath + "region{0}/ScientaSpectrum".format(index + 1), "RegionName", region['name'])
        set_attribute(get_exec_pars().scanPath + "region{0}/ScientaImage".format(index + 1), "RegionName", region['name'])
    set_attribute(get_exec_pars().scanPath, "Regions", [region['name'] for region in regions])

try:
    setup_live_plots(REGIONS)
    task = fork(update_live_plots)
    do_scan(MOTORS, POSITIONS, REGIONS, LATENCY)
finally:
    if CLOSE_SHUTTER_AT_END:
        after_scan()