# For controlling experiments for the ion trap lab led by Prof. Yiheng Lin
# The code is written by Yintai Zhang, School of Physical Sciences, USTC
# Last updated: March 29th, 2019

import time
import DataType
from matplotlib import pyplot as plt
import device
from device import num_to_bytes, bytes_to_num, bytes_to_hexstr
import dds
import time
import numpy as np
from scipy import optimize
import scipy

##
#  This script is for containing the functions may use.
##

def GetNowTime():
    return time.strftime("%Y-%m-%d %H:%M:%S",time.localtime(time.time()))

def CopyFile(dir1, dir2):
    inputfile = open(dir1, 'r')
    outputfile = open(dir2, 'w+')

    text = inputfile.readlines()

    for line in text:
        print(line, outputfile)

def StringSeparate(s):
    s_list = []
    s_list.append('')

    for char in s:
        if char != ' ':
            s_list[len(s_list) - 1] = s_list[len(s_list) - 1] + char
        else:
            s_list.append('')
    
    return s_list

def RemoveSpace(s):
    return s.replace(" ","")

def scan_generate(lb, ub, step):

    scan_list = []

    if step == 0:
        return [lb]
    elif step > 0:
        i = lb
        while i <= ub:
            scan_list.append(i)
            i = i + step
    else:
        i = ub
        while i >= lb:
            scan_list.append(i)
            i = i + step

    # print(scan_list)
    
    return scan_list

def Sort_Data_List(channel_list, data_list):
    if len(channel_list) != len(data_list):
        return [channel_list, data_list]

    cp_list = list(zip(channel_list, data_list))
    cp_list.sort(key = lambda x : x[0])

    return [list(list(zip(*cp_list))[0]), list(list(zip(*cp_list))[1])]

# The functions for experiment
#---------------------------------------------------------------------------------------

# Doppler cooling
def IonTrap_Doppler(amp, freq, DDS, channel, initial = True, stop = False, duration = -1):

    # amp:
        # Data Type: float
        # The amplitude for driving Doppler cooling
    # freq:
        # Data Type: float
        # The frequency for driving Doppler cooling
    # DDS:
        # Data Type: device
        # The dds for operation
    # channel:
        # Data Type: integer
        # The channel for driving Doppler cooling
    # initial:
        # Data Type: bool
        # whether initialize the device before using
    # stop:
        # Data Type: bool
        # Stop Doppler cooling
    # Duration:
        # Data Type: float
        # The duration of the Doppler cooling

    if initial:
        DDS.initial_dds()
    else:
        return

    if stop:
        DDS.cw_play(channel, 0, freq, 0)
        return

    DDS.cw_play(channel, amp, freq, 0)

    if duration != -1:
        time.sleep(duration / 1000000)
        DDS.cw_play(channel, 0, 0, 0)

# Ramsey Sequence


