from datetime import datetime
from time import sleep, time
import numpy as np
#from epics import PV
#from slic.utils import nice_arange
#from slic.devices.general.motor import Motor
import matplotlib.pyplot as plt
import epics

from slic.devices.xoptics.aramis_attenuator import Attenuator
attenuator = Attenuator("SAROP31-OATA150", description="Cristallina attenuator OATA150")

from slic.devices.xoptics.kb import KBBase,KBHor,KBVer
kbHor = KBHor(
    "SAROP31-OKBH154",
    description="Cristallina horizontal KB mirror"
)

kbVer = KBVer(
   "SAROP31-OKBV153",
   description="Cristallina vertical KB mirror"
)

do_not_move_benders = True
testing_flag = True
pos = np.array([0,0])

# parameters
n_same_holes = 7
attenuations = np.logspace(0,-6,num=7)
KBv_home = [1.5,1.7]
KBh_home = [1.6,1.8]

# Spacings
between_same_shots = 100
between_attenuations = 150
between_KB_settings = 500
KBvs = [[1,2],[1,3],[1,2],[1,3],[1,2]]
KBhs = [[1,2],[1,2],[1,2],[1,3],[1,2]]

# Time estimates for total time calculation
t_kb_change = 30
t_shot = 5
t_atten_change = 20

def shoot(pos=pos,testing=testing_flag):
    if testing:
        print(f'Shot at: {pos}')
        return pos
    else:
        epics.caput("SAR-CCTA-ESC:seq0Ctrl-Start-I",1)


def change_benders(bender_1,bender_2,KB = None,do_not_move_benders = do_not_move_benders):
    check_KB_value(KB)
    
    current_values = get_bender_values(KB)
    if current_values[0] > bender_1 or current_values[1] > bender_2:
        print('Unbending first because of backlash')
        if do_not_move_benders != True:
            # Move home first
            if KB == 'h' or KB == 'H':
                kbHor.bend1.set_target_value(KBh_home[0]).wait()
                sleep(1)
                kbHor.bend2.set_target_value(KBh_home[1]).wait()
            else:
                kbVer.bend1.set_target_value(KBv_home[0]).wait()
                sleep(1)
                kbVer.bend2.set_target_value(KBv_home[1]).wait()
    

    if do_not_move_benders:
        print(f'Bender 1 to: {bender_1}')
        print(f'Bender 2 to: {bender_2}')
        return 

    # Move to the new position
    print(f'Changing {KB} bender to: [{bender_1} , {bender_2}]')
    if KB == 'h' or KB == 'H':
        kbHor.bend1.set_target_value(bender_1).wait()
        sleep(1)
        kbHor.bend2.set_target_value(bender_2).wait()
    else:
        kbVer.bend1.set_target_value(bender_1).wait()
        sleep(1)
        kbVer.bend2.set_target_value(bender_2).wait()
    

def check_KB_value(KB):
    if KB not in ['H','h','V','v']:
        raise KeyError(f"KB can only be horizontal (H) or vertical (V), not {KB}")    

def get_bender_values(KB=None):
    check_KB_value(KB)
    if KB == 'h' or KB == 'H':
        return [kbHor.bend1.get(),kbHor.bend2.get()]
    else:
        return [kbVer.bend1.get(),kbVer.bend2.get()]


def move_x_rel(distance,testing=testing_flag,pos=pos):
    if testing == True:
        pos = pos + np.array([distance,0])
        return pos
    else:
        attocube.X.set_target_value(distance, relative=True).wait()

def move_y_rel(distance,testing=testing_flag,pos=pos):
    if testing == True:
        pos = pos + np.array([0,distance])
        return pos
    else:
        attocube.Y.set_target_value(distance, relative=True).wait()

def move_x(value,testing=testing_flag,pos=pos):
    if testing == True:
        pos[0]=value
        return pos
    else:
        attocube.X.set_target_value(value, relative=False).wait()

def move_y(value,testing=testing_flag,pos=pos):
    if testing == True:
        pos[1]=value
        return pos
    else:
        attocube.Y.set_target_value(value, relative=False).wait()


def move(target_pos,testing=True,pos=pos):
    if testing == True:
        pos = target_pos
        return pos
    else:
        attocube.X.set_target_value(target_pos[0]).wait()
        attocube.Y.set_target_value(target_pos[1]).wait()


def make_attenuations(attenuations,testing=testing_flag,pos=pos):
    if testing == True:
        original_position = pos
    else:
        original_position = [attocube.X.get_current_value(),attocube.Y.get_current_value()]
    
    # Make all attenuations
    for attenuation in attenuations:
        print(f'Setting attenuation to: {attenuation}')
        if testing_flag != True:
            attenuator.trans1st(attenuation).wait()
        print('Making same shots')
        make_same_shots(n_same_holes,pos=pos)
        pos = move_y_rel(between_attenuations,pos=pos)
        
    # Return back to where you started
    if testing == True:
        pos = original_position
    else:
        attocube.X.set_target_value(original_position, relative=False)
    return pos

def make_same_shots(n_same_holes,testing=testing_flag,pos=pos):
    if testing == True:
        original_position = pos
    else:
        original_position = [attocube.X.get_current_value(),attocube.Y.get_current_value()]
    
    # Make holes
    for shot in range(n_same_holes):
        shoot(pos=pos)
        pos = move_x_rel(between_same_shots,pos=pos)
        
    # Return back to where you started
    move(original_position)
    

# Estimate total time
total_time = len(KBhs)*len(KBvs)*(t_kb_change+len(attenuations)*(t_atten_change+n_same_holes*t_shot) )

# Get the starting x-position
if testing_flag == True:
    starting_x_pos = pos[0]
else:
    starting_x_pos = attocube.X.get_current_value()

# The actual loop to make inprints
for i,KBv in enumerate(KBvs):
    change_benders(KBv[0],KBv[1],KB = 'v')

    for ind,KBh in enumerate(KBhs):
        change_benders(KBh[0],KBh[1],KB = 'h')
        
        print(f'Progress so far: KBv loop: {i+1}/{len(KBvs)}. KBh loop:{ind+1}/{len(KBhs)}')
        make_attenuations(attenuations,pos=pos)

        print(f'Moving to a new KBh setting')
        # Move to the last shot of the same shot + the spacing between KB settings 
        pos = move_x_rel(between_KB_settings+between_same_shots*(n_same_holes-1),pos=pos)

    print('KBh set done, returning to starting_x_pos')
    # Move to the last shot of the same shot + the spacing between KB settings 

    pos = move_x(starting_x_pos,pos=pos)

    print('#################################################################################')
    print('Moving to a new KBv setting')
    # Move to the last shot of the same shot + the spacing between KB settings 

    pos = move_y_rel(between_KB_settings+between_attenuations*(len(attenuations)-1),pos=pos)


print('Inprints are done')
print(f'Total time estimate: {(total_time/60):.1f} minutes or {(total_time/60/60):.1f} hours')       

# To do:
# Fix movement of the attocubes in real time
# Load lut files into the controller and check that x is x and y is y