# *****************************************************************************
# 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:
#   Daniel Margineda <daniel.margineda@psi.ch>
#   Oksana Shliakhtun <oksana.shliakhtun@psi.ch>
# *****************************************************************************
"""Signal Recovery SR7270: lockin amplifier for AC susceptibility"""

import re
from frappy.core import Readable, Parameter, FloatRange, TupleOf, \
    HasIO, StringIO, BoolType, EnumType
from frappy.errors import RangeError


class SR_IO(StringIO):
    end_of_line = b'\x00'
    identification = [('ID', r'72.*')]  # Identification; causes the lock-in amplifier to respond with the model number

    def communicate(self, cmd):  # remove dash from terminator
        reply = super().communicate(cmd)
        status = self._conn.readbytes(2, timeout=0.1)  # get the 2 status bytes
        return reply + ';%d;%d' % tuple(status)


class XY(HasIO, Readable):
    value = Parameter('X, Y', datatype=TupleOf(FloatRange(unit='V'), FloatRange(unit='V')))
    freq = Parameter('oscill. frequen. control', FloatRange(0, 250e3), unit='Hz', readonly=False)
    amp = Parameter('oscill. amplit. control', FloatRange(0.00, 5), unit='V_rms', readonly=False)
    range = Parameter('sensitivity value', FloatRange(0.00, 1), unit='V', default=1, readonly=False)
    autorange = Parameter('autorange_on', EnumType('autorange', off=0, soft=1, hard=2),
                          readonly=False, default=0)

    sen_range = {index + 1: name for index, name in enumerate(
        ['2nV', '5nV', '10nV', '20nV', '50nV', '100nV', '200nV', '500nV', '1uV',
         '2uV', '5uV', '10uV', '20uV', '50uV', '100uV', '200uV', '500uV', '1mV',
         '2mV', '5mV', '10mV', '20mV', '50mV', '100mV', '200mV', '500mV', '1V']
    )}

    irange = Parameter('sensitivity index', EnumType('sensitivity index range', sen_range))
    nm = Parameter('noise mode on', BoolType(), readonly=False)
    phase = Parameter('reference phase control', FloatRange(-360, 360), unit='deg', readonly=False)
    vmode = Parameter('control mode', EnumType(both_grounded=0, A=1, B=2, A_B_diff=3), readonly=False)

    ioClass = SR_IO

    def doPoll(self):
        super().doPoll()
        if self.autorange == 1:  # soft auto range
            if max(abs(x), abs(y)) >= 0.9 * self.range and self.irange < 27:
                self.write_irange(self.irange + 1)
            elif max(abs(x), abs(y)) <= 0.3 * self.range and self.irange > 1:
                self.write_irange(self.irange - 1)

    def comm(self, cmd):
        reply, status, overload = self.communicate(cmd).split(';')  # try/except
        reply = reply.rstrip('\n')
        if overload != '0':
            self.status = (self.Status.WARN, f'overload {overload}')
        self.status = (self.Status.IDLE, '')
        return reply

    def string_to_value(self, value):
        """
        This method is used for writing methods (sensitivity range and time constant in particular). As the dictionaries
        with names were used (pop-up list in gui) this method transforms names into float values, that can be used
        further.
        :param value:
        :return:
        """
        value_with_unit = re.compile(r'(\d+)([pnumkMG]?)')
        value, pfx = value_with_unit.match(value).groups()
        pfx_dict = {'p': 1e-12, 'n': 1e-9, 'u': 1e-6, 'm': 1e-3, 'k': 1e3, 'M': 1e6, 'G': 1e9}
        if pfx in pfx_dict:
            value = round(float(value) * pfx_dict[pfx], 12)
        return float(value)

    def read_value(self):
        reply = self.comm('XY.').split(',')
        x = float(reply[0])
        y = float(reply[1])
        return x, y

    def read_freq(self):
        return float(self.comm('OF.'))

    def write_freq(self, freq):
        self.comm(f'OF. {freq}')
        return freq

    def read_autorange(self):
        reply = self.comm('AUTOMATIC')
        # determine hardware autorange
        if reply == 1:  # soft
            return 2  # hard
        if self.autorange == 0:  # soft
            return self.autorange  # read autorange
        return self.autorange  # off

    def write_autorange(self, value):
        if value == 2:  # hard
            self.comm('AS')  # put hardware autorange on
            self.comm('AUTOMATIC. 1')
        else:
            self.comm('AUTOMATIC. 0')
        return value

    def read_amp(self):
        return float(self.comm('OA.'))

    def write_amp(self, amp):
        self.comm(f'OA. {amp}')
        return amp

    def read_irange(self):
        reply = self.comm('SEN')
        return int(reply)

    def write_irange(self, irange):
        value = int(irange)
        self.comm(f'SEN {value}')
        self.read_range()
        return value

    def read_range(self):
        reply = self.comm('SEN.')  # range value
        return float(reply)

    def write_range(self, target):
        cl_idx = None

        for idx, name in self.sen_range.items():

            if idx < len(self.sen_range):
                value_l = self.string_to_value(self.sen_range.get(idx))
                value_r = self.string_to_value(self.sen_range.get(idx + 1))
                if value_l < target <= value_r:
                    cl_idx = idx + 1
                    break
            cl_idx = idx
        self.write_irange(cl_idx)
        return self.read_range()

    def read_nm(self):
        reply = self.comm('NOISEMODE')
        return reply

    def write_nm(self, value):
        self.comm('NOISEMODE %d' % int(value))
        self.read_nm()
        return value

    # phase and autophase
    def read_phase(self):
        reply = self.comm('REFP.')
        return float(reply)

    def write_phase(self, value):
        self.comm(f'REFP {round(1000 * value)}')
        return self.read_phase()

    def aphase(self):
        """auto phase"""
        self.read_phase()
        return self.comm('AQN')

    def read_vmode(self):
        reply = self.comm('VMODE')
        return int(reply)

    def write_vmode(self, vmode):
        value = int(vmode)
        self.comm(f'VMODE {value}')
        return value


