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.core.device import Device, SimpleDevice
from slic.devices.general.motor import Motor
from slic.utils.printing import printable_table

from constraints import ExtraConstraint

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


class QSpace3D(Device):

    def __init__(self, ID, mu, chi, phi, nu, 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,
            mu = mu,
            chi = chi,
            phi = phi,
            nu = nu
        )

        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)

    def add_orientation(self, *args, **kwargs):
        self.dc.ub.add_orientation(*args, **kwargs)

    def del_orientation(self, *args):
        self.dc.ub.del_orientation(*args)

    def add_reflection(self, *args, **kwargs):
        self.dc.ub.add_reflection(*args, **kwargs)

    def del_reflection(self, *args):
        self.dc.ub.del_reflection(*args)

    def calc_ub(self, *args):
        self.dc.ub.calc_ub(*args) # not sure whether this needs to be called explicitly?

    def fit_ub(self, *args, **kwargs):
        self.dc.ub.fit_ub(*args, **kwargs) 

    @property
    def UB(self):
        self.dc.ub.UB

    def clear_constraints(self):
        self.dc.cons.clear()

    def set_constraints(self, **kwargs):
        self.clear_constraints()
        for name, value in kwargs.items():
            setattr(self.dc.cons, name, value)

    def get_position(self, *args, extra_cons):
        wl = self.wavelength.get_current_value()
        res = self.dc.hkl.get_position(*args, wl)
        if extra_cons:
            res = [r for r in res if extra_cons(_flatten_get_position_result(r))]
        return res

    def hard_constraints(x):
        if x["betain"]<0:
            return False
        if x["betaout"]<0: 
            return False
        if x["nu"]<0: 
            return False
        if x["nu"]>180:    
            return False
        if x["mu"]<0:    
            return False
        if x["mu"]>90:    
            return False
        if x["chi"]<-30:
            return False
        if x["chi"]>30:
            return False
        if x["phi"]<-60:
            return False
        if x["phi"]>60:
            return False
        return True

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

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

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

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

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

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

    def _set_angles(self, mu, chi, phi, nu):
        ms = self.motors
        t_mu = ms.mu.set_target_value(90+mu)
        t_chi = ms.chi.set_target_value(chi-180)
        t_phi = ms.phi.set_target_value(phi)
        t_nu = ms.nu.set_target_value(180+nu) # We may need a correct offset value
        # wait for all motors to arrive
        tasks = (t_mu, t_chi, t_phi, t_nu)
        for t in tasks:
            t.wait()

#    def _set_angles(self, mu, chi, phi, nu):
#        ms = self.motors
#        t_mu = ms.mu.set_target_value(90-mu)
#        t_chi = ms.chi.set_target_value(chi)
#        t_phi = ms.phi.set_target_value(phi)
#        t_nu = ms.nu.set_target_value(nu)
        # wait for all motors to arrive
#        tasks = (t_mu, t_chi, t_phi, t_nu)
#        for t in tasks:
#            t.wait()

#    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.mu, pos.chi, pos.phi, pos.nu

    def _calc_angles(self, h, k, l, extra_cons):
        wl = self.get_wavelength()
        pos, _virtual_angles = next(iter(
            self.get_position(h, k, l, extra_cons=extra_cons)
        ))
        return pos.mu, pos.chi, pos.phi, pos.nu

def _flatten_get_position_result(x):
    d0, d1 = x
    d0 = d0.asdict
    return dict(**d0, **d1)



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, extra_cons):
        # get all current indices
        hkl = self._get_hkl()
        hkl = list(hkl)
        i = self._get_index()
        # insert the target value into the right spot
        hkl[i] = value
        self._set_hkl(*hkl,extra_cons)

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


    # some helpful things:

    def get_position(self, *args, extra_cons):
        wl = self.wavelength.get_current_value()
        res = self.dc.hkl.get_position(*args, wl)
        if extra_cons:
            res = [r for r in res if extra_cons(_flatten_get_position_result(r))]
        return res

    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, extra_cons):
        angles = self._calc_angles(h, k, l, extra_cons)
        self._set_angles(*angles)

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

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

    def _calc_hkl(self, mu, chi, phi, nu):
        wl = self.get_wavelength()
        pos = Position(mu=mu, chi=chi, phi=phi, nu=nu)
        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.mu, pos.chi, pos.phi, pos.nu

    def _calc_angles(self, h, k, l, extra_cons):
        wl = self.get_wavelength()
        pos, _virtual_angles = next(iter(
            self.get_position(h, k, l, extra_cons)
        ))
        return pos.mu, pos.chi, pos.phi, pos.nu


h = 6.62607015e-34 # J s
c = 299792458 # m / s
e = 1.60217663e-19 # C
f = h * c / e * 1e-3 * 1e9 *1e4# E [eV] = f / lambda [nm]


class Wavelength(Adjustable):

    def __init__(self, energy):
        self.energy = energy
#        assert self.energy.units == "keV" # otherwise conversion is wrong
        super().__init__(energy.ID + "_WL", name=energy.name + " as Wavelength", units="nm")

    def get_current_value(self):
        E = self.energy.get_current_value()
        return f / E

    def set_target_value(self, value):
        E = f / value
        return self.energy.set_target_value(E)

    def is_moving(self):
        return self.energy.is_moving()



class HistoryDummy(Adjustable):

    def __init__(self, ID, initial_value=None, **kwargs):
        super().__init__(ID, **kwargs)
        self.history = [] if initial_value is None else [initial_value]

    @classmethod
    def init_from(cls, adj):
        ID = "HD:" + adj.ID
        initial_value = adj.get_current_value()
        name = cls.__name__ + " for " + adj.name
        units = adj.units
        return cls(ID, initial_value=initial_value, name=name, units=units)

    def get_current_value(self):
        try:
            return self.history[-1]
        except IndexError:
            return None

    def set_target_value(self, value):
        self.history.append(value)

    def is_moving(self):
        return False


def print_history(*adjs, enumerate_lines=True, make_legend=True):
    labels = []
    histories = []
    for a in adjs:
        try:
            h = a.history
        except AttributeError:
            continue
        labels.append(repr(a))
        histories.append(h)
    print(printable_table(histories, labels, enumerate_lines=enumerate_lines, make_legend=make_legend))



## these point to the different motors
#mu = Motor("SATES30-RIXS:MOT_SRY.VAL")
#chi = Motor("SATES30-RIXS:MOT_SRZ.VAL")
#phi = Motor("SATES30-RIXS:MOT_SRX.VAL")
#nu = Motor("SATES30-RIXS:MOT_DRY.VAL")

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

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

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