#!/usr/bin/python
import datetime
import optparse
import os
import numpy
import math
import scipy.integrate
import sys
import time
from vesna.spectrumsensor import SpectrumSensor, SweepConfig
from vesna.rftest import DeviceUnderTest, SignalGenerator, parse_test_kwargs

log_f = None

def log(msg):
	log_f.write(msg + "\n")
	print msg

class LocalDeviceUnderTest(DeviceUnderTest):
	def add_options(self, parser):
		parser.add_option("-d", "--device", dest="device", metavar="DEVICE",
				help="Use VESNA spectrum sensor attached to serial DEVICE.",
				default="/dev/ttyUSB0")
		parser.add_option("--no-calibration", dest="calibration", action="store_false",
				help="Don't use calibration tables.",
				default=True)

	def setup(self, options):
		self.spectrumsensor = SpectrumSensor(options.device, calibration=options.calibration)

		self.config_list = self.spectrumsensor.get_config_list()
		if not self.config_list.configs:
			raise Exception("Device returned no configurations. "
					"It is still scanning or not responding.")

		self.config = self.config_list.get_config(self.device_id, self.config_id)

	def get_fw_version(self):
		return self.spectrumsensor.get_fw_version()

	def get_status(self):
		resp = self.spectrumsensor.get_status(self.config)
		resp += [ "", "Calibration: %s" % (self.spectrumsensor.calibration,) ]
		return resp

	def measure_ch_impl(self, ch, n):
		sweep_config = SweepConfig(self.config, ch, ch+1, 1)

		measurements = []
		timestamps = []

		def cb(sweep_config, sweep):
			assert len(sweep.data) == 1
			measurements.append(sweep.data[0])
			timestamps.append(sweep.timestamp)
			return len(measurements) < n

		self.spectrumsensor.run(sweep_config, cb)

		return timestamps, measurements

def chop(p_dbm_list, pout_dbm_list, min_dbm, max_dbm):
	p_dbm_list2 = []
	pout_dbm_list2 = []
	for p_dbm, pout_dbm in zip(p_dbm_list, pout_dbm_list):
		if p_dbm >= min_dbm and p_dbm <= max_dbm:
			p_dbm_list2.append(p_dbm)
			pout_dbm_list2.append(pout_dbm)

	return p_dbm_list2, pout_dbm_list2

def max_error(pin_dbm_list, pout_dbm_list):
	return max(	abs(i - o)
			for i, o in zip(pin_dbm_list, pout_dbm_list) )

def get_linear_range(k, n, pin_dbm_list, pout_dbm_list):

	emax = -1
	for i, (pin_dbm, pout_dbm) in enumerate(zip(pin_dbm_list, pout_dbm_list)):

		pout_dbm_lin = pin_dbm * k + n

		e = abs(pout_dbm - pout_dbm_lin)

		if pin_dbm < -40:
			emax = max(emax, e)
		else:
			if e > emax * 1.5:
				break

	log("      saturation @ Pin = %.1f dBm" % (pin_dbm_list[i-1],))

def test_power_ramp(dut, gen):
	"""Check measurement errors versus input power
	"""

	log("Start power ramp test")

	N = 100

	nruns = 3
	ch_num = dut.config.num
	ch_list = [ int(ch_num*(i+0.5)/nruns) for i in xrange(nruns) ]

	for ch in ch_list:
		f_hz = dut.config.ch_to_hz(ch)

		log("  f = %d Hz" % (f_hz))
		
		gen.rf_off()

		nf = dut.measure_ch(ch, N, "power_ramp_%dhz_off" % (f_hz,))
		nf_mean = numpy.mean(nf)

		log("    N = %f dBm, u = %f" % (nf_mean, numpy.std(nf)))

		p_dbm_start = int(round(nf_mean/10)*10-30)
		p_dbm_step = 5

		if p_dbm_start < -85:
			p_dbm_list =	range(p_dbm_start, -85, p_dbm_step) + \
					range(-85, -75, p_dbm_step/5) + \
					range(-75, p_dbm_step, p_dbm_step)
		else:
			p_dbm_list = range(p_dbm_start, p_dbm_step, p_dbm_step)

		pout_dbm_list = []
		for p_dbm in p_dbm_list:
			log("    Pin = %d dBm" % (p_dbm))
			gen.rf_on(f_hz, p_dbm)
			s = dut.measure_ch(ch, N, "power_ramp_%dhz_%ddbm" % (f_hz, p_dbm))
			s_mean = numpy.mean(s)
			log("       Pout = %f dBm, u = %f" % (s_mean, numpy.std(s)))
			pout_dbm_list.append(s_mean)

		gen.rf_off()


		f = open("%s/%s_power_ramp_%dhz.log" % (dut.log_path, dut.name, f_hz,), "w")
		f.write("# Pin [dBm]\tPout [dBm]\n")
		for p_dbm, pout_dbm in zip(p_dbm_list, pout_dbm_list):
			f.write("%f\t%f\n" % (p_dbm, pout_dbm))
		f.close()

		r, m = chop(p_dbm_list, pout_dbm_list, nf_mean+20, 0)
		log("    Range %.1f - %.1f dBm" % (r[0], r[-1]))
		log("      max absolute error %.1f dBm" % (max_error(r, m)))

		r, m = chop(p_dbm_list, pout_dbm_list, nf_mean+20, -40)
		log("    Range %.1f - %.1f dBm" % (r[0], r[-1]))
		log("      max absolute error %.1f dBm" % (max_error(r, m)))

		A = numpy.array([numpy.ones(len(r))])
		n = numpy.linalg.lstsq(A.T, numpy.array(m) - numpy.array(r))[0][0]
		log("      offset = %f dBm" % (n,))

		A = numpy.array([r, numpy.ones(len(r))])
		k, n = numpy.linalg.lstsq(A.T, m)[0]
		log("      linear regression: k = %f, n = %f dBm" % (k, n))

		r, m = chop(p_dbm_list, pout_dbm_list, nf_mean+20, 0)
		get_linear_range(k, n, r, m)

	log("End power ramp test")

