import numpy as np
import keithley as ku
import nt252 as nt
import time

NPULSES = 500
NUM_SAMPLES = 10

laser = nt.LaserSerialInterface()
keith = ku.Keithley(resnum=0)

#wavelength_grid = np.arange(650, 1205, 5)
wavelength_grid = np.arange(518, 521, 1)
max_charge = np.zeros((wavelength_grid.size, NUM_SAMPLES))

laser.set_wavelength(wavelength_grid[0])
print('Setting laser to starting wavelength...')
time.sleep(10)


for i,w in enumerate(wavelength_grid):
    print('Moving to wavelength {}'.format(w))
    laser.set_wavelength(w)
    for n in range(NUM_SAMPLES):
        t, q = keith.get_charge_timeseries(npulses=NPULSES)
        max_charge[i, n] = q[-1]
        print('Sample {} Charge: {:6e}'.format(n, q[-1]))

arr = np.column_stack((wavelength_grid, max_charge))
#np.savetxt('Laser_charge_500pulses_10avg.txt', arr)
np.savetxt('my_first_CBP_test.txt', arr)