from setup_cbp import *
import cbpassist as ca
import time
import pyds9

#ds9 = pyds9.DS9()

def show(f):
    if True:
        fid = pyfits.open(f)
        ds9.set_pyfits(fid)
        fid.close()

def move_shoot(d, v, h, outdir):
    d[v] = d[v] + h
    f=sched.shoot(d, outdir)
#    show(f)

d['sbig.EXPTIME'] = 1
d['sbig.filter'] = 'EMPTY'
d['mount.alpha'] = sched.mount.get_alpha()
d['mount.delta'] = sched.mount.get_delta()
d['cbpmount.alt'] = sched.cbpmount.get_alt()
d['cbpmount.az'] = sched.cbpmount.get_az()
d['keithley'] = False
d['spectro'] = False
d['laser.wavelength'] = 450
d['cbpwheel.pinhole'] = 'grid'
d['cbpwheel.mask'] = 2
d['ndfilter.pos'] = 'IDLE'


outdir = 'mondrik_temp_test/focal_plane_scan_cbpaz'

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

wl_grid = np.array([400, 450, 550, 650, 750, 850,  950, 1050])

e,_ = ca.calc_exp_time(wl_grid)
t_expected = 2.*((np.sum(e) + 10.*len(e)) / 60.)
print('Approximate time for scan: {} min'.format(t_expected))
print('This might be made longer by frequent ND filter rotations')

focal_plane_shifts = -np.ones(16).astype(np.int) # - dir for az
focal_plane_shifts[0] = 0.
#shift_key = 'cbpmount.alt'
shift_key = 'cbpmount.az'
total_shift = 0.

for f in focal_plane_shifts:
    total_shift += f
    print('SHIFTING BY {}, TOTAL SHIFT: {}'.format(f, total_shift))
    d[shift_key] += int(f)
    for w in wl_grid:
        exptime, filter_pos = ca.calc_exp_time(w)
        d['sbig.EXPTIME'] = exptime+10.
        d['keithley.exptime'] = str(exptime)
        d['laser.wavelength'] = str(w)
        d['ndfilter.pos'] = str(filter_pos)
        d['spectro.exptime'] = 1.
    
        #  Do dark
        d['keithley'] = False
        print('Taking Dark')
        sched.shoot(d, outdir)
        d['keithley'] = True
        print('Taking Light')
        sched.shoot(d, outdir)