def test_freq_sweep(dut, gen, p_dbm_list=None):
	"""Check measurement errors versus tuned frequency
	"""

	log("Start frequency sweep test")

	N = 100

	if p_dbm_list is None:
		p_dbm_list = [ -90, -75, -60, -45 ]

	for p_dbm in p_dbm_list:

		log("  Pin = %d dBm" % (p_dbm,))

		nruns = 50
		ch_num = dut.config.num
		ch_list = [ int((ch_num-1)*i/(nruns-1)) for i in xrange(nruns) ]

		f_hz_list = []
		pout_dbm_list = []
		for ch in ch_list:
			f_hz = dut.config.ch_to_hz(ch)
			f_hz_list.append(f_hz)

			log("    f = %d Hz" % (f_hz))
			gen.rf_on(f_hz, p_dbm)
			s = dut.measure_ch(ch, N, "freq_sweep_%ddbm_%dhz" % (p_dbm, f_hz))
			s_mean = numpy.mean(s)
			log("       Pout = %f dBm, u = %f" % (s_mean, numpy.std(s)))
			pout_dbm_list.append(s_mean)

		gen.rf_off()

		path = ("%s/%s_freq_sweep_%ddbm.log" % (dut.log_path, dut.name, p_dbm)).replace("-", "m")
		f = open(path, "w")
		f.write("# f [Hz]\tPout [dBm]\n")
		for f_hz, pout_dbm in zip(f_hz_list, pout_dbm_list):
			f.write("%f\t%f\n" % (f_hz, pout_dbm))
		f.close()

		log("    Range %.1f - %.1f Hz" % (f_hz_list[0], f_hz_list[-1]))
		log("      max absolute error %.1f dBm" % (max_error([p_dbm]*len(f_hz_list), pout_dbm_list)))

	log("End frequency sweep test")

def get_settle_time(measurements, settled):
	mmax = max(settled)
	mmin = min(settled)

	for n, v in enumerate(measurements):
		if v <= mmax and v >= mmin:
			return n

def test_settle_time(dut, gen):
	"""Measure automatic gain control settling time"""
	
	log("Start settle time test")

	if hasattr(dut.spectrumsensor, "run"):
		N = 1000

		nruns = 3
		ch_num = dut.config.num
		ch_list = [ int(ch_num*(i+0.5)/nruns) for i in xrange(nruns) ]

		p_dbm_list = [ -90, -50, -10 ]
	else:
		log("  Skipping test - cannot be done remotely")
		p_dbm_list = []

	for p_dbm in p_dbm_list:
		log("  Pin = %d dBm" % (p_dbm,))

		for ch in ch_list:
			f_hz = dut.config.ch_to_hz(ch)

			log("    f = %d Hz" % (f_hz,))

			step = 200

			gen.rf_off()

			sweep_config = SweepConfig(dut.config, ch, ch+1, 1)

			timestamps = []
			measurements = []

			def cb(sweep_config, sweep):
				assert len(sweep.data) == 1

				timestamps.append(sweep.timestamp)
				measurements.append(sweep.data[0])

				if len(measurements) == 1*step:
					gen.rf_on(f_hz, p_dbm)
					return True
				elif len(measurements) == 2*step:
					gen.rf_off()
					return True
				elif len(measurements) == 3*step:
					return False
				else:
					return True

			name = "settle_time_%ddbm_%dhz" % (p_dbm, f_hz)
			if dut.is_replay():
				measurements = dut._measure_ch_replay(name)
			else:
				dut.spectrumsensor.run(sweep_config, cb)
				dut._measure_ch_save(name, timestamps, measurements)

			on_settled = measurements[2*step-step/4:2*step]
			t = get_settle_time(measurements[1*step:], on_settled)

			log("      settled up in %d samples" % (t,))
			if t >= dut._extra:
				log("        WARNING: settle time too long for other tests!")

			off_settled = measurements[-step/4:]
			t = get_settle_time(measurements[2*step:], off_settled)

			log("      settled down in %d samples" % (t,))
			if t >= dut._extra:
				log("        WARNING: settle time too long for other tests!")

	log("End settle time test")

