# *****************************************************************************
#
# 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:
#   Markus Zolliker <markus.zolliker@psi.ch>
#
# *****************************************************************************

"""drivers for CCU4, the cryostat control unit at SINQ"""
import time
import math
from frappy.lib.enum import Enum
# the most common Frappy classes can be imported from frappy.core
from frappy.core import HasIO, Parameter, Command, Readable, Writable, Drivable, \
    Property, StringIO, BUSY, IDLE, WARN, ERROR, DISABLED, Attached
from frappy.datatypes import BoolType, EnumType, FloatRange, StructOf, \
    StatusType, IntRange, StringType, TupleOf
from frappy.errors import CommunicationFailedError
from frappy.states import HasStates, status_code, Retry


M = Enum(idle=0, opening=1, closing=2, opened=3, closed=5, no_motor=6)
A = Enum(disabled=0, manual=1, auto=2)


class CCU4IO(StringIO):
    """communication with CCU4"""
    # for completeness: (not needed, as it is the default)
    end_of_line = '\n'
    # on connect, we send 'cid' and expect a reply starting with 'CCU4'
    identification = [('cid', r'cid=CCU4.*')]


class CCU4Base(HasIO):
    ioClass = CCU4IO

    def command(self, **kwds):
        """send a command and get the response

        :param kwds: <parameter>=<value> for changing a parameter <parameter>=<type> for querying a parameter
        :returns: the (new) values of the parameters
        """
        types = {}
        cmds = []
        for key, value in kwds.items():
            if isinstance(value, type):
                types[key] = value
                cmds.append(key)
            elif isinstance(value, str):
                types[key] = str
                cmds.append(f'{key}={value}')
            else:
                types[key] = float
                cmds.append(f'{key}={value:g}')
        reply = self.io.communicate(' '.join(cmds)).split()
        if len(reply) != len(types):
            raise CommunicationFailedError('number of reply items does not match')
        result = []
        for (given, typ), res in zip(types.items(), reply):
            name, txtvalue = res.split('=')
            if given != name:
                raise CommunicationFailedError('result keys do not match given keys')
            result.append(typ(txtvalue))
        if len(kwds) == 1:
            return result[0]
        return result


class HeLevel(CCU4Base, Readable):
    """He Level channel of CCU4"""

    value = Parameter(unit='%')
    empty_length = Parameter('warm length when empty', FloatRange(0, 2000, unit='mm'),
                             readonly=False)
    full_length = Parameter('warm length when full', FloatRange(0, 2000, unit='mm'),
                            readonly=False)
    sample_rate = Parameter('sample rate', EnumType(slow=0, fast=1), readonly=False)

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

    # conversion of the code from the CCU4 parameter 'hsf'
    STATUS_MAP = {
        0: (IDLE, 'sensor ok'),
        1: (ERROR, 'sensor warm'),
        2: (ERROR, 'no sensor'),
        3: (ERROR, 'timeout'),
        4: (ERROR, 'not yet read'),
        5: (DISABLED, 'disabled'),
    }

    def read_value(self):
        return self.command(h=float)

    def read_status(self):
        return self.STATUS_MAP[int(self.command(hsf=int))]

    def read_sample_rate(self):
        value, self.empty_length, self.full_length = self.command(hf=int, hem=float, hfu=float)
        return value

    def write_sample_rate(self, value):
        return self.command(hf=value)

    def write_empty_length(self, value):
        return self.command(hem=value)

    def write_full_length(self, value):
        return self.command(hfu=value)


class Valve(CCU4Base, Writable):
    value = Parameter('relay state', BoolType())
    target = Parameter('relay target', BoolType())
    ioClass = CCU4IO
    STATE_MAP = {0: (0, (IDLE, 'off')),
                 1: (1, (IDLE, 'on')),
                 2: (0, (ERROR, 'no valve')),
                 3: (0, (WARN, 'timeout')),  # timeout in filling process (takes too long to fill)
                 4: (0, (WARN, 'timeout1')),  # timeout in filling process (takes too long to start)
                 5: (1, (IDLE, 'boost')),
                 }
    _open_command = None
    _close_command = None
    _query_state = None

    def write_target(self, target):
        if target:
            self.command(**self._open_command)
        else:
            self.command(**self._close_command)

    def read_status(self):
        state = int(self.command(**self._query_state))
        self.value, status = self.STATE_MAP[state]
        return status


class HeFillValve(Valve):
    _open_command = {'hcd': 1, 'hf': 1}
    _close_command = {'hcd': 0,  'hf': 0}
    _query_state = {'hv': int}


class N2FillValve(Valve):
    _open_command = {'nc': 1}
    _close_command = {'nc': 0}
    _query_state = {'nv': int}


class AuxValve(Valve):
    channel = Property('valve number', IntRange(1, 12))

    def initModule(self):
        self._open_command = {f'vc{self.channel}': 1}
        self._close_command = {f'vc{self.channel}': 0}
        self._query_state = {f'v{self.channel}': int}


