from pyepics import PV
from slic.core.adjustable import Adjustable, PVAdjustable
from slic.utils import json_load
from model import parallel2gap


UND_PERIOD = 38.0


def check_phase(phase):
    assert -180 <= phase <= 180 #TODO: modulo to be in the correct range?

def convert_phase_to_shift(phase):
    ratio = UND_PERIOD / 360.0
    return phase * ratio / 2



class UndPhases(Adjustable):
    """
    Set of several UndPhase objects
    allows to set the same phase to all undulators
    """

    def __init__(self, ID, params, und_names=None, **kwargs):
        super().__init__(ID, **kwargs)
        self.params = params

        if und_names is None:
            und_names = params.keys()

        SUFFIX = "-UIND030:"
        self.phases = [UndPhase(i + SUFFIX, params[i]) for i in und_names]


    def set_target_value(self, value):
        tasks = [p.set_target_value(value) for p in self.phases]
        for t in tasks:
            t.wait()



class UndPhase(Adjustable):
    """
    Combination of UndShift, UndRadial and UndTotalK
    allows to set the phase of one undulator
    """

    def __init__(self, ID, params, **kwargs):
        super().__init__(ID, **kwargs)
        self.params = params
        self.shift = UndShift(ID)
        self.radial = UndRadial(ID)
        self.totalk = UndTotalK(ID)


    def get_current_value(self):
        raise NotImplementedError #TODO: how to do that?


    def set_target_value(self, value):
        phase = value
        shift = convert_phase_to_shift(phase)

        k = self.totalk.get()
        radial = parallel2gap(k, phase, self.params)
        radial = round(radial, 4) #TODO: why?

        self.shift.set_target_value(shift).wait()
        self.radial.set_target_value(radial).wait()


    def is_moving(self):
        return self.shift.is_moving() or self.radial.is_moving()



class UndShiftRadialBase(PVAdjustable): #TODO: better name?
    """
    Base class with functionality shared between UndShift and UndRadial
    """

    def __init__(self, ID, accuracy=0.001):
        super().__init__(ID + "-SET", ID + "-TL", accuracy=accuracy)
        self.pv_on = PV(ID + "-ON")
        self.pv_go = PV(ID + "-GO")

    def set_target_value(self, value):
        t = super().set_target_value(value)
        sleep(0.3)
        self.pv_on.put(1)
        sleep(0.3)
        self.pv_go.put(1)
        t.wait()


class UndShift(UndShiftRadialBase):

    def __init__(self, ID):
        super().__init__(ID + "SHIFT")


class UndRadial(UndShiftRadialBase):

    def __init__(self, ID):
        super().__init__(ID + "RADIAL")


class UndTotalK:
    """
    Helper class to get the total K from set value and taper
    """

    def __init__(self, ID):
        self.pv_kset = PV(ID + "K_SET")
        self.pv_ktaperset = PV(ID + "K_TAPER_SET")


    def get(self):
        k = self.pv_kset.get()
        ktaper = self.pv_ktaperset.get()

        #TODO why?
        k = round(k, 4)
        ktaper = round(ktaper, 4)

        return k + ktaper





if __name__ == "__main__":
    phase = 123
    check_phase(phase)

    params = json_load("UE38_all_parallel_parameters.json")

    ups = UndPhases("SATUN-PHASES", params)
    ups.set_target_value(phase).wait()