def IonTrap_Ramsey(scan_list, scan_params, scan_type, DDS, channels, t4 = 50, t5 = 50,
                    amp_detect = 0.5, freq_detect = 300,
                    amp_pumping = 0.5, freq_pumping = 300, spin_echo = False, echo_phi = 0.5,
                    initial = True, scan_apply = True,
                    trigger_channel = -1,
                    test_mode = False, test_channel = -1):

    # scan_list:
        # Data Type: List
        # containing the points that needs to be scanned
        # The elements in the list are floats
    # scan_params:
        # Data Type: List
        # Format: [amp, freq, t1, t2, t3, phi, times]
        # t1, t2 are the time of the ramsey pulse; t3 is the time for spin-echo
        # freq, amp, t1, t2 and phi are floats. The parameter times is an integer.
    # scan_type:
        # Data Type: String
        # "freq", "t1", "t2", "t3", "phi", "amp"
    # DDS:
        # Data Type: device
        # The DDS for this operation
    # channels:
        # Data Type: List
        # Format: [channel_for_ramsey_waveform, channel_for_detection, channel_for_optical_pumping]
        # The elements in the list are integers
    # t4:
        # Data Type: float
        # The time for optical-pumping, default: 50 microseconds
    # t5:
        # Data Type: float
        # t5 is the time for detection, default: 50 microseconds
    # amp_detect:
        # Data Type: float
        # The amplitude for detecting, default: 0.5
    # amp_pumping:
        # Data Type: float
        # The amplitude for optical pumping, default: 0.5
    # spin_echo:
        # Data Type: bool
        # If spin_echo is applied, default: False
    # echo_phi:
        # Data Type: float
        # The phase for the echo pulse
    # scan_apply:
        # Data Type: bool
        # Whether the scan is applied or nor.
    # test_mode:
        # Data Type: bool
        # If use test mode. For test mode, signals are generated by a TLL. default: False
    # test_channel:
        # Data_Type: bool
        # The channel for generating test data.

    if not scan_apply:
        return [[]]

    fpga = DDS

    if len(scan_params) != 7 or len(channels) != 3:
        return [[]]

    if scan_list == []:
        return [[]]

    amp, freq, t1, t2, t3, phi, times = scan_params # Change the position of phi and (t1, t2, t3) result to bug
    t6 = 10 # standby time for every single-shot

    if spin_echo:
        echo = 1
    else:
        echo = 0

    data_list = [[], [], []]
    if test_mode:
        data_list.append([])

    play_list = []

    result = []

    if initial:
        fpga.initial_dds()
    else:
        pass

    if trigger_channel == -1:
        pass
    else:
        pass

    for channel in channels:
        fpga.cw_play(channel, 0, 0, 0)

    fpga.phase_clear_dds(channels)

    test_mode_count = 0

    if test_mode:
        if test_channel == -1:
            return [[]]
        else:
            channels.append(test_channel)

    for scan_var in scan_list:

        play_list.clear()
        data_list[0].clear()
        data_list[1].clear()
        data_list[2].clear()

        if scan_type == "t1":
            t1 = scan_var
        elif scan_type == "t2":
            t2 = scan_var
        elif scan_type == "t3":
            t3 = scan_var
        elif scan_type == "amp":
            amp = scan_var
        elif scan_type == "phi":
            phi = scan_var
        elif scan_type == "freq":
            freq = scan_var
        else:
            return [[]]

        ## The pulse for ramsey
        data_list[0].append([0, [0, 0, 0], ['low', t4]])
        data_list[0].append([0, [amp, freq, 0], ['high', t1]])
        data_list[0].append([0, [0, freq, 0], ['low', t2]])
        if spin_echo:
            data_list[0].append([0, [amp, freq, echo_phi], ['high',t3]])
            data_list[0].append([0, [0, freq, 0], ['low', t2]])
        data_list[0].append([0, [amp, freq, 0], ['high', t1]])
        data_list[0].append([0, [0, 0, 0], ['low', t5]])
        data_list[0].append([0, [0, 0, 0], ['low', t6]])

        # The channel for detection
        if len(data_list[1]) == 0:
            flag = 0
        else:
            flag = 1
        data_list[1].append([0, [0, 0, 0], ['low', t4 + t1 * 2 + t2 + echo * (t3 + t2)]])
        data_list[1].append([0, [amp_detect, freq_detect, 0], ['high', t5]])
        data_list[1].append([0, [0, 0, 0], ['low', t6]])

        # The channel for optical pumping
        data_list[2].append([0, [amp_pumping, freq_pumping, 0], ['high', t4]])
        data_list[2].append([0, [0, freq_pumping, 0], ['low', t1 * 2 + t2 + t5 + t6 + echo * (t3 + t2)]])

        if test_mode:
            data_list[3].clear()
            test_mode_count = test_mode_count + 1
            print("test mode counter =", test_mode_count)
            test_time_0 = (t4 + t1 * 2 + t2 + t5 + t6 + echo * (t3 + t2)) / test_mode_count
            for j in range (0, test_mode_count):
                data_list[3].append([0, [0, 0, 0], ['low', test_time_0 / 2]])
                data_list[3].append([0, [0, 0, 0], ['high', test_time_0]])
                data_list[3].append([0, [0, 0, 0], ['low', test_time_0 / 2]])

        # make_play_list

        play_list.append([times, 0, 0])
        channels, data_list = Sort_Data_List(channels, data_list)
        fpga.sequence_data_download(channels, data_list)
        result.append(fpga.play(0, play_list))

    print(data_list[0])

    return [channels, data_list, result]