class N2TempSensor(Readable):
    value = Parameter('LN2 T sensor', FloatRange(unit='K'), default=0)


class N2Level(CCU4Base, Readable):
    valve = Attached(Writable, mandatory=False)
    lower = Attached(Readable, mandatory=False)
    upper = Attached(Readable, mandatory=False)

    value = Parameter('vessel state', EnumType(empty=0, ok=1, full=2))
    status = Parameter(datatype=StatusType(Readable, 'DISABLED', 'BUSY'))
    mode = Parameter('auto mode', EnumType(A), readonly=False, default=A.manual)

    threshold = Parameter('threshold triggering start/stop filling',
                          FloatRange(unit='K'), readonly=False)
    cool_delay = Parameter('max. minutes needed to cool the lower sensor',
                           FloatRange(unit='s'), readonly=False)
    fill_timeout = Parameter('max. minutes needed to fill',
                             FloatRange(unit='s'), readonly=False)
    names = Property('''names of attached modules

                     configure members as empty strings to disable the creation
                     ''',
                     StructOf(valve=StringType(), lower=StringType(), upper=StringType()),
                     default={'valve': '$_valve', 'lower': '$_lower', 'upper': '$_upper'})

    # conversion of the code from the CCU4 parameter 'ns'
    STATUS_MAP = {
        0: (IDLE, 'sensor ok'),
        1: (ERROR, 'no sensor'),
        2: (ERROR, 'short circuit'),
        3: (ERROR, 'upside down'),
        4: (ERROR, 'sensor warm'),
        5: (WARN, 'empty'),
    }

    def initialReads(self):
        self.command(nav=1)  # tell CCU4 to activate LN2 sensor readings
        super().initialReads()

    def read_status(self):
        auto, nstate = self.command(na=int, ns=int)
        if not self.valve or not auto:
            if self.mode == A.auto:
                # no valve assigned
                self.mode = A.manual
        if self.mode == A.disabled:
            return DISABLED, ''
        status = self.STATUS_MAP[nstate]
        if status[0] // 100 != IDLE // 100:
            return status
        if self.mode == A.manual:
            return IDLE, ''
        vstatus = self.valve.status
        if vstatus[0] // 100 == WARN // 100:
            return ERROR, vstatus[1]
        if vstatus[0] // 100 != IDLE // 100:
            return vstatus
        if self.valve.value:
            return BUSY, 'filling'
        return IDLE, 'watching'

    def read_value(self):
        # read sensors
        lower, upper = self.command(nl=float, nu=float)
        if self.lower:
            self.lower.value = lower
        if self.upper:
            self.upper.value = upper
        if upper < self.threshold:
            return 'full'
        if lower < self.threshold:
            return 'ok'
        return 'empty'

    def write_mode(self, mode):
        if mode == A.auto:
            if self.isBusy():
                return mode
            # set to watching
            self.command(nc=3)
        else:
            # set to off
            self.command(nc=2)
        return mode

    def read_threshold(self):
        value, self.cool_delay, self.fill_timeout = self.command(nth=float, ntc=float, ntm=float)
        return value

    def write_threshold(self, value):
        return self.command(nth=value)

    def write_cool_delay(self, value):
        return self.command(ntc=value)

    def write_fill_timeout(self, value):
        return self.command(ntm=value)

    @Command()
    def fill(self):
        """start filling"""
        self.mode = A.auto
        self.command(nc=1)

    @Command()
    def stop(self):
        """stop filling"""
        if self.mode == A.auto:
            # set to watching
            self.command(nc=3)
        else:
            # set to off
            self.command(nc=0)


class FlowPressure(CCU4Base, Readable):
    value = Parameter(unit='mbar')
    mbar_offset = Parameter(unit='mbar', default=0.8, readonly=False)
    pollinterval = Parameter(default=0.25)

    def read_value(self):
        return self.filter(self.command(f=float)) - self.mbar_offset