def find_zero(i_start, i_step, x, y):
	i = i_start
	try:
		while y[i] > 0:
			i += i_step
	except IndexError:
		return None

	return numpy.interp(0, y[i-1:i+1], x[i-1:i+1])

def get_channel_start_stop(fc_hz, f_hz_list, pout_dbm_list, config):

	pout_dbm_max = max(pout_dbm_list)
	i_max = pout_dbm_list.index(pout_dbm_max)

	pout_diff_list = numpy.array(pout_dbm_list) - pout_dbm_max + 3.0

	f_hz_start = find_zero(i_max, -1, f_hz_list, pout_diff_list)
	f_hz_stop = find_zero(i_max, +1, f_hz_list, pout_diff_list)

	if f_hz_start is None or f_hz_stop is None:
		log("      Can't calculate channel filter pass band!")
	else:
		bw = f_hz_stop - f_hz_start
		fc_hz_real = (f_hz_stop + f_hz_start) / 2.0

		log("    Poutmax = %f dBm" % (pout_dbm_max,))
		log("    fmin = %f Hz" % (f_hz_start,))
		log("    fmax = %f Hz" % (f_hz_stop,))
		log("    BW(-3 dB) = %f Hz (should be %d Hz, %.1f %%)" % (
				bw, config.bw, 100.0 * (config.bw - bw) / bw))

		efc = fc_hz - fc_hz_real
		log("    Efc = %f Hz (%.1f %% channel)" % (
				efc, 100.0 * efc / config.bw))

def get_noise_figure(f_hz_list, pout_dbm_list, p_dbm, nf_mean):
	pout_list = (10 ** ((numpy.array(pout_dbm_list)-p_dbm) / 10)) * 1e-3
	i = 10*math.log10(scipy.integrate.trapz(pout_list, f_hz_list)/1e-3)
	log("    NF = %f" % ((nf_mean - i) + 174))

def test_ch_filter(dut, gen):
	"""Test channel filter, local oscillator accuracy and noise figure
	"""

	log("Start channel filter test")

	N = 100

	nruns = 3
	ch_num = dut.config.num
	ch_list = [ int(ch_num*(i+0.5)/nruns) for i in xrange(nruns) ]

	p_dbm = -30

	log("  Pin = %d dBm" % (p_dbm,))

	for ch in ch_list:
		fc_hz = dut.config.ch_to_hz(ch)

		log("  fc = %d Hz" % (fc_hz,))

		gen.rf_off()

		nf = dut.measure_ch(ch, N, "channel_filter_%dhz_off" % (fc_hz,))
		nf_mean = numpy.mean(nf)

		log("    N = %f dBm, u = %f" % (nf_mean, numpy.std(nf)))

		npoints = 80

		f_hz_list = [	(fc_hz - 1.5*dut.config.bw) + 3.0 * dut.config.bw / (npoints - 1) * n
				for n in xrange(npoints) ]

		pout_dbm_list = []

		for f_hz in f_hz_list:
			log("    f = %d Hz" % (f_hz))
			gen.rf_on(f_hz, p_dbm)
			s = dut.measure_ch(ch, N, "channel_filter_%dhz_%dhz" % (fc_hz, f_hz))
			s_mean = numpy.mean(s)
			log("      Pout = %f dBm, u = %f" % (s_mean, numpy.std(s)))
			pout_dbm_list.append(s_mean)

		gen.rf_off()

		path = ("%s/%s_channel_filter_%dhz.log" % (dut.log_path, dut.name, fc_hz))
		f = open(path, "w")
		f.write("# f [Hz]\tPout [dBm]\n")
		for f_hz, pout_dbm in zip(f_hz_list, pout_dbm_list):
			f.write("%f\t%f\n" % (f_hz, pout_dbm))
		f.close()

		get_channel_start_stop(fc_hz, f_hz_list, pout_dbm_list, dut.config)
		get_noise_figure(f_hz_list, pout_dbm_list, p_dbm, nf_mean)

	log("End channel filter test")

