#!/usr/bin/env python

from time import sleep
from datetime import datetime
import numpy as np
from tqdm import trange
from epics import PV

from slic.gui import GUI
from slic.core.adjustable import Adjustable, PVAdjustable, DummyAdjustable
from slic.core.acquisition import SFAcquisition#, PVAcquisition
from slic.core.condition import PVCondition
from slic.core.scanner import Scanner
from slic.devices.simpledevice import SimpleDevice
from slic.devices.general.motor import Motor
from slic.devices.general.delay_stage import DelayStage
from slic.devices.general.smaract import SmarActAxis
from slic.utils import devices, Marker, as_shortcut
#from slic.utils import Channels, Config, Elog, Screenshot, PV
#from slic.core.acquisition.fakeacquisition import FakeAcquisition

from devices.attenuator import Attenuator
from devices.undulator import Undulators, TwoColorChicane, Phases
from devices.magnet import Magnet

from channels import channels, pvs, detectors
from spreadsheet import overview, print_overview, print_line_for_spreadsheet
import colors


dummy = DummyAdjustable(units="au")

#jet_delay = PVAdjustable("SATES20-CVME-EVR0:Pul6-Delay-SP", "SATES20-CVME-EVR0:Pul6-Delay-RB", accuracy=1, name="Jet Delay")

source_delay = PVAdjustable("SATES20-CVME-EVR0:Pul14-Delay-SP", "SATES20-CVME-EVR0:Pul14-Delay-RB", accuracy=1, name="Source Delay")

waveplate = Motor("SLAAT21-LMOT-M702:MOT", name="Laser Waveplate")
TT_delay = Motor("SLAAT21-LMOT-M704:MOT", name="TT_delay")
#laser_comp21 = Motor("SLAAT21-LMOT-M701:MOT", name="Laser Comp T2")
#laser_comp11 = Motor("SLAAT21-LMOT-M705:MOT", name="Laser Comp T1 1")
pump_probe_delay = Motor("SLAAT21-LMOT-M703:MOT", name="pump_probe_delay")
#laser_comp12 = Motor("SLAAT21-LMOT-M706:MOT", name="Laser Comp T1 2")
#laser_delay = DelayStage("SLAAT21-LMOT-M708:MOT", name="Laser Delay")
laser_delay = Motor("SLAAT21-LMOT-M708:MOT", name="Laser Global Delay")
lxt = PVAdjustable("SLAAT01-LTIM-PDLY:DELAY", pvname_done_moving="SLAAT01-LTIM-PDLY:WAITING", name="LXT")

source_y = Motor("SATES20-MANIP1:MOTOR_1", name="Aerosol Injector x")
source_x = Motor("SATES20-MANIP1:MOTOR_2", name="Aerosol Injector y")

att = Attenuator("SATFE10-OATT064")
chic_delay = TwoColorChicane("zz Two Color Chicane")


und = Undulators(name="z Athos Undulators (both colors)")
und1 = Undulators([6, 7, 8, 9, 10, 11, 12, 13],     n_und_ref=10, name="z Athos Undulators  6-13 (first color)")
und2 = Undulators([15, 16, 17, 18, 19, 20, 21, 22], n_und_ref=19, name="z Athos Undulators 15-22 (second color)")

und.set_limits(390, 1100)
und1.set_limits(390, 1100)
und2.set_limits(390, 1100)


xota_z_trans = Motor("SATES21-XOTA166:W_Z",     name="Table 1: Z coordinated")
xota_y_trans = Motor("SATES21-XOTA166:W_Y",     name="Table 1: Y coordinated")
#xota_hori_z  = Motor("SATES21-XOTA166:MOTOR_Z", name="Table 1: Z raw")

mot1_x = SmarActAxis("SATES21-XSMA166:MOT4", name="Table 1: sma X")
mot1_y = SmarActAxis("SATES21-XSMA166:MOT5", name="Table 1: sma Y")
mot1_z = SmarActAxis("SATES21-XSMA166:MOT6", name="Table 1: sma Z")

mot3_z = SmarActAxis("SATES23-XSMA169:MOT3", name="Table 3: sma Z")


standa = Motor("SLAAT21-LMOT-M707:MOT", name="Standa Motor")

mark_microscope = Marker(standa, 141.344, "Microscope")
mark_gascell    = Marker(standa,  88.334, "Gas Cell")
mark_tof        = Marker(standa,  28.330, "TOF")


magn = PV("SATUN14-MBND100:CYCLE")

@as_shortcut
def cycle_magnet():
    magn.put(1, wait=True)
    for _ in trange(250):
        sleep(1)
    # set the current to 1 and 5 consecutively,
    # since the results after the first step always looks strange
    chic_delay.set(1).wait()
    sleep(1)
    chic_delay.set(5).wait()
    sleep(1)


# add some more devices to the overview
overview.standa = standa
overview.TT_delay = TT_delay
overview.pump_probe_delay = pump_probe_delay
overview.LXT = lxt





instrument = "maloja"
#pgroup = "p18493" # commissioning june 2020
#pgroup = "p18722" # commissioning sept 2020
#pgroup = "p18863" # JF commissioning april 2021
#pgroup = "p19241" # two colors
#pgroup = "p19342" # XMCD
#pgroup = "p19343" # slic dev
#pgroup = "p19373" # laser
#pgroup = "p19483" # pilot 1 -- TRXPS
#pgroup = "p19488" # pilot 2 -- Non-Linear Spectroscopy
#pgroup = "p19510" # pilot 2 -- Non-Linear Spectroscopy
#pgroup = "p19727" # Laser only timing tests
#pgroup = "p19509" # slic dev
#pgroup = "p19743" # Focus tests and timing tests
#pgroup = "p19582" # pilot 3 -- SPI
pgroup = "p19583" # permalloy 2.0
#pgroup = "p20073" # alvise

daq = SFAcquisition(instrument, pgroup, default_channels=channels, default_pvs=pvs, default_detectors=detectors, rate_multiplicator=1, append_user_tag_to_data_dir=True)
#daq = FakeAcquisition(instrument, pgroup)

#daqPV = PVAcquisition(instrument, pgroup, default_channels=channels_ks) # workaround for KS not going to DB

check_intensity = PVCondition("SATBD01-DBPM060:Q2", vmin=5, vmax=None, wait_time=1, required_fraction=0.8)
#check_intensity = None

scan = Scanner(scan_info_dir=f"/sf/{instrument}/data/{pgroup}/res/scan_info", default_acquisitions=[daq], condition=check_intensity)

gui = GUI(scan, show_goto=True, show_spec=True)