from time import sleep
import numpy as np
from scipy.interpolate import Akima1DInterpolator as Akima
from epics import PV

from slic.core.adjustable import Adjustable
from slic.core.adjustable import PVAdjustable
from slic.core.scanner.scanbackend import wait_for_all #, stop_all


n_unds = [
    8, 9, 10, 11, 12, 13,
    15, 16, 17, 18, 19, 20, 21, 22
]

und_names = [f"SATUN{n:02}-UIND030" for n in n_unds]

und_name_cal = "SATUN12-UIND030"


energies = [
    0.5457977557372193,
    0.5429495415906611,
    0.5401552739843691,
    0.537355500731781,
    0.534605067464375,
    0.5318574959825555,
    0.5291128804749312,
    0.5264167029579244,
    0.5237231992520139,
    0.5210645197823921,
]

k_values = [
    2.673,
    2.681888888888889,
    2.690777777777778,
    2.699666666666667,
    2.708555555555556,
    2.7174444444444443,
    2.7263333333333333,
    2.735222222222222,
    2.744111111111111,
    2.753,
]


#energies = [
#    0.549610621973092,
#    0.5479937271701031,
#    0.546395126105196,
#    0.544814685256097,
#    0.543218728674525,
#    0.541660979262817,
#    0.540074069530098,
#    0.538501771379368,
#    0.536939177980873,
#    0.535392458877575,
#    0.533842152732307,
#    0.5323060597843,
#    0.530767931907046,
#    0.5292440325790481,
#    0.5277161991879861,
#    0.526199323076255,
#    0.524699672767738,
#]

#k_values = [
#    2.673,
#    2.678,
#    2.683,
#    2.688,
#    2.693,
#    2.698,
#    2.703,
#    2.708,
#    2.713,
#    2.718,
#    2.723,
#    2.728,
#    2.733,
#    2.738,
#    2.743,
#    2.748,
#    2.753,
#]


#energies = [
#    1.1743862662276334,
#    1.1561455825400897,
#    1.1381577488951293,
#    1.1204686653636284,
#    1.1031371551015796,
#    1.0860714511221887,
#    1.069262314526579,
#    1.05279814063446,
#    1.0366146962391598,
#    1.0206897775380315,
#]

#k_values = [
#    1.5,
#    1.52222222,
#    1.54444444,
#    1.56666667,
#    1.58888889,
#    1.61111111,
#    1.63333333,
#    1.65555556,
#    1.67777778,
#    1.7,
#]



class Undulators(Adjustable):

    def __init__(self, energies=energies, k_values=k_values):
        super().__init__(name="Athos Undulators")
        self.adjs = {name: Undulator(name) for name in und_names}
        self.eV_to_K = make_calib(energies, k_values)
        self.K_to_eV = make_calib(k_values, energies)


    def set_target_value(self, value, hold=False):
        value /= 1000 # eV -> keV
        k = self.eV_to_K(value)
        if np.isnan(k):
            print("K is nan for", value)
            return
        print(f"{k} <- {value}")

        def change():
            # replace by set_all_target_values_and_wait when print not needed anymore
            tasks = []
            for name, a in self.adjs.items():
                print(f"{name} <- {k}")
                t = a.set_target_value(k, hold=False)
                tasks.append(t)
            wait_for_all(tasks)

        return self._as_task(change, hold=hold)


    def get_current_value(self):
        a = self.adjs[und_name_cal]
        k = a.get_current_value()
        energy = self.K_to_eV(k)
        return float(energy) * 1000 # keV -> eV


    def is_moving(self):
        return any(a.is_moving() for a in self.adjs)



class Undulator(PVAdjustable):

    def __init__(self, name, accuracy=0.001):
        pvname_setvalue = name + ":K_SET"
        pvname_readback = name + ":K_READ"
        super().__init__(pvname_setvalue, pvname_readback=pvname_readback, accuracy=accuracy, name=name)
        self.energy = PVAdjustable(name + ":FELPHOTENE")



def make_calib(x, y):
    x, y = list(zip(*sorted(zip(x, y))))
    return Akima(x, y)


def get_map(k_min=min(k_values), k_max=max(k_values)):
    und = Undulator(und_name_cal)
    start_k_value = und.get_current_value()

    k_values = np.linspace(k_min, k_max, 10)
    energies = []
    print()
    print("mapping = {")
    for k in k_values:
        und.set_target_value(k).wait()
        energy = und.energy.get_current_value()
        print(f"    {energy}: {k},")
        energies.append(energy)
    print("}")
    energies = np.array(energies)

    print()
    print("energies = [")
    for energy in energies:
        print(f"    {energy},")
    print("]")

    print()
    print("k_values = [")
    for k in k_values:
        print(f"    {k},")
    print("]")

    und.set_target_value(start_k_value).wait()
    return energies, k_values