import scheduler
from stardice import configfiles
from glob import glob
import astropy.io.fits as pyfits
import logging
logging.basicConfig(format='%(asctime)s %(message)s', level=logging.DEBUG)
import numpy as np
from notification import notify
import os
import pickle
outdir = '../data/spectro_bench/zwo_qe_32'

if not os.path.exists(outdir):
    os.mkdir(outdir)

dkd_pos = 145000
zwo_pos = 92140

config = configfiles.Config('spectro_bench.cfg')

sched = scheduler.Scheduler(config)


d = sched.get_default()

#leds = np.load('../drivers/ardice/led16_1.npy')

d['zwo.exptime'] = 1.0
d['phd.range'] = 2e-4
d['phd.nsamples'] = 10
d['dkd.range'] = 2e-10
d['dkd.nsamples'] = 10

with open('led_32.pkl', 'rb') as fid:
    wavelengthes = pickle.load(fid)
    currents = pickle.load(fid)
    cwls = pickle.load(fid)

leds = np.array(list(cwls))
wls = np.array([cwls[led] for led in leds]).squeeze()
led_order = np.argsort(wls)

notify('wavelength scan started')
try:
    for led in leds[led_order]:
        notify(f'Starting led {led}')
        d['xy.x'] = int(dkd_pos)
        d['ar.led'] = int(led)
        d['ar.iled'] = int(currents[led])
        d['mono.grating'] = 3  if cwls[led] > 600 else 2
        d['ar.shutter'] = 'close'
        d['zwo'] = False
        # on phd
        for i in range(30):
            sched.shoot(d, outdir)
        d['ar.shutter'] = 'open'
        for wavelen in wavelengthes[led]:
            d['mono.wavelength'] = int(wavelen)
            sched.shoot(d, outdir)
            sched.shoot(d, outdir)
        # on camera
        d['xy.x'] = int(zwo_pos)
        d['ar.shutter'] = 'close'
        d['zwo'] = True
        for i in range(10):
            sched.shoot(d, outdir)
        d['ar.shutter'] = 'open'            
        for wavelen in wavelengthes[led]:
            d['mono.wavelength'] = int(wavelen)
            sched.shoot(d, outdir)
            sched.shoot(d, outdir)
        d['ar.shutter'] = 'close'
        for i in range(10):
            sched.shoot(d, outdir)
        #on phd
        d['ar.shutter'] = 'open'
        d['xy.x'] = int(dkd_pos)
        d['zwo'] = False
        for wavelen in wavelengthes[led]:
            d['mono.wavelength'] = int(wavelen)
            sched.shoot(d, outdir)
            sched.shoot(d, outdir)
        d['ar.shutter'] = 'close'
        for i in range(10):
            sched.shoot(d, outdir)
finally:
    notify('Wavelength scan stopped')
        
#wl = [pyfits.getheader(i)['wavelength'] for i in l[-32:]]
#g = [pyfits.getheader(i)['grating'] for i in l[-32:]]
#f = [pyfits.getdata(i).sum() for i in l[-32:]]