# return the command for running  a python programme
def IonTrap_ToRunCommand(s):
    # s:
        # Data Type: string
        # The python script directory that needs to run by python
    return "python \"" + s + "\""

# n-th order polynomial fitting
def IonTrap_PolyFit(n, x_list, y_list, init = []):
    # n
        # Data Type: integer
        # The aimed order of the polynomial
    # x_list
        # Data Type: list
        # The values of x that needs to be fitted
    # y_list:
        # Data Type: list
        # The values of y that needs to be fitted
    # init:
        # Data Type: list
        # The initial values for the coefficients

    a = []

    for i in range (0, n + 1):
        a.append(0)

    if len(init) != n + 1:
        return []

    def IonTrap_PolyFit_func(coef, x):
        sum = coef[0]
        for i in range(1, n + 1):
            sum = sum + x ** i * coef[i]

        return sum

    def IonTrap_PolyFit_err_func(coef, x, y):
        return y - IonTrap_PolyFit_func(coef, x)


    return optimize.leastsq(IonTrap_PolyFit_err_func, np.array(init), args = (np.array(x_list), np.array(y_list)))

# Run the DAC
def IonTrap_DAC5371(fpga, ch_list, data_list, initial = False):

    # fpga:
        # Data Type: device
        # The FPGA that controlls the DAC5371
    # ch_list:
        # Data Type: list
        # Format: [ch1, ch2, ch3, ...]
        # The channels that needs to be played
    # data_list:
        # Data Type: list
        # Format[[ch1_1, ch1_2, ch1_3, ...], [ch2_1, ch2_2, ch2_3, ...], ...]
    # initial:
        # Data Type: bool
        # Initial the dds or not. Default: True.

    data_list = []

    if ch_list == []:
        return []

    if (len(ch_list) != len(data_list)):
        return []

    test_len = len(data_list[0])
    for i in range (1, len(ch_list)):
        if len(data_list[i]) != test_len:
            return []

    if initial:
        fpga.initial_dds()
    data_list_ = []
    for i in range (len(data_list[0])):
        data_list_.append([])
        for j in range (0, len(ch_list)):
            data_list_[i].append(data_list[j][i])

    fpga.ad5371_play(ch_list, data_list_)

# Set DDS to ZERO
def IonTrap_DAC5371_setZero(fpga, ch_list):
    # fpga:
        # Data Type: device
        # The FPGA that controlls the DAC5371
    # ch_list:
        # Data Type: list
        # Format: [ch1, ch2, ch3, ...]
        # The channels that needs to be played
    if len(ch_list) == 0:
        return []

    data_list = []
    for i in range (0, len(ch_list)):
        data_list.append([0])

    fpga.ad5371_play(ch_list, data_list)

def IonTrap_Update(name, data, file):

    # name:
        # The name of the variable that needs to be changed
    # data:
        # The data that needs to be replaced with
    # file:
        # The directory of the file

    pass

def IonTrap_Average(data_list):

    # data_list:
        # Data Type: list
        # Format : [[data_1_1, data_1_2, data_1_3,...], [data_2_1, data_2_2, data_2_3,...], ..., ...]
        # The data that needs to be average out

    aver_data_list = []
    cov_data_list = []
    for temp in data_list:
        print(temp)
        aver_data_list.append(np.average(temp))
        cov_data_list.append(np.cov(temp))

    return [aver_data_list, cov_data_list]


if __name__ == '__main__':

    def f(s): # Something like a functional

        b = s()
        print(b + ' ' + b)

    @f
    def g():
        return 'hello world!'

    print("start")


