import re
import warnings
from io import BytesIO
import numpy as np
from pymeasure.instruments import Instrument, RangeException, discreteTruncate
from pyvisa import VisaIOError

from pymeasure.adapters import PrologixAdapter

adapter = PrologixAdapter("ASRL4::INSTR", address=16, visa_library="@py")


class HP8753E(Instrument):
    """Represents the HP8753E Vector Network Analyzer
    and provides a high-level interface for taking scans of the
    scattering parameters.
    """

    SCAN_POINT_VALUES = [3, 11, 21, 26, 51, 101, 201, 401, 801, 1601]
    SCATTERING_PARAMETERS = ("S11", "S12", "S21", "S22")
    S11, S12, S21, S22 = SCATTERING_PARAMETERS

    start_frequency = Instrument.control(
        "STAR?",
        "STAR %e Hz",
        """ A floating point property that represents the start frequency
        in Hz. This property can be set.
        """,
    )
    stop_frequency = Instrument.control(
        "STOP?",
        "STOP %e Hz",
        """ A floating point property that represents the stop frequency
        in Hz. This property can be set.
        """,
    )
    sweep_time = Instrument.control(
        "SWET?",
        "SWET%.2e",
        """ A floating point property that represents the sweep time
        in seconds. This property can be set.
        """,
    )
    averages = Instrument.control(
        "AVERFACT?",
        "AVERFACT%d",
        """ An integer representing the number of averages to take. Note that
        averaging must be enabled for this to take effect. This property can be set.
        """,
        cast=lambda x: int(
            float(x)
        ),  # need float() to convert scientific notation in strings
    )
    averaging_enabled = Instrument.control(
        "AVERO?",
        "AVERO%d",
        """ A bool that indicates whether or not averaging is enabled. This property
        can be set.""",
        cast=bool,
    )

    def __init__(
        self, adapter, name="Hewlett Packard 8753E Vector Network Analyzer", **kwargs
    ):
        super().__init__(adapter, name, **kwargs)

    @property
    def id(self):
        """Get the identification of the instrument."""
        if self.SCPI:
            result = self.ask("*IDN?").strip("\n")
            return result
        else:
            raise NotImplementedError(
                "Non SCPI instruments require implementation in subclasses"
            )

    def set_fixed_frequency(self, frequency):
        """Sets the scan to be of only one frequency in Hz"""
        self.start_frequency = frequency
        self.stop_frequency = frequency
        self.scan_points = 3

    @property
    def parameter(self):
        for parameter in HP8753E.SCATTERING_PARAMETERS:
            if int(self.values("%s?" % parameter)) == 1:
                return parameter
        return None

    @parameter.setter
    def parameter(self, value):
        if value in HP8753E.SCATTERING_PARAMETERS:
            self.write("%s" % value)
        else:
            raise Exception(
                "Invalid scattering parameter requested" " for Hewlett Packard 8753E"
            )

    @property
    def scan_points(self):
        """Gets the number of scan points"""
        search = re.search(r"\d\.\d+E[+-]\d{2}$", self.ask("POIN?"), re.MULTILINE)
        if search:
            return int(float(search.group()))
        else:
            raise Exception("Improper message returned for the" " number of points")

    @scan_points.setter
    def scan_points(self, points):
        """Sets the number of scan points, truncating to an allowed
        value if not properly provided
        """
        points = discreteTruncate(points, HP8753E.SCAN_POINT_VALUES)
        if points:
            self.write("POIN%d" % points)
        else:
            raise RangeException(
                "Maximum scan points (1601) for" " Hewlett Packard 8753E exceeded"
            )

    def set_IF_bandwidth(self, bandwidth):
        """Sets the resolution bandwidth (IF bandwidth)"""
        allowedBandwidth = [10, 30, 100, 300, 1000, 3000, 3700, 6000]
        bandwidth = discreteTruncate(bandwidth, allowedBandwidth)
        if bandwidth:
            self.write("IFBW%d" % bandwidth)
        else:
            raise RangeException(
                "Maximum IF bandwidth (6000) for Hewlett Packard " "8753E exceeded"
            )

    def set_averaging(self, averages):
        """Sets the number of averages and enables/disables averaging. Should be
        between 1 and 999"""
        averages = int(averages)
        if not 1 <= averages <= 999:
            assert RangeException("Set", averages, "must be in the range 1 to 999")
        self.averages = averages
        self.averaging_enabled = averages > 1

    def disable_averaging(self):
        """Disables averaging"""
        warnings.warn(
            "Don't use disable_averaging(), use averaging_enabled = False instead",
            FutureWarning,
        )
        self.averaging_enabled = False

    def enable_averaging(self):
        """Enables averaging"""
        warnings.warn(
            "Don't use enable_averaging(), use averaging_enabled = True instead",
            FutureWarning,
        )
        self.averaging_enabled = True

    def is_averaging(self):
        """Returns True if averaging is enabled"""
        warnings.warn(
            "Don't use is_averaging(), use averaging_enabled instead", FutureWarning
        )
        return self.averaging_enabled

    def restart_averaging(self, averages):
        warnings.warn(
            "Don't use restart_averaging(), use scan_single() instead", FutureWarning
        )
        self.scan_single()

    def scan(self, averages=None, blocking=None, timeout=None, delay=None):
        """Initiates a scan with the number of averages specified and
        blocks until the operation is complete.
        """
        if (
            averages is not None
            or blocking is not None
            or timeout is not None
            or delay is not None
        ):
            warnings.warn(
                "averages, blocking, timeout, and delay arguments are no longer used by scan()",
                FutureWarning,
            )
        self.write("*CLS")
        self.scan_single()
        # All queries will block until the scan is done, so use NOOP? to check.
        # These queries will time out after several seconds though,
        # so query repeatedly until the scan finishes.
        while True:
            try:
                self.ask("NOOP?")
            except VisaIOError as e:
                if e.abbreviation != "VI_ERROR_TMO":
                    raise e
            else:
                break

    def scan_single(self):
        """Initiates a single scan"""
        if self.averaging_enabled:
            self.write("NUMG%d" % self.averages)
        else:
            self.write("SING")

    def scan_continuous(self):
        """Initiates a continuous scan"""
        self.write("CONT")

    @property
    def frequencies(self):
        """Returns a list of frequencies from the last scan"""
        return np.linspace(
            self.start_frequency, self.stop_frequency, num=self.scan_points
        )

    @property
    def data_complex(self):
        """Returns the complex power from the last scan"""
        # TODO: Implement binary transfer instead of ASCII
        self.ask("FORM4;OUTPDATA")
        data = np.loadtxt(
            BytesIO(self.read().encode()),
            delimiter=" ",
            dtype=np.float32,
        )

        print(data)
        self.read()
        # data_complex = data[:, 0] + 1j * data[:, 1]
        return data

    @property
    def data_log_magnitude(self):
        """Returns the absolute magnitude values in dB from the last scan"""
        return 20 * np.log10(self.data_magnitude)

    @property
    def data_magnitude(self):
        """Returns the absolute magnitude values from the last scan"""
        return np.abs(self.data_complex)

    @property
    def data_phase(self):
        """Returns the phase in degrees from the last scan"""
        return np.degrees(np.angle(self.data_complex))

    @property
    def data(self):
        """Returns the real and imaginary data from the last scan"""
        warnings.warn(
            "Don't use this function, use data_complex instead", FutureWarning
        )
        data_complex = self.data_complex
        return data_complex.real, data_complex.complex

    def log_magnitude(self, real, imaginary):
        """Returns the magnitude in dB from a real and imaginary
        number or numpy arrays
        """
        warnings.warn(
            "Don't use log_magnitude(), use data_log_magnitude instead", FutureWarning
        )
        return 20 * np.log10(self.magnitude(real, imaginary))

    def magnitude(self, real, imaginary):
        """Returns the magnitude from a real and imaginary
        number or numpy arrays
        """
        warnings.warn("Don't use magnitude(), use data_magnitude", FutureWarning)
        return np.sqrt(real**2 + imaginary**2)

    def phase(self, real, imaginary):
        """Returns the phase in degrees from a real and imaginary
        number or numpy arrays
        """
        warnings.warn("Don't use phase(), use data_phase instead", FutureWarning)
        return np.arctan2(imaginary, real) * 180 / np.pi


