from time import sleep
from epics import caget, caput
from slic.utils import as_shortcut


KSET_FMT     = "SATUN{:02}-UIND030:K_SET"
RADIALGO_FMT = "SATUN{:02}-UIND030:RADIAL-GO"
#PHASESET_FMT = "SATUN{:02}-CHIC:PHASE"
N_UND_CHIC = 14
N_UNDS = list(range(6, 22+1))
N_UNDS.remove(N_UND_CHIC)


#FMT_G1 = "SATUN{:02}-CHIC:G1"
#FMT_G2 = "SATUN{:02}-CHIC:G2"
#FMT_GD = "SATUN{:02}-UDLY060:GAP-D-SET"
#FMT_GU = "SATUN{:02}-UDLY060:GAP-U-SET"


#ksets = {n: caget(KSET_FMT.format(n)) for n in n_unds}
#ksets


color1 = {
    6: 2.0618975071726737,
    7: 2.0618998489503135,
    8: 2.061926961690368,
    9: 2.0618991835864264,
    10: 2.061919185767039,
    11: 2.0619176190165964,
    12: 2.0618790520140906,
    13: 2.052086641409292,

    15: 1.1842616466870595,
    16: 1.2842673674062317,
    17: 1.3842579267233892,
    18: 1.484294303739585,
    19: 1.6843694577800838,
    20: 1.584281331329752,
    21: 1.2842820179639072,
    22: 1.3842844240055432
}


color2 = {
    6: 1.0843075706646517,
    7: 1.1843037624109625,
    8: 1.3842958393058413,
    9: 1.4843045803089876,
    10: 1.5843035457018615,
    11: 1.8400123990046287,
    12: 1.6843112146639572,
    13: 1.7742896734281286,
 
    15: 2.124329777953559,
    16: 2.124339587297864,
    17: 2.124344707174983,
    18: 2.1243781537048076,
    19: 2.1243433061478303,
    20: 2.124360736030826,
    21: 2.124337879383864,
    22: 2.11434429180146
}


colors_both = {
    6: 2.061898658676767,
    7: 2.0619010015723753,
    8: 2.061836025091557,
    9: 2.0619014907947957,
    10: 2.0619007931739812,
    11: 2.0619164451260863,
    12: 2.061841664062066,
    13: 2.0520064916331067,
 
    15: 2.124321241942164,
    16: 2.1243286922141595,
    17: 2.124315652168558,
    18: 2.124332151134097,
    19: 2.124354200526172,
    20: 2.1243170649754792,
    21: 2.1243234280187435,
    22: 2.11434068217272
}


#colors_reset = colors_both
#phases_reset = {
#    6: 6.271830001886656,
#    7: 4.2,
#    8: 11.700000000000001,
#    9: 4.2,
#    10: 333.76497589650967,
#    11: 5.9,
#    12: 5.800000000000001,

#    13: 303.1292805364285,
#    15: 26.939664715067693,
#    16: 131.94293079970862,
#    17: 276.5058000202555,
#    18: 287.6856327168946,
#    19: 207.38583858800666,
#    20: 128.34546121694189,
#    21: 61.3287322695885
#}



def put_color(cols):
    for n, k in cols.items():
        pv = KSET_FMT.format(n)
        caput(pv, k)
    # press the go buttons ...
    for n in cols:
        pv = RADIALGO_FMT.format(n)
        caput(pv, 1)
    print("color done")

#def put_phases(phs):
#    for n, k in phs.items():
#        pv = PHASESET_FMT.format(n)
#        caput(pv, k)
#    print("phases done")



@as_shortcut
def enable_color1():
    put_color(color1)

@as_shortcut
def enable_color2():
    put_color(color2)

@as_shortcut
def enable_colors_both():
    put_color(colors_both)

#@as_shortcut
#def reset_undulators():
#    put_color(colors_reset)
#    sleep(2)
#    put_phases(phases_reset)