# *****************************************************************************
# 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: Oksana Shliakhtun <oksana.shliakhtun@psi.ch>
# *****************************************************************************
"""Stanford Research Systems SR830 DS Lock-in Amplifier"""

import re
import time
from frappy.core import StringIO, HasIO, Parameter, EnumType, FloatRange, TupleOf, ERROR, IDLE, WARN, StatusType, \
    Readable, BUSY

from frappy.errors import IsBusyError


def string_to_value(value):
    """
    Converting the value to float, removing the units, converting the prefix into the number.
    :param value: value
    :return: float value without units
    """
    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)


class SR830_IO(StringIO):
    end_of_line = b'\r'  # should be <lf> or <cr>
    identification = [('*IDN?', r'Stanford_Research_Systems,.*')]


class StanfRes(HasIO, Readable):
    def set_par(self, cmd, *args):
        """
        Set parameter.
        Query commands are the same as setting commands, but they have a question mark.
        :param cmd: command
        :param args: value(s)
        :return: reply
        """
        head = ','.join([cmd] + [a if isinstance(a, str) else f'{a:g}' for a in args])
        tail = cmd.replace(' ', '? ')
        new_tail = re.sub(r'[0-9.]+', '', tail)

        reply = self.communicate(f'{head};{new_tail}')
        result = []
        for num in reply.split(','):
            try:
                result.append(float(num))
            except ValueError:
                result.append(num)
        if len(result) == 1:
            return result[0]
        return result

    def get_par(self, cmd):
        reply = self.communicate(cmd)
        result = []
        for num in reply.split(','):
            try:
                result.append(float(num))
            except ValueError:
                result.append(num)
        if len(result) == 1:
            return result[0]
        return result


