from datetime import datetime
from time import sleep, time
from tracemalloc import start
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 cristallina import Newport_large, attocube, attenuator
from smaract_device_def import smaract_mini_XYZ
from slic.devices.xoptics.aramis_attenuator import Attenuator
from tqdm import tqdm

attenuator = Attenuator("SAROP31-OATA150", description="Cristallina attenuator OATA150")

from slic.devices.xoptics.kb import 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 = False
pos = np.array([0, 0])

# parameters
n_same_holes = 3
attenuations = np.linspace(0.15, 0.35, 9)  # np.linspace(0.1,0.35,11)#np.linspace(0.1,0.35,6)#np.logspace(-3,0,num=7)
KBv_home = [1.3, 1.6]
KBh_home = [1.4, 1.7]

# Spacings
between_same_shots = 150
between_attenuations = 150
between_KB_settings = 500  # 500


### 23.02
def kb_settings_list(best, step, no_steps):
    for i in range(no_steps):
        if i == 0:
            l = []
        l.append([best[0] + i * step, best[1]])
    for i in range(no_steps):
        if i == 0:
            pass
        l.append([best[0], best[1] + i * step])
    return l


KBvs = kb_settings_list([1.595, 1.874], 0.05, 3)
KBhs = kb_settings_list([1.799000, 2.100000], 0.05, 3)


### 24.02
def kb_settings_list_new(best, step, no_steps_one_way):
    l = []
    first = np.array(best) - np.array([step * no_steps_one_way, step * no_steps_one_way])
    for i in range(no_steps_one_way * 2 + 1):
        l.append([first[0] + i * step, first[1] + i * step])
    return l


KBvs = kb_settings_list_new([1.595, 1.874], 0.0075, 3)
KBhs = kb_settings_list_new([1.799000, 2.100000], 0.0075, 3)


# Time estimates for total time calculation (time in seconds)
t_kb_change = 30
t_shot = 4
t_atten_change = 10

# Choose motor stage (set true for smaract_stage_xyz, false if attocubes)
smaract_stage = True
# Change to mm from um if smaract stage selected
if smaract_stage:
    between_same_shots = between_same_shots / 1000
    between_attenuations = between_attenuations / 1000
    between_KB_settings = between_KB_settings / 1000


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


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()
                sleep(1)
            else:
                kbVer.bend1.set_target_value(KBv_home[0]).wait()
                sleep(1)
                kbVer.bend2.set_target_value(KBv_home[1]).wait()
                sleep(1)

    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()
        sleep(1)
    else:
        kbVer.bend1.set_target_value(bender_1).wait()
        sleep(1)
        kbVer.bend2.set_target_value(bender_2).wait()
        sleep(1)


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:
        if smaract_stage:
            smaract_mini_XYZ.x.mvr(distance).wait()
        else:
            attocube.X.set_target_value(distance, relative=True).wait()
        pos = pos + np.array([distance, 0])
        return pos


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


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


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


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


def get_original_position(testing=testing_flag):
    if testing == True:
        original_position = pos
    else:
        if smaract_stage:
            original_position = [smaract_mini_XYZ.x.get_current_value(), smaract_mini_XYZ.y.get_current_value()]
        else:
            original_position = [attocube.X.get_current_value(), attocube.Y.get_current_value()]
    return original_position


def set_attenuation(value):
    attenuator.set_transmission(value)
    sleep(1)
    while attenuator.motors.any_is_moving():
        sleep(1)
    return value


def make_attenuations(attenuations, testing=testing_flag, pos=pos):
    original_position = get_original_position(testing=testing)

    # Make all attenuations
    for attenuation in attenuations:
        print(f"Setting attenuation to: {attenuation}")
        if testing:
            print("Testing: no attenuator change")
        else:
            set_attenuation(attenuation)
        print("Making same shots")
        make_same_shots(n_same_holes, pos=pos, testing=testing)
        pos = move_y_rel(between_attenuations, pos=pos, testing=testing)

    # Return back to where you started
    if testing == True:
        pos = original_position
    else:
        pos = move(original_position, testing=testing)
    return pos


def make_same_shots(n_same_holes, testing=testing_flag, pos=pos):
    original_position = get_original_position(testing=testing)

    # Make holes
    for shot in range(n_same_holes):
        sleep(1)
        shoot(pos=pos, testing=testing)
        pos = move_x_rel(between_same_shots, pos=pos, testing=testing)

    # Return back to where you started
    move(original_position, testing=testing)


def estimate_total_time(
    KBhs=KBhs,
    KBvs=KBvs,
    attenuations=attenuations,
    n_same_holes=n_same_holes,
    t_kb_change=t_kb_change,
    t_atten_change=t_atten_change,
    t_shot=t_shot,
):
    total_time = len(KBhs) * len(KBvs) * (t_kb_change + len(attenuations) * (t_atten_change + n_same_holes * t_shot))
    print(f"Total time estimate: {(total_time/60):.1f} minutes or {(total_time/60/60):.1f} hours")
    return total_time


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


def make_everything(KBvs, KBhs, attenuations, n_same_holes, testing=testing_flag, pos=pos, do_not_move_benders=True):
    # The loop to make inprints
    for i, KBv in enumerate(tqdm(KBvs)):
        change_benders(KBv[0], KBv[1], KB="v", do_not_move_benders=do_not_move_benders)

        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, testing=testing)

            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, testing=testing)

        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, testing=testing)

        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, testing=testing)

    print("Inprints are done")
    set_attenuation(1e-6)
    status_pv = PV("SF-OP:ESC-MSG:STATUS")
    status_pv.put(0)

    # 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


def make_z_scan(z_list, n_same_holes, attenuations, vertical_spacing=0.150):
    original_position = get_original_position(testing=False)
    for attenuation in attenuations:
        set_attenuation(attenuation)
        for z in z_list:
            Newport_large.mv(z).wait()
            sleep(1)
            make_same_shots(n_same_holes, testing=False)
            move_y_rel(vertical_spacing, testing=False)
        move(original_position, testing=False)
        sleep(1)
        move_x_rel(1, testing=False)