from types import SimpleNamespace

from diffcalc.hkl.calc import HklCalculation
from diffcalc.hkl.constraints import Constraints
from diffcalc.hkl.geometry import Position
from diffcalc.ub.calc import UBCalculation

from slic.core.adjustable import Adjustable, PVAdjustable
from slic.devices.general.motor import Motor
from slic.devices.device import Device
from slic.devices.simpledevice import SimpleDevice


INDICES = {
    "h": 0,
    "k": 1,
    "l": 2
}


class QSpace3D(Device):

    def __init__(self, ID, eta, chi, phi, wavelength, **kwargs): # if the diffcalc objects need some initial parameters, add them here and fill them in below...
        super().__init__(ID, **kwargs)

        # collect the motors in a device for nicer printing
        self.motors = motors = SimpleDevice(
            ID,
            eta = eta,
            chi = chi,
            phi = phi
        )

        self.wavelength = wavelength

        # some diffcalc objects ...
        ub = UBCalculation()
        cons = Constraints()
        hkl = HklCalculation(ub, cons)

        # ... collected in a namespace
        self.dc = dc = SimpleNamespace(
            ub = ub,
            cons = cons,
            hkl = hkl
        )

        # and the individual coordinates:
        self.h = QSpace1D(ID + "-H", "h", dc, motors, wavelength)
        self.k = QSpace1D(ID + "-K", "k", dc, motors, wavelength)
        self.l = QSpace1D(ID + "-L", "l", dc, motors, wavelength)


    # it might be nice (but optional) to add some shortcuts like this:
    def set_lattice(self, *args, **kwargs):
        self.dc.ub.set_lattice(*args, **kwargs)
        self.dc.ub.calc_ub() # not sure whether this needs to be called explicitly?

    # if I understand the examples/code correctly, then some more method calls are mandatory?
    # those should probably all get shortcuts...



class QSpace1D(Adjustable):

    def __init__(self, ID, index, dc, motors, wavelength):
        super().__init__(ID)
        self.index = index
        self.dc = dc
        self.motors = motors
        self.wavelength = wavelength


    # the following three methods are mandatory:

    def get_current_value(self):
        hkl = self._get_hkl()
        i = self._get_index()
        return hkl[i]

    def set_target_value(self, value):
        # get all current indices
        hkl = self._get_hkl()
        i = self._get_index()
        # insert the target value into the right spot
        hkl[i] = value
        self._set_hkl(self, *hkl)

    def is_moving(self):
        ms = self.motors
        return ms.eta.moving or ms.chi.moving or ms.phi.moving


    # some helpful things:

    def get_wavelength(self):
        return self.wavelength.get_current_value()

    def _get_index(self):
        i = self.index
        if isinstance(i, str):
            i = i.casefold()
            i = INDICES.get(i, i)
        if i not in INDICES.values():
            allowed = INDICES.keys() | INDICES.values()
            allowed = sorted(allowed, key=str)
            raise ValueError(f"index must be from {allowed} but is {repr(i)}")
        return i

    def _get_hkl(self):
        angles = self._get_angles()
        return self._calc_hkl(*angles)

    def _set_hkl(self, h, k, l):
        angles = self._calc_angles(self, h, k, l)
        self._set_angles(self, *angles)

    def _get_angles(self):
        ms = self.motors
        eta = ms.eta.get_current_value()
        chi = ms.chi.get_current_value()
        phi = ms.phi.get_current_value()
        return eta, chi, phi

    def _set_angles(self, eta, chi, phi):
        ms = self.motors
        t_eta = ms.eta.set_target_value(eta)
        t_chi = ms.chi.set_target_value(chi)
        t_phi = ms.phi.set_target_value(phi)
        # wait for all motors to arrive
        tasks = (t_eta, t_chi, t_phi)
        for t in tasks:
            t.wait()

    def _calc_hkl(self, eta, chi, phi):
        wl = self.get_wavelength()
        pos = Position(eta=eta, chi=chi, phi=phi)
        hkl = self.dc.hkl.get_hkl(pos, wl)
        return hkl

    def _calc_angles(self, h, k, l):
        wl = self.get_wavelength()
        pos, _virtual_angles = next(iter(
            self.dc.hkl.get_position(h, k, l, wl)
        ))
        return pos.eta, pos.chi, pos.phi





# these point to the different motors
eta = Motor("SOMETHING:ETA")
chi = Motor("SOMETHING:CHI")
phi = Motor("SOMETHING:PHI")

# and this to the machine wavelength (maybe this needs to be calculated from the energy? then we should add a Wavelength wrapper...)
wl = PVAdjustable("MACHINE:WAVELENGTH")

# put it all together:
q = QSpace3D("SOMETHING:Q", eta, chi, phi, wl)

## in ipython
#q.set_lattice("SiO2", 4.913, 5.405)
#q.set_target_value(1.5)