class XY70(XY):

    time_const = {name: value for value, name in enumerate(
        ['10us', '20us', '50us', '100us', '200us', '500us', '1ms', '2ms', '5ms', '10ms',
         '20ms', '50ms', '100ms', '200ms', '500ms', '1s', '2s', '5s', '10s', '20s', '50s',
         '100s', '200s', '500s', '1ks', '2ks', '5ks', '10ks', '20ks', '50ks', '100ks']
    )}

    tc = Parameter('time const. value', FloatRange(0.00001, 100000), unit='s', readonly=False)
    itc = Parameter('time const. index', EnumType('time const. index range', time_const), readonly=False)

    ioClass = SR_IO

    def comm(self, cmd):
        reply, status, overload = self.communicate(cmd).split(';')
        reply = reply.rstrip('\n')
        if overload != '0':
            self.status = (self.Status.WARN, f'overload {overload}')
        self.status = (self.Status.IDLE, '')
        return reply

    def read_tc(self):
        reply = self.comm('TC.')
        return float(reply)

    def write_tc(self, target):
        cl_idx = None
        cl_diff = float('inf')

        for name, idx in self.time_const.items():
            value = self.string_to_value(name)
            diff = abs(value - target)  # range is the actual value, like SEN.
            if diff < cl_diff:
                cl_idx = idx
                cl_diff = diff

        if self.nm is False or (self.nm is True and 5 <= cl_idx <= 9):
            self.comm(f'TC {cl_idx}')
            return self.read_tc()
        raise RangeError('Not allowed with noisemode=1')

    def read_itc(self):
        return int(self.comm('TC'))

    def write_itc(self, itc):
        value = int(itc)
        self.comm(f'TC {value}')
        self.read_tc()
        return value


class XY65(XY):

    time_const = {name: value for value, name in enumerate(
        ['10us', '20us', '40us', '80us', '160us', '320us', '640us', '5ms', '10ms', '20ms',
         '50ms', '100ms', '200ms', '500ms', '1s', '2s', '5s', '10s', '20s', '50s', '100s',
         '200s', '500s', '1ks', '2ks', '5ks', '10ks', '20ks', '50ks', '100ks']
    )}

    def read_tc(self):
        reply = self.comm('TC.')
        return float(reply)

    def write_tc(self, target):
        cl_idx = None
        cl_diff = float('inf')

        for name, idx in self.time_const.items():
            value = self.string_to_value(name)
            diff = abs(value - target)  # range is the actual value, like SEN.
            if diff < cl_diff:
                cl_idx = idx
                cl_diff = diff

        if self.nm is False or (self.nm is True and 5 <= cl_idx <= 9):
            self.comm(f'TC {cl_idx}')
            return self.read_tc()
        raise RangeError('Not allowed with noisemode=1')

    def read_itc(self):
        return int(self.comm('TC'))

    def write_itc(self, itc):
        value = int(itc)
        self.comm(f'TC {value}')
        self.read_tc()
        return value
