import VNA
import pyvisa
import time
import numpy
import math
import cmath

import dut_measurement


class VnaMeasurement(dut_measurement.DutMeasurement):
    def __init__(self, config_data, target_accuracy):
        self.delta_mag = config_data['delta_mag']
        self.delta_phase = config_data['delta_phase']
        self.reference_signal_names = config_data['reference_signals']

        self.vna = VNA.create_vna(config_data['vna_ip'], target_accuracy)

        self.vna.load_config(config_data['vna_config_file'])
        self.vna.create_new_trace("Trace1", "S11")
        self.vna.create_new_trace("Trace2", "S12")
        self.vna.create_new_trace("Trace3", "S21")
        self.vna.create_new_trace("Trace4", "S22")

        self.frequencies = {}
        for f in config_data['frequencies']:
            self.frequencies[str(f / 1e9) + 'GHz'] = f

    def _get_trace_data(self, trace):
        return self.vna.get_list_of_measurement_values(trace, "SDAT")

    def get_dut_measurements(self, set_name):
        # FIXME: The try/catch should be way down in the VNA class
        for iteration in range(10):
            try:
                power = self.vna.get_current_power()
                frequency = self.frequencies[set_name]
                self.vna.set_cw_frequency(frequency)
                self.vna.do_single_sweep()
                s11 = self._get_trace_data("Trace1")
                s12 = self._get_trace_data("Trace2")
                s21 = self._get_trace_data("Trace3")
                s22 = self._get_trace_data("Trace4")

                return {'RF_POWER': power, 'RF_FREQUENCY': frequency,
                        'S11_MAGNITUDE': self.calculate_mean_magnitude_db(s11),
                        'S11_PHASE': self.calculate_mean_phase(s11),
                        'S12_MAGNITUDE': self.calculate_mean_magnitude_db(s12),
                        'S12_PHASE': self.calculate_mean_phase(s12),
                        'S21_MAGNITUDE': self.calculate_mean_magnitude_db(s21),
                        'S21_PHASE': self.calculate_mean_phase(s21),
                        'S22_MAGNITUDE': self.calculate_mean_magnitude_db(s22),
                        'S22_PHASE': self.calculate_mean_phase(s21),
                        'SET_NAME': set_name}

            # exception for pyvisa error or timeout
            except pyvisa.errors.VisaIOError:
                # call reset function for VNA status register and error queue
                self.vna.reset_status()
                print('An error occurred during VNA read out')
                # wait one second for next try
                time.sleep(1)
        # FIXME: In case we did not succeed we need something like a stop_measurement exception

        raise Exception('FIXME: Throw stop_measurement here, and handle it. Don\'t crash!')

    def get_dut_signal_names(self):
        return ['RF_POWER', 'RF_FREQUENCY', 'S11_MAGNITUDE', 'S11_PHASE', 'S12_MAGNITUDE',
                'S12_PHASE', 'S21_MAGNITUDE', 'S21_PHASE', 'S22_MAGNITUDE', 'S22_PHASE', 'SET_NAME']

    def get_dut_reference_signal_names(self):
        return self.reference_signal_names

    def get_dut_max_delta_signals(self):
        return [self.delta_phase, self.delta_mag]

    def calculate_complex_numbers(self, values_list):
        real_num = values_list[::2]
        imaginary_num = values_list[1::2]
        complex_numbers = []
        for i, q in zip(real_num, imaginary_num):
            complex_numbers.append(complex(i, q))
        return complex_numbers

    def calculate_magnitudes(self, values_list):
        complex_numbers = self.calculate_complex_numbers(values_list)
        magnitudes = [abs(val) for val in complex_numbers]
        return magnitudes

    def calculate_mean_magnitude(self, values_list):
        magnitudes = self.calculate_magnitudes(values_list)
        return numpy.mean(magnitudes)

    def calculate_mean_magnitude_db(self, values_list):
        return 20*math.log10(self.calculate_mean_magnitude(values_list))

    def calculate_phases(self, values_list):
        complex_numbers = self.calculate_complex_numbers(values_list)
        phases = [math.degrees(cmath.phase(val)) for val in complex_numbers]
        return phases

    def calculate_mean_phase(self, values_list):
        phases = self.calculate_phases(values_list)
        return numpy.mean(phases)

    def get_measurement_set_names(self):
        return list(self.frequencies.keys())