def test_ident(dut, gen, **kwargs):
	"""Identify device under test and testing harness
	"""

	log("Start identification")
	log("  Device under test: %s" % (dut.name,))
	if dut.is_replay():
		log("  *** REPLAY ***")

	log("    Firmware version:")
	log("      %s" % (dut.get_fw_version()))

	log("    Device status:")
	resp = dut.get_status()
	for line in resp:
		log("      %s" % (line.strip(),))

	log("    Device configurations:")
	for line in str(dut.config_list).split('\n'):
		log("      %s" % (line,))

	log("    Using config %d,%d" % (dut.config.device.id, dut.config.id))

	log("  Signal generator: %s" % (gen.get_name(),))
	log("End identification")

def load_dut(options):

	device_id, config_id = map(int, options.vesna_config.split(","))

	dut_opts = options.dut_opts
	if options.dut_name is None:
		if options.vesna_device is not None:
			dut_opts = "--device=%s,%s" % (options.vesna_device, dut_opts)

		dut_class = LocalDeviceUnderTest
	else:
		f = options.dut_name.split('.')
		dut_module_name = '.'.join(f[:-1])
		dut_name = '.'.join(f[-1:])

		__import__(dut_module_name, globals(), locals(), [], -1)
		dut_class = getattr(sys.modules[dut_module_name], dut_name)

	return dut_class(dut_opts, options.name,
			replay=options.replay, log_path=options.log_path,
			device_id=device_id, config_id=config_id)

def run_tests(options):

	dut = load_dut(options)

	test_kwargs = parse_test_kwargs(options.test_kwargs)

	gen = SignalGenerator(options.usbtmc_device)

	run_all = not any(getattr(options, name) for name, testfunc in iter_tests())

	log("Command line: %s" % (' '.join(sys.argv),))
	log("Session started at %s" % (datetime.datetime.now()))

	for name, testfunc in sorted(iter_tests(), key=lambda x:"ident" not in x[0]):
		if run_all or getattr(options, name):
			testfunc(dut, gen, **test_kwargs)

	log("Session ended at %s" % (datetime.datetime.now()))

def iter_tests():
	for name, testfunc in globals().iteritems():
		if name.startswith("test_") and callable(testfunc):
			yield name, testfunc

def main():
	default_log_path = datetime.datetime.now().strftime("log_%Y%m%d")

	parser = optparse.OptionParser()
	parser.add_option("-d", "--vesna-device", dest="vesna_device", metavar="DEVICE",
			help="Use VESNA spectrum sensor attached to DEVICE.")
	parser.add_option("-g", "--usbtmc-device", dest="usbtmc_device", metavar="DEVICE",
			help="Use signal generator attached to DEVICE.", default="/dev/usbtmc3")
	parser.add_option("-i", "--id", dest="name", metavar="ID",
			help="Use ID as identification for device under test.")
	parser.add_option("-o", "--log", dest="log_path", metavar="PATH", default=default_log_path,
			help="Write measurement logs under PATH.")
	parser.add_option("-n", "--replay", dest="replay", action="store_true",
			help="Replay measurement from logs.")
	parser.add_option("--vesna-config", dest="vesna_config", metavar="CONFIG",
			help="Manually choose a specific device configuration", default="0,0")
	parser.add_option("-R", "--remote", dest="dut_name", metavar="OBJECT",
			help="Remotely access Device Under Test using OBJECT.")
	parser.add_option("-O", "--remote-option", dest="dut_opts", metavar="OPTIONS", default="",
			help="Any additional options needed for the remote access object. Separate "
			"multiple options with a comma. Use -O,--help with -R to list available options.")
	parser.add_option("-T", "--test-option", dest="test_kwargs", metavar="OPTIONS", default="",
			help="Any additional options to pass to specified tests. Separate "
			"multiple options with a comma.")

	group = optparse.OptionGroup(parser, "Tests", description="Choose only specific tests to run "
			"(default is to run all tests)")

	for name, testfunc in iter_tests():
		opt = "--" + name.replace("_", "-")
		group.add_option(opt, dest=name, action="store_true", help=testfunc.__doc__)

	parser.add_option_group(group)

	(options, args) = parser.parse_args()

	if not options.name:
		print "Please specify ID for device under test using the \"-i\" option"
		return

	try:
		os.mkdir(options.log_path)
	except OSError:
		pass

	global log_f
	log_f = open("%s/%s.log" % (options.log_path, options.name), "w")

	run_tests(options)

main()
