#  -*- coding: utf-8 -*-
# *****************************************************************************
# 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>
# *****************************************************************************
"""vector field"""

import math
from frappy.core import Drivable, Done, BUSY, IDLE, ERROR, Parameter, TupleOf, ArrayOf, FloatRange
from frappy.errors import RangeError
from frappy_psi.vector import Vector
from frappy.states import HasStates, Retry, status_code


DECREASE = 1
INCREASE = 2


class VectorField(HasStates, Vector, Drivable):
    sphere_radius = Parameter('max. sphere', datatype=FloatRange(0, 0.7, unit='T'), readonly=True, default=0.6)
    cylinders = Parameter('allowed cylinders (list of radius and height)',
                          datatype=ArrayOf(TupleOf(FloatRange(0, 0.6, unit='T'), FloatRange(0, 5.2, unit='T')), 1, 9),
                          readonly=True, default=((0.23, 5.2), (0.45, 0.8)))

    def initModule(self):
        super().initModule()
        # override check_limits of the components with a check for restrictions on the vector
        for idx, component in enumerate(self.components):

            def outer_check(target, vector=self, i=idx, inner_check=component.check_target):
                inner_check(target)
                value = [c.value - math.copysign(c.tolerance, c.value)
                         for c in vector.components]
                value[i] = target
                vector._check_limits(value)

            component.check_target = outer_check

    def merge_status(self):
        return self.status

    def _check_limits(self, value):
        """check if value is within one of the safe shapes"""
        if sum((v ** 2 for v in value)) <= self.sphere_radius ** 2:
            return
        for r, h in self.cylinders:
            if sum(v ** 2 for v in value[0:2]) <= r ** 2 and abs(value[2]) <= h:
                return
        raise RangeError('vector %s does not fit in any limiting shape' % repr(value))

    def write_target(self, target):
        """initiate target change"""
        # check limits  first
        for component_target, component in zip(target, self.components):
            # check against limits of individual components
            component.check_target(component_target, vector=None)  # no outer check here!
        self._check_limits(target)
        for component_target, component in zip(target, self.components):
            if component_target * component.value < 0:
                # change sign: drive to zero first
                component_target = 0
            if abs(component_target) > abs(component.value):
                continue  # do not drive yet
            component.write_target(component_target)
        self.start_machine(self.ramp_down, target=target)
        return target

    @status_code(BUSY)
    def ramp_down(self, state):
        for target, component in zip(state.target, self.components):
            if component.isDriving():
                return Retry()
        for target, component in zip(state.target, self.components):
            component.write_target(target)
        return self.final_ramp

    @status_code(BUSY)
    def final_ramp(self, state):
        for component in self.components:
            if component.isDriving():
                return Retry()
        return self.final_status()