vna = HP8753E(adapter)

vna.id, vna.scan_points

import time

points = vna.scan_points = 101
print(f'VNA is sweeping {points}')
buffer = []
buffer.append(vna.ask("OUTPDATA").strip('\n'))
adapter.gpib_read_timeout = 1
start = time.time()
i = 0
while len(buffer) < (points):
    if vna.adapter.connection.bytes_in_buffer >= 50:
        try:
            buffer.append(vna.read().strip('\n'))
        except:
            print(f'miss {i} {vna.adapter.connection.bytes_in_buffer} {len(buffer)}', end='\r')
            i += 1
            time.sleep(0.001)
    else:
        print(f'empty? {i} {vna.adapter.connection.bytes_in_buffer} {len(buffer)}', end='\r')
        i += 1
        time.sleep(0.001)
    if time.time() - start > 10:
        break

print()
print(f'VNA is sweeping {points}, retreived {len(buffer)} points of data\n')
print(f'{buffer[0]}, {buffer[1]}, {buffer[-1]}\n')
print(f'Bytes in prologix buffer: {vna.adapter.connection.bytes_in_buffer}')

adapter.gpib_read_timeout = 2000

VNA is sweeping 101

VNA is sweeping 101, retreived 101 points of data

   1.021729000000000E+00,   1.497803000000000E-01,    5.731812000000000E-01,   8.118286000000000E-01,   -3.325806000000000E-01,  -9.866333000000000E-01

Bytes in prologix buffer: 0

VNA is sweeping 201
empty? 7090 42 163
VNA is sweeping 201, retreived 163 points of data

   1.022217000000000E+00,   1.495361000000000E-01,    9.062500000000000E-01,   4.926758000000000E-01,   -1.855469000000000E-01,  -9.708862000000000E-01

Bytes in prologix buffer: 42