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.9733,
    7: 2.9733,
    8: 2.9733,
    9: 2.9733,
    10: 2.9733,
    11: 2.9733,
    12: 2.9733,
    13: 2.9633,

    15: 3.3349,
    16: 3.3282,
    17: 3.2344,
    18: 3.3517,
    19: 3.3169,
    20: 3.3407,
    21: 3.3477,
    22: 3.2467
}


color2 = {
    6: 2.7249,
    7: 2.818,
    8: 2.766,
    9: 2.7539,
    10: 2.8251,
    11: 2.7392,
    12: 2.7545,
    13: 2.7353,
 
    15: 3.4958,
    16: 3.4958,
    17: 3.4958,
    18: 3.4958,
    19: 3.4958,
    20: 3.4958,
    21: 3.4958,
    22: 3.4808
}


colors_both = {
    6: 2.9733,
    7: 2.9733,
    8: 2.9733,
    9: 2.9733,
    10: 2.9733,
    11: 2.9733,
    12: 2.9733,
    13: 2.9633,
 
    15: 3.4958,
    16: 3.4958,
    17: 3.4958,
    18: 3.4958,
    19: 3.4958,
    20: 3.4958,
    21: 3.4958,
    22: 3.4808
}


kill_both = {
    6: 2.7249,
    7: 2.818,
    8: 2.766,
    9: 2.7539,
    10: 2.8251,
    11: 2.7392,
    12: 2.7545,
    13: 2.7353,

    15: 3.3349,
    16: 3.3282,
    17: 3.2344,
    18: 3.3517,
    19: 3.3169,
    20: 3.3407,
    21: 3.3477,
    22: 3.2467
}


#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 kill_both_colors():
    put_color(kill_both)


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