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

import dut_measurement

class VnaData:
    def __init__(self, power, frequency, s11, s21, s12, s22):
        self.power = power
        self.frequency = frequency
        self.s11 = s11
        self.s21 = s21
        self.s12 = s12
        self.s22 = s22


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.vna = VNA.create_vna(config_data['vna_ip'], target_accuracy)

        self.vna.load_config(config_data['vna_config_file'],  config_data['frequency'])
        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")

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

    def get_dut_measurements(self):
        # 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.vna.get_current_cw_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)}

            # 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']

    def get_dut_reference_signal_names(self):
        return ['S21_MAGNITUDE', 'S21_PHASE']

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

    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)