#!/usr/bin/env python3

import pyvisa
import time

import pyaudio as pa
import numpy as np

phase = 0.0 # in cycles
scaled_frequency = 0.0
# actual frequency = scaled_frequency * sample_rate
# For repeatability, scaled_frequency should ideally fit in
# much fewer than 53 bits.
# More precisely, if scaled_frequency * frame_count is always exact, we get
# negligible phase errors over 1hr.

def audio_callback(in_data, frame_count, time_info, status):
	global phase
	L = np.sin(2 * np.pi * (phase + np.arange(frame_count) * scaled_frequency)).astype(np.float32) # cast to float
	delta_phase  = scaled_frequency * frame_count # exact
	delta_phase -= np.floor(delta_phase) # exact
	phase += delta_phase # error at most ulp(1) = 2^-52
	if(phase >= 1):
		phase -= int(phase)
	outbytes = L.tobytes()
	#print(L, frame_count)
	return (outbytes, pa.paContinue)

paudio = pa.PyAudio()
stream = paudio.open(format=pa.paFloat32,
	channels=1,
	rate=48000,
	output=True,
	stream_callback=audio_callback)

def do_frequency(scaled_f, osci, hscale):
	scaled_frequency = scaled_f
	time.sleep(0.1)
	osci.write("acquire:state run") # Restart acquisition sequence
	time.sleep(10 * hscale)

	osci.write("data:source CH1")
	A_samps = osci.query_binary_values("curve?",datatype='b',container=numpy.array)
	A_samps = A_samps.astype(np.float64)

	osci.write("data:source CH2") # For now, do not adjust for zeros in the resp.
	B_samps = osci.query_binary_values("curve?",datatype='b',container=numpy.array)
	B_samps = B_samps.astype(np.float64)

	A_samps = np.rfft(A_samps) # No window, because of the precise frequency
	B_samps = np.rfft(B_samps)

	# We have (10 * hscale)/Tcurv periods
	return B_samps[np.round((10 * hscale) * (scaled_f * 4800))]
        /A_samps[np.round((10 * hscale) * (scaled_f * 4800))]

def find_oscillo():
	rm = pyvisa.ResourceManager()
	R = rm.list_resources("USB?*")
	if R.empty():
		print("No device found")
		exit(1)
	return rm.open_resource(R[0])

osci = find_oscillo()
print(osci.query("*IDN?"))

osci.write("*RST")
time.sleep(5)
osci.write("horizontal:main:scale 5.0e-4")

osci.write("ch1:scale 1.0e-1")
osci.write("ch1:bandwidth ON")
osci.write("ch1:coupling AC")
osci.write("select:ch1 1")

osci.write("ch2:scale 5.0e-2")
osci.write("ch2:bandwidth ON")
osci.write("ch2:coupling AC")
osci.write("select:ch2 1")

# 8-bit data
osci.write("data:source CH1")
osci.write("data:encdg ascii")

osci.write("trigger:main:edge:source CH1")
osci.write("trigger:main:mode NORMAL")
osci.write("trigger:main:type EDGE")

osci.write("acquire:mode average")
osci.write("acquire:numavg 128")
osci.write("acquire:stopafter sequence")
osci.write("acquire:state RUN")

print(osci.query("curve?"))