# *****************************************************************************
# 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>
# *****************************************************************************
import math

from frappy.core import Readable, Parameter, IntRange, FloatRange, \
    StringIO, HasIO, StringType, Property, Writable, Drivable, IDLE, ERROR, \
    StructOf, WARN, Done, BoolType, Enum
from frappy.errors import RangeError
from frappy_psi.convergence import HasConvergence
from frappy.mixins import HasOutputModule, HasControlledBy


class Ls340IO(StringIO):
    """communication with 340CT"""
    end_of_line = '\r'
    wait_before = 0.05
    identification = [('*IDN?', r'LSCI,MODEL340,.*')]


class LakeShore(HasIO):
    def set_par(self, cmd, *args):
        head = ','.join([cmd] + [a if isinstance(a, str) else f'{a:g}' for a in args])
        tail = cmd.replace(' ', '? ')
        reply = self.communicate(f'{head};{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 Sensor340(LakeShore, Readable):
    """A channel of 340TC"""

    # define the communication class to create the IO module
    ioClass = Ls340IO
    channel = Property('lakeshore channel', StringType())
    alarm = Parameter('alarm limit', FloatRange(unit='K'), readonly=False)
    # define or alter the parameters
    # as Readable.value exists already, we give only the modified property 'unit'
    value = Parameter(unit='K')

    def read_value(self):
        return self.get_par(f'KRDG? {self.channel}')

    def read_status(self):
        c = int(self.get_par(f'RDGST? {self.channel}'))
        if c >= 128:
            return ERROR, 'units overrange'
        if c >= 64:
            return ERROR, 'units zero'
        if c >= 32:
            return ERROR, 'temperature overrange'
        if c >= 16:
            return ERROR, 'temperature underrange'
        # do not check for old reading -> this happens regularly on NTCs with T comp
        if c % 2:
            return ERROR, 'invalid reading'
        # ask for high alarm status and return warning
        if 1 in self.get_par(f'ALARMST? {self.channel}'):
            return WARN, 'alarm triggered'
        return IDLE, ''

    def write_alarm(self, alarm):
        return self.set_par(f'ALARM {self.channel}', 1, 1, alarm, 0, 0, 2)[2]

    def read_alarm(self):
        return self.get_par(f'ALARM? {self.channel}')[2]


class HeaterOutput(LakeShore, HasControlledBy, HasIO, Writable):
    max_power = Parameter('max heater power', datatype=FloatRange(0, 100), unit='W', readonly=False)
    value = Parameter('heater output', datatype=FloatRange(0, 100), unit='W')
    target = Parameter('manual heater output', datatype=FloatRange(0, 100), unit='W')
    loop = Property('lakeshore loop', IntRange(1, 2), default=1)  # output
    channel = Property('attached channel', StringType())  # input
    resistance = Property('heater resistance', datatype=FloatRange(10, 100))
    _range = 0
    _max_power = 50

    MAXCURRENTS = {1: 0.25, 2: 0.5, 3: 1.0, 4: 2.0}
    RANGES = {1: 1e4, 2: 1e3, 3: 1e2, 4: 1e1, 5: 1}

    SETPOINTLIMS = 999

    STATUS_MAP = {
        0: (IDLE, ''),
        1: (ERROR, 'Power supply over voltage'),
        2: (ERROR, 'Power supply under voltage'),
        3: (ERROR, 'Output digital-to-analog Converter error'),
        4: (ERROR, 'Current limit digital-to-analog converter error'),
        5: (ERROR, 'Open heater load'),
        6: (ERROR, 'Heater load less than 10 ohms')
    }

    def earlyInit(self):
        super().earlyInit()
        self.CHOICES = sorted([(maxcurrent ** 2 * factor, icurrent, irange)
                               for irange, factor in self.RANGES.items()
                               for icurrent, maxcurrent in self.MAXCURRENTS.items()])

    def write_max_power(self, max_power):
        prev = 0
        for i, (factor, icurrent, irange) in enumerate(self.CHOICES):
            power = min(factor * self.resistance, 2500 / self.resistance)
            if power >= max_power:
                if prev >= max_power * 0.9 or prev == power:
                    icurrent, irange = self.CHOICES[i - 1][1:3]
                break
            prev = power
        self._range = irange
        self.set_par(f'CLIMIT {self.loop}', self.SETPOINTLIMS, 0, 0, icurrent, irange)
        self.set_par(f'RANGE {irange}')
        self.set_par(f'CDISP {self.loop}', 1, self.resistance, 1, 0)

    def read_max_power(self):
        setplimit, _, _, icurrent, irange = self.get_par(f'CLIMIT? {self.loop}')
        # max_power from codes disregarding voltage limit:
        self._max_power = self.MAXCURRENTS[icurrent] ** 2 * self.RANGES[irange] * self.resistance
        # voltage limit = 50V:
        max_power = min(self._max_power, 2500 / self.resistance)
        return max_power

    def set_range(self):
        self.set_par('RANGE ', self._range)

    def percent_to_power(self, percent):
        return min((percent / 100) ** 2 * self._max_power,
                   2500 / self.resistance)

    def power_to_percent(self, power):
        return (power / self._max_power) ** (1 / 2) * 100  # limit

    def read_status(self):
        return self.STATUS_MAP[self.get_par(f'HTRST?')]

    def write_target(self, target):
        self.self_controlled()
        self.write_max_power(self.max_power)
        self.set_heater_mode(3)  # 3=open loop
        self.set_range()
        percent = self.power_to_percent(target)
        reply = self.set_par(f'MOUT {self.loop}', percent)
        return self.percent_to_power(reply)

    def set_heater_mode(self, mode):
        self.set_par(f'CSET {self.loop}', self.channel, 1, 1, 0)
        self.set_par(f'CMODE {self.loop}', int(mode))
        return self.get_par(f'RANGE?')

    def read_value(self):
        return self.percent_to_power(self.get_par(f'HTR?{self.loop}'))


class HeaterOutput340(HeaterOutput):
    resistance = Property('heater resistance', datatype=FloatRange(10, 100))

    MAXCURRENTS = {1: 0.25, 2: 0.5, 3: 1.0, 4: 2.0}
    RANGES = {1: 1e4, 2: 1e3, 3: 1e2, 4: 1e1, 5: 1}

    STATUS_MAP = {
        0: (IDLE, ''),
        1: (ERROR, 'Power supply over voltage'),
        2: (ERROR, 'Power supply under voltage'),
        3: (ERROR, 'Output digital-to-analog Converter error'),
        4: (ERROR, 'Current limit digital-to-analog converter error'),
        5: (ERROR, 'Open heater load'),
        6: (ERROR, 'Heater load less than 10 ohms')
    }

    def read_value(self):
        return self.percent_to_power(self.get_par(f'HTR?'))  # no loop to be given on 340


class HeaterOutput336(HeaterOutput):

    power = 20

    STATUS_MAP = {
        0: (IDLE, ''),
        1: (ERROR, 'Open heater load'),
        2: (ERROR, 'Heater short')
    }

    def write_max_power(self, max_power):
        max_current = min(math.sqrt(self.power / self.resistance), 2500 / self.resistance)
        if self.loop == 1:
            max_current_limit = 2
        else:
            max_current_limit = 1.414
        if max_current > max_current_limit:
            raise RangeError('max_power above limit')
        if max_current >= max_current_limit / math.sqrt(10):
            self._range = 3
            user_current = max_current
        elif max_current >= max_current_limit / 10:
            self._range = 2
            user_current = max_current * math.sqrt(10)
        else:
            self._range = 1
            user_current = max_current * math.sqrt(100)
        self.set_par(f'HTRSET {self.loop}', 1 if self.resistance < 50 else 2, 0, user_current, 1)
        max_power = max_current ** 2 * self.resistance
        self._max_power = max_power
        self.set_range()
        return max_power


class TemperatureLoop340(HasConvergence, HasOutputModule, Sensor340, Drivable, LakeShore):
    Status = Enum(
        Drivable.Status,
        RAMPING=370,
        STABILIZING=380,
    )

    target = Parameter(unit='K')
    ctrlpars = Parameter('PID parameters',
                         StructOf(p=FloatRange(0, 1000), i=FloatRange(0, 1000), d=FloatRange(0, 1000)),
                         readonly=False)
    loop = Property('lakeshore loop', IntRange(1, 2), default=1)
    ramp = Parameter('ramp rate', FloatRange(min=0, max=100), unit='K/min', readonly=False)
    ramp_used = Parameter('whether ramp is used or not', BoolType(), readonly=False)
    setpoint = Parameter('setpoint', datatype=FloatRange, unit='K')

    def write_target(self, target):
        out = self.output_module
        if not self.control_active:
            if self.ramp_used:
                self.set_par(f'RAMP {self.loop}', 0, self.ramp)
                self.set_par(f'SETP {self.loop}', self.value)
                self.set_par(f'RAMP {self.loop}', 1, self.ramp)
            out.write_target(0)
            out.write_max_power(out.max_power)
            out.set_heater_mode(1)  # closed loop
        self.activate_output()
        self.start_state()  # start the convergence check
        out.set_range()
        self.set_par(f'SETP {self.loop}', target)
        return target

    def read_setpoint(self):
        setpoint = self.get_par(f'SETP?{self.loop}')
        status = self.get_par(f'RAMPST? {self.loop}')
        if status == 0:
            self.target = setpoint
        return setpoint

    def write_ctrlpars(self, ctrlpars):
        p, i, d = self.set_par(f'PID {self.loop}', ctrlpars['p'], ctrlpars['i'], ctrlpars['d'])
        return {'p': p, 'i': i, 'd': d}

    def read_ctrlpars(self):
        p, i, d = self.get_par(f'PID? {self.loop}')
        return {'p': p, 'i': i, 'd': d}

    def read_ramp(self):
        self.ramp_used, rate = self.get_par(f'RAMP? {self.loop}')
        return rate

    def write_ramp(self, ramp):
        self.ramp_used = True
        ramp = self.set_par(f'RAMP {self.loop}', self.ramp_used, ramp)[1]
        # if self.control:
        #     self.ramp = ramp
        #     self.write_target(self.target)
        #     return Done
        return ramp

    def write_ramp_used(self, ramp_used):
        ramp_used = self.set_par(f'RAMP {self.loop}', ramp_used, self.ramp)[0]
        if self.ramp_used and not ramp_used:
            self.write_target(self.target)
        return ramp_used

    def read_status(self):
        statuscode, statustext = self.status
        if self.ramp_used:
            if self.read_setpoint() == self.target:
                statuscode = self.Status.STABILIZING
            else:
                statuscode = self.Status.RAMPING
                statustext = 'ramping'
        if statuscode != ERROR:
            return Done
        if self.convergence_state.is_active:
            self.stop_machine((statuscode, statustext))
        return ERROR, statustext



