"""
Arguments:

"""

LATENCY = 0.0

PHI_RANGE = (0, 160.0) # (tuple (min, max))
THETA_RANGE = (-9, 81) #  (tuple (min, max))
STEPS = (40.0, 1.0) # (tuple (phi, theta))
ZIGZAG = True

REGION1 = {'name': 'Pt4f', 'elo': 1016, 'ehi': 1028, 'estep': 0.2, 'epass': 100, 'tstep': 0.05, 'iter': 1, 'cis': True, 'slit': 40}
REGION2 = {'name': 'B1s', 'elo': 901, 'ehi': 909, 'estep': 0.2, 'epass': 100, 'tstep': 0.5, 'iter': 1, 'cis': True, 'slit': 210}
REGION3 = {'name': 'C1s', 'elo': 808, 'ehi': 813, 'estep': 0.2, 'epass': 100, 'tstep': 0.5, 'iter': 1, 'cis': True, 'slit': 210}
REGION4 = {'name': 'N1s', 'elo': 694, 'ehi': 701, 'estep': 0.2, 'epass': 100, 'tstep': 0.5, 'iter': 1, 'cis': True, 'slit': 210}
REGIONS = [REGION1, REGION2, REGION3, REGION4]

CLOSE_SHUTTER_AT_END = True

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

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"
        self.slit_dataset_name = path + "ExitSlit"
        
        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')
        create_dataset(self.slit_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)
            
        elo = self.region['elo']
        ehi = self.region['ehi']

        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'])
        ExitSlit.write(self.region['slit'])
        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'])
        append_dataset(self.slit_dataset_name, self.region['slit'])

    def read(self):
        global current_region_index
        current_region_index = self.region_index
        self.setup()
        print "Acquiring spectrum {elo}-{ehi} eV.".format(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(regions, latency):
    global SENSORS, PHI_RANGE, THETA_RANGE
    
    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()

    ascan((ManipulatorPhi, ManipulatorTheta), SENSORS, (PHI_RANGE[0], THETA_RANGE[0]), (PHI_RANGE[1], THETA_RANGE[1]), STEPS, LATENCY, False, zigzag = ZIGZAG, 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(REGIONS, LATENCY)
finally:
    if CLOSE_SHUTTER_AT_END:
        after_scan()