from datetime import datetime
from time import sleep, time
from tracemalloc import start
import numpy as np
from collections import defaultdict

# 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
from standa import standa

from slic.devices.timing.events import CTASequencer

cta = CTASequencer("SAR-CCTA-ESC")

# setup logging
import logging

logger = logging.getLogger("slic.hole_drilling")
logger.setLevel(logging.INFO)

# formatter = logging.Formatter(fmt='%(asctime)s %(levelname)-8s %(message)s',
#                              datefmt='%Y-%m-%d %H:%M:%S')

# create file handler which logs
try:
    fh = logging.FileHandler("/sf/cristallina/applications/slic/cristallina/log/hole_drilling.log")
    fh.setLevel(logging.DEBUG)
    # fh.setFormatter(formatter)
    logger.addHandler(fh)
except PermissionError as e:
    logger.warning("Cannot write log file.")
    logger.warning(e)


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(n_shots=1):
    epics.caput("SAR-CCTA-ESC:seq0Ctrl-Cycles-I", n_shots)
    sleep(0.1)
    epics.caput("SAR-CCTA-ESC:seq0Ctrl-Start-I", 1)


def move_x_rel(distance):
    standa.x.mvr(distance)


def move_y_rel(distance):
    standa.z.mvr(distance)


def move_absolute(x, y):
    standa.x.mv(x).wait()
    standa.z.mv(y).wait()


def make_holes(quantity, spacing=0.25, wait=1):
    start_x = standa.x.get()
    start_y = standa.z.get()

    print(f"Starting postion:\n {standa}")

    for i in range(quantity):
        print(f"Shot {i}")
        shoot()
        sleep(wait)
        move_x_rel(spacing)
    print("Done")
    print("Returning to the start")
    standa.x.mv(start_x)
    standa.z.mv(start_y)


def drill_holes_grid(x_start, x_stop, n_cols, y_start, y_stop, n_rows, wait_time=0.5):
    # make grid
    x_pos = np.linspace(x_start, x_stop, n_cols)
    x_spacing = (x_stop - x_start) / (n_cols - 1)

    y_pos = np.linspace(y_start, y_stop, n_rows)
    y_spacing = (y_stop - y_start) / (n_rows - 1)

    # TODO: add settings we want to change per row
    row_settings = defaultdict(lambda: {"n_pulses": 100})
    row_settings[0] = {"n_pulses": 500}
    row_settings[1] = {"n_pulses": 100}
    row_settings[2] = {"n_pulses": 10}
    row_settings[3] = {"n_pulses": 1}

    x_start = standa.x.get()
    y_start = standa.z.get()

    logger.info(
        f"Begin hole pattern at x:{x_start:.5f}, y:{y_start:.5f}, with transmission: {attenuator.get_transmission():.5f}."
    )

    for row, y in enumerate(y_pos):
        settings = row_settings[row]
        n_pulses = settings["n_pulses"]

        for x in x_pos:
            move_absolute(x, y)
            sleep(wait_time)

            shoot(n_pulses)

            sleep(0.2)
            while cta.cta_client.is_running():
                # wait until sequence is done
                print("waiting for CTA ...")
                sleep(0.1)

            logger.info(
                f"At {time():.5f}s, drilled {n_pulses} hole(s) at x:{x:.5f}, y:{y:.5f}, with transmission: {attenuator.get_transmission():.5f}."
            )

    move_absolute(x_start, y_start)


def drill_here():
    x_start = standa.x.get()
    x_stop = x_start + 2
    n_cols = 11

    y_start = standa.z.get()
    y_stop = y_start + 1
    n_rows = 4

    drill_holes_grid(x_start, x_stop, n_cols, y_start, y_stop, n_rows, wait_time=0.5)
