# *****************************************************************************
# This program is free software; you can redistribute it and/or modify it under
# the terms of the GNU General Public License as published by the Free Software
# Foundation; either version 2 of the License, or (at your option) any later
# version.
#
# This program is distributed in the hope that it will be useful, but WITHOUT
# ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
# FOR A PARTICULAR PURPOSE.  See the GNU General Public License for more
# details.
#
# You should have received a copy of the GNU General Public License along with
# this program; if not, write to the Free Software Foundation, Inc.,
# 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
#
# Module authors:
#   Anik Stark <anik.stark@psi.ch>
# *****************************************************************************
""" function generator agilent 33210a """

import re
from frappy.core import StringIO, HasIO, Writable, Parameter, FloatRange, EnumType, \
    IDLE, WARN, nopoll


class IO(StringIO):
    end_of_line = '\n'
    identification = [('*IDN?', 'Agilent Technologies,33210A.*')]     # Agilent Technologies,33210A,0,f.ff-b.bb-aa-p

class Frequency(HasIO, Writable):

    ioClass = IO

    target = Parameter('target frequency', datatype=FloatRange(unit='Hz'), readonly=False)
    value = Parameter('frequency', datatype=FloatRange(unit='Hz'))
    function = Parameter('function', datatype=EnumType(SIN=0, SQU=1, RAMP=2), readonly=False)
    voltage = Parameter('peak-to-peak voltage', datatype=FloatRange(unit='V'), readonly=False)
    offset = Parameter('voltage offset', datatype=FloatRange(unit='V'), readonly=False)
    width = Parameter('width', datatype=FloatRange(unit='s'), readonly=False)
    burstcycles = Parameter('burst cycles', datatype=FloatRange(), readonly=False)
    burstperiod = Parameter('burst period', datatype=FloatRange(unit='s'), readonly=False)
    mode = Parameter('device mode', datatype=EnumType(continuous=0, burst=1, off=2), readonly=False)
    _prev_mode = ''

    def parse_reply(self, reply):
        match = re.match(r'"([A-Z]*) (.*),(.*),(.*)"', reply)
        self.function = match.group(1)
        self.value = float(match.group(2))
        self.voltage = float(match.group(3))
        self.offset = float(match.group(4))

    def read_value(self):
        reply = self.communicate('APPL?')
        self.parse_reply(reply)
        return self.value

    @nopoll
    def read_target(self):
        return self.read_value()

    def write_target(self, target):
        cmd = f'APPL:{self.function.name} {float(target)} HZ, {float(self.voltage)} VPP, {float(self.offset)}'
        reply = self.communicate(f'{cmd}\nAPPL?')
        self.parse_reply(reply)
        return self.value

    def write_function(self, function):
        self.function = function
        self.write_target(self.target)
        self.read_value()

    def write_voltage(self, voltage):
        self.voltage = voltage
        self.write_target(self.target)
        self.read_value()

    def write_offset(self, offset):
        self.offset = offset
        self.write_target(self.target)
        self.read_value()

    def read_width(self):
        return float(self.communicate('FUNC:PULS:WIDT?'))

    def write_width(self, width):
        self.communicate(f'FUNC:PULS:WIDT {width}')
        return self.read_width()

    def read_burstcycles(self):
        return int(float(self.communicate('BURS:NCYC?')))

    def write_burstcycles(self, burstcycles):
        self.communicate(f'BURS:NCYC {int(burstcycles)}')
        return self.read_burstcycles()

    def read_burstperiod(self):
        return float(self.communicate('BURS:INT:PER?'))

    def write_burstperiod(self, burstperiod):
        self.communicate(f'BURS:INT:PER {burstperiod}')
        return self.read_burstperiod()

    def read_mode(self):
        reply = self.communicate('OUTP?')
        if reply == '0':
            return 'off'
        reply = self.communicate('BURS:STAT?')
        if reply == 'ON':
            return 'burst'
        return 'continuous'

    def write_mode(self, mode):
        if mode == 'off':
            reply = self.communicate('OUTP OFF;OUTP?')
            if reply == '0':
                self.status = WARN, 'device is turned off'
            else:
                self.status = WARN, 'error when turning device off'
        if mode == 'burst':
            reply = self.communicate('OUTP ON;BURS:STAT ON\nBURS:STAT?')
            if reply == '1':
                self.status = IDLE, 'burst mode'
            else:
                self.status = WARN, 'error when turning burst mode on'
        if mode == 'continuous':
            reply = self.communicate('OUTP ON;BURS:STAT OFF\nBURS:STAT?')
            if reply == '0':
                self.status = IDLE, 'continuous mode'
            else:
                self.status = WARN, 'error when turning continuous mode on'
        self._prev_mode = mode
        return mode

    # def read_function(self):
    #     return self.communicate('FUNC?')

    # def write_function(self, function):
    #     self.communicate(f'FUNC {function}')
    #     return self.read_function()

    # def read_voltage(self):
    #     txtval, unit = self.communicate('VOLT?').split('')
    #     return float(txtval)

    # def write_voltage(self, voltage):
    #     self.communicate(f'VOLT {voltage:.1f} VPP')
    #     return self.read_voltage()

    # def read_offset(self):
    #     return float(self.communicate('VOLT:OFFS?'))

    # def write_offset(self, offset):
    #     self.communicate(f'VOLT:OFFS {offset:.1f}')
    #     return self.read_offset()
    