class XY(StanfRes):
    value = Parameter('X, Y', datatype=TupleOf(FloatRange(unit='V'), FloatRange(unit='V')))
    amp = Parameter('oscill. amplit. control', FloatRange(4e-3, 5), unit='V', readonly=False)
    freq = Parameter('oscill. frequen. control', FloatRange(1e-3, 102000), unit='Hz', readonly=False)
    phase = Parameter('reference phase control', FloatRange(-360, 729), unit='deg', readonly=False)
    autorange = Parameter('autorange_on', EnumType('autorange', off=0, soft=1, hard=2),
                          readonly=False, default=0)

    status = Parameter(datatype=StatusType(Readable, 'BUSY'))

    SEN_RANGE = ['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',
                                                     {name: idx for idx, name in enumerate(SEN_RANGE)}), readonly=False)

    range = Parameter('sensitivity value', FloatRange(2e-9, 1), unit='V', default=1, readonly=False)

    TIME_CONST = ['10us', '30us', '100us', '300us', '1ms', '3ms', '10ms', '30ms', '100ms', '300ms',
                  '1s', '3s', '10s', '30s', '100s', '300s', '1ks', '3ks', '10ks', '30ks']

    tc = Parameter('time const. value', FloatRange(1e-6, 3e4), unit='s', readonly=False)
    itc = Parameter('time const. index', EnumType(
        'time const. index range', {name: value for value, name in enumerate(TIME_CONST)}), readonly=False)

    SEN_RANGE_values = [string_to_value(value) for value in SEN_RANGE]
    TIME_CONST_values = [string_to_value(value) for value in TIME_CONST]

    ioClass = SR830_IO

    _autogain_started = 0

    # status = serial poll status byte, standard event status byte, lock-in status byte, error status byte
    def read_status(self):
        if time.time() < self._autogain_started + self.tc * 10:
            return BUSY, 'changing gain'

        stb = int(self.communicate('*STB?'))  # serial poll status byte
        esr = int(self.communicate('*ESR?'))  # standard event status byte
        lias = int(self.communicate('LIAS?'))  # lock-in status byte
        errs = int(self.communicate('ERRS?'))  # error status byte

        if lias & (1 << 2):
            return ERROR, 'output overload'
        if lias & (1 << 1):
            return ERROR, 'tc overload'
        if lias & (1 << 0):
            return ERROR, 'reserve/input overload'
        if esr & (1 << 5):
            return ERROR, 'illegal command'
        if esr & (1 << 4):
            return ERROR, 'execution error'
        if errs & (1 << 1):
            return ERROR, 'backup error'
        if errs & (1 << 2):
            return ERROR, 'RAM error'
        if errs & (1 << 4):
            return ERROR, 'ROM error'
        if errs & (1 << 5):
            return ERROR, 'GRIB error'
        if errs & (1 << 6):
            return ERROR, 'DSP error'
        if errs & (1 << 7):
            return ERROR, 'internal math error'
        if esr & (1 << 0):
            return WARN, 'input queue overflow, cleared'
        if esr & (1 << 2):
            return WARN, 'output queue overflow, cleared'
        if lias & (1 << 3):
            return WARN, 'reference unlock'
        if lias & (1 << 4):
            return WARN, 'freq crosses 200 Hz'
        if not stb & (1 << 0):
            return BUSY, 'scan in progress'
        if not stb & (1 << 1):
            return BUSY, 'command execution in progress'
        return IDLE, ''

    def read_value(self):
        """
        Read XY. The manual autorange implemented.
        :return:
        """
        if self.read_status()[0] == BUSY:
            raise IsBusyError('changing gain')
        reply = self.get_par('SNAP? 1, 2')
        value = tuple(float(x) for x in reply)
        x, y = value
        maxxy = max(abs(x), abs(y))
        if self.autorange == 1:
            if maxxy >= 0.9 * self.range and self.irange < 26:
                self.write_irange(self.irange + 1)
            elif maxxy <= 0.3 * self.range and self.irange > 0:
                self.write_irange(self.irange - 1)
        return value

    def read_irange(self):
        return int(self.get_par('SENS?'))

    def read_range(self):
        """Sensitivity range value"""
        idx = self.read_irange()
        name = self.SEN_RANGE[idx]
        return string_to_value(name)

    def write_irange(self, irange):
        """Index of sensitivity from the range"""
        value = int(irange)
        self.set_par(f'SENS {value}')
        self._autogain_started = time.time()
        self.read_range()
        return value

    def write_range(self, target):
        """
        Setting the sensitivity range.
        cl_idx/cl_value is the closest index/value from the range to the target
        :param target:
        :return: closest value of the sensitivity range
        """
        target = float(target)
        cl_idx = None
        cl_value = float('inf')

        for idx, sen_value in enumerate(self.SEN_RANGE_values):
            if sen_value >= target:
                diff = sen_value - target

                if diff < cl_value:
                    cl_value = sen_value
                    cl_idx = idx

        self.set_par(f'SENS {cl_idx}')
        return cl_value

    def read_itc(self):
        """Time constant index from the range"""
        return int(self.get_par(f'OFLT?'))

    def write_itc(self, itc):
        value = int(itc)
        return self.set_par(f'OFLT {value}')

    def read_tc(self):
        """Read time constant value from the range"""
        idx = self.read_itc()
        name = self.TIME_CONST[idx]
        return string_to_value(name)

    def write_tc(self, target):
        """
        Setting the time constant from the range.
        cl_idx/cl_value is the closest index/value from the range to the target
        :param target: time constant
        :return: closest time constant value
        """
        target = float(target)
        cl_idx = None
        cl_value = float('inf')

        for idx, time_value in enumerate(self.TIME_CONST_values):
            if time_value >= target:
                diff = time_value - target

                if diff < cl_value:
                    cl_value = time_value
                    cl_idx = idx

        self.set_par(f'OFLT {cl_idx}')
        return cl_value

    def read_phase(self):
        return float(self.get_par('PHAS?'))

    def write_phase(self, value):
        return self.set_par(f'PHAS {value}')

    def read_freq(self):
        return float(self.get_par('FREQ?'))

    def write_freq(self, value):
        return self.set_par(f'FREQ {value}')

    def read_amp(self):
        return float(self.get_par('SLVL?'))

    def write_amp(self, value):
        return self.set_par(f'SLVL {value}')

    def auto_phase(self):
        return self.set_par('APHS')