import astropy.units as u
from astropy.units import Quantity

import hpib
import serial.threaded

serial_port = serial.Serial("/dev/ttyUSB0")
command_thread = serial.threaded.ReaderThread(serial_port, hpib.GPIBProtocol)

class CamrasHpibDevice(object):
    """Wrapper around HPIB commands"""

    def __init__(self, address):
        if not command_thread.is_alive():
            command_thread.start()
        self.hpib_address = address

    @property
    def frequency(self):
        """Returns the frequency of the device as an astropy Quantity.
        Throws a RuntimeError if no response"""
        freq_str = command_thread.protocol.command("freq?", address=self.hpib_address)
        if len(freq_str)==0:
            raise RuntimeError("Camras device at address {} is not responding".format(self._hpib_address))
        return int(float(freq_str)) * u.Hz

    @frequency.setter
    def frequency(self, freq):
        """Set the device to the specified frequency. Throws a RuntimeError if failed"""
        if isinstance(freq, Quantity):
            freq = freq.to(u.Hz).value
        command_thread.protocol.command("freq {:d} Hz".format(int(freq)), address=self.hpib_address)
        new_freq = self.frequency
        if new_freq.to(u.Hz).value != int(freq):
            raise RuntimeError("Setting frequency failed: tried to set to {}, it is now {}".format(
                               freq, new_freq))

class Receiver(CamrasHpibDevice):
    """Wrapper around HPIB commands for the Rohde & Schwartz receiver"""

    def __init__(self, **kwargs):
        super(Receiver, self).__init__(1, **kwargs)

class LocalOscillator(CamrasHpibDevice):
    """Wrapper around HPIB commands for the local oscillator"""

    def __init__(self, **kwargs):
        super(LocalOscillator, self).__init__(28, **kwargs)

class ClockGenerator(CamrasHpibDevice):
    """Wrapper around HPIB commands for the clock generator (should be at 140MHz)"""

    def __init__(self, **kwargs):
        super(SignalGenerator, self).__init__(14, **kwargs)