class NeedleValve(HasStates, CCU4Base, Drivable):
    flow = Attached(Readable, mandatory=False)
    flow_pressure = Attached(Readable, mandatory=False)

    value = Parameter(unit='ln/min')
    target = Parameter(unit='ln/min')

    lnm_per_mbar = Parameter(unit='ln/min/mbar', default=0.6, readonly=False)
    use_pressure = Parameter('use flow from pressure', BoolType(),
                             default=False, readonly=False)
    motor_state = Parameter('motor_state', EnumType(M))
    tolerance = Parameter('tolerance', FloatRange(0), value=0.25, readonly=False)
    tolerance2 = Parameter('tolerance limit above 2 lnm', FloatRange(0), value=0.5, readonly=False)
    prop = Parameter('proportional term', FloatRange(unit='s/lnm'), readonly=False)
    deriv = Parameter('min progress time constant', FloatRange(unit='s'),
                      default=30, readonly=False)
    settle = Parameter('time within tolerance before getting quiet', FloatRange(unit='s'),
                       default=30, readonly=False)
    step_factor = Parameter('factor (no progress time) / (min step size)', FloatRange(), default=300)
    control_active = Parameter('control active flag', BoolType(), readonly=False)
    pollinterval = Parameter(default=1)
    _ref_time = 0
    _ref_dif = 0
    _last_cycle = 0
    _last_progress = 0
    _step = 0

    def initModule(self):
        super().initModule()
        if self.flow_pressure:
            self.flow_pressure.addCallback('value', self.update_flow_pressure)
        if self.flow:
            self.flow.addCallback('value', self.update_flow)
        self.write_tolerance(self.tolerance)

    def write_tolerance(self, tolerance):
        if hasattr(self.flow_pressure, 'tolerance'):
            self.flow_pressure.tolerance = tolerance / self.lnm_per_mbar
        if hasattr(self.flow, 'tolerance'):
            self.flow.tolerance = tolerance

    def read_use_pressure(self):
        if self.flow_pressure:
            if self.flow:
                return self.use_pressure
            return True
        return False

    def update_flow(self, value):
        if not self.use_pressure:
            self.value = value
            self.cycle_machine()

    def update_flow_pressure(self, value):
        if self.use_pressure:
            self.value = value * self.lnm_per_mbar
            self.cycle_machine()

    def write_target(self, value):
        self.start_machine(self.controlling, in_tol_time=0,
                           ref_time=0, ref_dif=0, prev_dif=0)

    @status_code(BUSY)
    def unblock_from_open(self, state):
        self.motor_state = self.command(fm=int)
        if self.motor_state == 'opened':
            self.command(mp=-60)
            return Retry
        if self.motor_state == 'closing':
            return Retry
        if self.motor_state == 'closed':
            if self.value > max(1, self.target):
                return Retry
            state.flow_before = self.value
            state.wiggle = 1
            state.start_wiggle = state.now
            self.command(mp=60)
            return self.unblock_open
        return self.approaching

    @status_code(BUSY)
    def unblock_open(self, state):
        self.motor_state = self.command(fm=int)
        if self.value < state.flow_before:
            state.flow_before_open = self.value
        elif self.value > state.flow_before + 1:
            state.wiggle = -state.wiggle / 2
            self.command(mp=state.wiggle)
            state.start_wiggle = state.now
            return self.unblock_close
        if self.motor_state == 'opening':
            return Retry
        if self.motor_state == 'idle':
            self.command(mp=state.wiggle)
            return Retry
        if self.motor_state == 'opened':
            if state.now < state.start_wiggle + 20:
                return Retry
            return self.final_status(ERROR, 'can not open')
        return self.controlling

    @status_code(BUSY)
    def unblock_close(self, state):
        self.motor_state = self.command(fm=int)
        if self.value > state.flow_before:
            state.flow_before_open = self.value
        elif self.value < state.flow_before - 1:
            if state.wiggle < self.prop * 2:
                return self.final_status(IDLE, '')
            state.wiggle = -state.wiggle / 2
            self.command(mp=state.wiggle)
            state.start_wiggle = state.now
            return self.unblock_open
        if self.motor_state == 'closing':
            return Retry
        if self.motor_state == 'idle':
            self.command(mp=state.wiggle)
            return Retry
        if self.motor_state == 'closed':
            if state.now < state.start_wiggle + 20:
                return Retry
            return self.final_status(ERROR, 'can not close')
        return self.final_status(WARN, 'unblock interrupted')

    def _tolerance(self):
        return min(self.tolerance * min(1, self.value / 2), self.tolerance2)

    @status_code(IDLE)
    def at_target(self, state):
        dif = self.target - self.value
        if abs(dif) > self._tolerance():
            state.in_tol_time = 0
            return self.unstable
        return Retry

    @status_code(IDLE, 'unstable')
    def unstable(self, state):
        return self.controlling(state)

    @status_code(BUSY)
    def controlling(self, state):
        delta = state.delta(0)
        dif = self.target - self.value
        difdif = dif - state.prev_dif
        state.prev_dif = dif
        self.motor_state = self.command(fm=int)
        if self.motor_state == 'closed':
            if dif < 0 or difdif < 0:
                return Retry
            return self.unblock_from_open
        elif self.motor_state == 'opened':  # trigger also when flow too high?
            if dif > 0 or difdif > 0:
                return Retry
            self.command(mp=-60)
            return self.unblock_from_open

        tolerance = self._tolerance()
        if abs(dif) < tolerance:
            state.in_tol_time += delta
            if state.in_tol_time > self.settle:
                return self.at_target
            return Retry
        expected_dif = state.ref_dif * math.exp((state.now - state.ref_time) / self.deriv)
        if abs(dif) < expected_dif:
            if abs(dif) < expected_dif / 1.25:
                state.ref_time = state.now
                state.ref_dif = abs(dif) * 1.25
            state.last_progress = state.now
            return Retry  # progress is fast enough
        state.ref_time = state.now
        state.ref_dif = abs(dif)
        state.step += dif * delta * self.prop
        if abs(state.step) < (state.now - state.last_progress) / self.step_factor:
            # wait until step size is big enough
            return Retry
        self.command(mp=state.step)
        return Retry