from keras.preprocessing.image import ImageDataGenerator, array_to_img, img_to_array, load_img
from matplotlib.pyplot import imread, imshow, show, savefig
import numpy as np
import sys
import os
import cv2

WALDO_DIR = os.path.dirname(os.path.realpath(__file__)) + '/../Hey-Waldo/'
WALDO_DUPLICATION_NB = 100
NOTWALDO_DUPLICATION_NB = 2
NOISE_DUPLICATION = 30

def get_images_data(dir):
    return np.array([np.array(imread(WALDO_DIR + dir + '/waldo/' + fname))
                     for fname in os.listdir(WALDO_DIR + dir + '/waldo')]),\
            np.array([np.array(imread(WALDO_DIR + dir + '/notwaldo/' + fname))
                      for fname in os.listdir(WALDO_DIR + dir + '/notwaldo')])

def data_augmentation(waldos, notWaldos, dir):
    datagen = ImageDataGenerator(
        rotation_range=40,
        width_shift_range=0.2,
        height_shift_range=0.2,
        rescale=1. / 255,
        horizontal_flip=True,
        fill_mode='nearest')

    print('Waldo data augmentation...')
    for waldo in waldos:
        waldo = waldo.reshape((1,) + waldo.shape)

        i = 0
        for batch in datagen.flow(waldo, batch_size=1,
                                  save_to_dir=WALDO_DIR + dir + '/waldoOther', save_prefix='waldo', save_format='jpeg'):
            i += 1
            if i >= WALDO_DUPLICATION_NB - 1:
                break

    print('Not waldo data augmentation...')
    for not_waldo in notWaldos:
        not_waldo = not_waldo.reshape((1,) + not_waldo.shape)

        i = 0
        for batch in datagen.flow(not_waldo, batch_size=1,
                                  save_to_dir=WALDO_DIR + dir + '/notWaldoOther', save_prefix='waldo', save_format='jpeg'):
            i += 1
            if i >= NOTWALDO_DUPLICATION_NB - 1:
                break

def get_waldos_data(path):
    waldos = []
    for fname in os.listdir(path + '/'):
        waldos.append(cv2.imread(path + '/' + fname))
    return np.array(waldos)

def data_add_noise(waldos, path):
    for idx, waldo in enumerate(waldos):
        for i in range(40):
            noise = (0.4 * np.random.randn(*waldo.shape)) * 10
            img = waldo + noise
            cv2.imwrite(path + '/noisy_' + str(idx) + '_' + str(i) + '.jpg', img)

if __name__ == '__main__':
    if len(sys.argv) > 2:
        if sys.argv[1] == 'aug':
            waldos, notWaldos = get_images_data(sys.argv[2])
            data_augmentation(waldos, notWaldos, sys.argv[2])
        elif sys.argv[1] == 'noise':
            waldos = get_waldos_data(sys.argv[2])
            data_add_noise(waldos, sys.argv[2])
    else:
        print('Usage: data_augmentation.py [aug or noise] [image directory]')


from collections import deque
import sys

input = sys.stdin.readline

f_inf = float("inf")

n, m = map(int, input().split())
to = [[] for _ in range(n)]

for _ in range(m):
    a, b = map(int, input().split())
    a -= 1
    b -= 1
    to[a].append(b)
    to[b].append(a)

que = deque([])
que.append(0)

dist = [f_inf for _ in range(n)]
dist[0] = 0

guide = [-1 for _ in range(n)]

while que:
    v = que.popleft()
    for u in to[v]:
        if dist[u] != f_inf: continue
        dist[u] = dist[v] + 1
        guide[u] = v
        que.append(u)

print("Yes")
for i in range(1, n):
    ans = guide[i] + 1
    print(ans)

Atcoder/20200720_BBC_40/D - .. (Double Dots).py


#!/usr/bin/python3
# -*- coding: utf-8 -*-
"""
Created on Wed Aug 24 11:09:19 2022

@author: user
"""

# importing required libraries
import sys
import cv2
import time
import random
import configparser
import matplotlib.pyplot as plt
from threading import Thread # library for implementing multi-threaded processing

import TIS

# deskewing
import numpy as np
import math
# from deskew import determine_skew # for detect angle
from typing import Tuple, Union

# OCR
import pytesseract
from libc import azure_ocr
# import pytessy

# modbus
import modbus_tk.defines as cst
import modbus_tk.modbus_tcp as modbus_tcp

# debug
image_debug = False
# image_file = "./MKS.png"
image_file = "./img/2022_09_01_Catupre_8.png"

# parameters
config = configparser.ConfigParser()
config.read('/home/smg/smg/ocr/v5/config.ini')

windows_base    = config['general'].getboolean('windows_base')
debug           = config['general'].getboolean('debug')

# ocr
azure_ocr_en    = config['ocr'].getboolean('azure_ocr_en')
win_psm         = config['ocr'].getint('win_psm')

# camera
cameraID = config['camera'].getint('cameraID')
wind_show = config['camera'].getboolean('windows_show')

# image pre-processing
angle           = config['img_preprocessing'].getfloat('angle')
region_topx     = config['img_preprocessing'].getint('region_topx')
region_topy     = config['img_preprocessing'].getint('region_topy')
region_bottomx  = config['img_preprocessing'].getint('region_bottomx')
region_bottomy  = config['img_preprocessing'].getint('region_bottomy')

# modbus
mod_client  = config['modbus']['mod_client']
mod_port    = config['modbus'].getint('mod_port')
slaveID     = config['modbus'].getint('slaveID')
size        = config['modbus'].getint('size')
scale       = config['modbus'].getint('scale')

list_KEY = ['GNITE','GHITE','ON','hour']

err_char = False
result = 0
key_type = 0
lastVal_ON = 0
lastVal_hour = 0
lastVal_IGNITE = 0
list_last = []

# list_READY = ['READY', 'EADY']
# list_IGNITE = ['IGNITE', 'GNITE']
# list_PLASMA = ['PLASMA', 'PLASHA']

# if(windows_base):
#     print("windows_base: {}".format(windows_base))
#     pytesseract.pytesseract.tesseract_cmd = r'C:\Program Files\Tesseract-OCR\tesseract.exe'

def rotate(
        image: np.ndarray, angle: float, background: Union[int, Tuple[int, int, int]]
) -> np.ndarray:
    old_width, old_height = image.shape[:2]
    angle_radian = math.radians(angle)
    width = abs(np.sin(angle_radian) * old_height) + abs(np.cos(angle_radian) * old_width)
    height = abs(np.sin(angle_radian) * old_width) + abs(np.cos(angle_radian) * old_height)

    image_center = tuple(np.array(image.shape[1::-1]) / 2)
    rot_mat = cv2.getRotationMatrix2D(image_center, angle, 1.0)
    rot_mat[1, 2] += (width - old_width) / 2
    rot_mat[0, 2] += (height - old_height) / 2
    return cv2.warpAffine(image, rot_mat, (int(round(height)), int(round(width))), borderValue=background)

def merge_rgb(R,G,B):
    img = cv2.merge([B,G,R])
    return img

def img_preprocessing(image, angle):
    global region_topx
    global region_topy
    global region_bottomx
    global region_bottomy

    #print("region_topx:{}".format(region_topx))
    grayscale = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)

    if(image_debug):
        # print("test")

        skewed = rotate(grayscale, 3.5, (0, 0, 0))

        # region
        region_topx = 200
        region_topy = 0
        region_bottomx = 320
        region_bottomy = 650
        Cropped = skewed[region_topx:region_bottomx+1, region_topy:region_bottomy+1]

        # print("grayscale shape:{}".format(grayscale.shape))

        # blur = Cropped
        blur = cv2.fastNlMeansDenoising(Cropped, h=7,templateWindowSize=7,searchWindowSize=21);

        result = blur

    else:

        # Cropped
        region_topx = 170
        region_topy = 70
        region_bottomx = 300
        region_bottomy = 550
        Cropped = frame[region_topx:region_bottomx, region_topy:region_bottomy]
        # show_img(Cropped)

        # skew
        skewed = rotate(Cropped, 3.3, (0, 0, 0))
        # show_img(skewed)

        # Cropped
        region_topx = 45
        region_topy = 30
        region_bottomx = 115
        region_bottomy = 445
        Cropped = skewed[region_topx:region_bottomx, region_topy:region_bottomy]
        # show_img(Cropped)

        # G channel
        (B, G, R) = cv2.split(Cropped)
        zeros = np.ones(Cropped.shape[:2],dtype=np.uint8)
        G_ch_img = merge_rgb(R=zeros, G=G, B=zeros)
        # show_img(G_ch_img)

        # gray
        grayscale = cv2.cvtColor(G_ch_img, cv2.COLOR_BGR2GRAY)

        # blur-median
        blur_median = cv2.medianBlur(grayscale,3)
        # show_img(blur_median)

        # enhance black
        enhanced = np.where((blur_median>50) & (blur_median<135), blur_median-40, blur_median)
        # show_img(enhanced)

        result = enhanced

    return result

def isfloat(num):
    try:
        float(num)
        return True
    except ValueError:
        return False

def modbus_preprocessing(ocr_text):
    global scale
    global lastVal_ON
    global lastVal_hour
    global lastVal_IGNITE
    global list_last
    global key_type
    global err_char

    global result
    scale = 1
    st_idx = 0

    # case sensitive
    for key in list_KEY: #choosen key
        if(key in ocr_text): #check key
            if(key=="ON"):
                st_idx = ocr_text.find(" ")+1
                scale = 100
                key_type = 1
                if(len(ocr_text.strip())>st_idx): # "PLASMA ON 3. 7kw"
                    tmp_val = ""
                    for i in range(len(ocr_text)-st_idx):
                        tmp_char = ocr_text[st_idx+i]
                        if(tmp_char==" " and i>3):
                            break
                        elif(tmp_char.isnumeric() or tmp_char=="."): #fetch if char is valid
                            tmp_val = tmp_val+tmp_char
                    if(tmp_val[-1:]=="."): #remove last char if non-numeric
                        tmp_val = tmp_val[:-1]
                    if(isfloat(tmp_val)): #chech float or not
                        list_last.append(float(tmp_val)*scale)
                        print(list_last)

            elif(key=="hour"):
                st_idx = ocr_text.find(" ")+1
                scale = 1
                key_type = 2
                if(len(ocr_text.strip())>st_idx): # "READY 14996. hour"
                    tmp_val = ""
                    for i in range(len(ocr_text)-st_idx):
                        tmp_char = ocr_text[st_idx+i]
                        if(tmp_char==" " and i>3):
                            break
                        elif(tmp_char.isnumeric() or tmp_char=="."): #fetch if char is valid
                            tmp_val = tmp_val+tmp_char

                    if(tmp_val[-1:]=="."): #remove last char if non-numeric
                        tmp_val = tmp_val[:-1]
                    if(isfloat(tmp_val)): #chech float or not
                        list_last.append(float(tmp_val)*scale)
                        print(list_last)

            elif(key=="GNITE" or key=='GHITE'):
                st_idx = ocr_text.find(" ")+1
                scale = 1
                key_type = 3
                if(len(ocr_text.strip())>st_idx): # "READY 14996. hour"
                    tmp_val = ""
                    for i in range(len(ocr_text)-st_idx):
                        tmp_char = ocr_text[st_idx+i]
                        if(tmp_char==" " and i>3):
                            break
                        elif(tmp_char.isnumeric() or tmp_char=="."): #fetch if char is valid
                            tmp_val = tmp_val+tmp_char
                        else:
                            if(err_char):
                                break
                            err_char = True

                    if(tmp_val[-1:]=="."): #remove last char if non-numeric
                        tmp_val = tmp_val[:-1]
                    if(isfloat(tmp_val)): #chech float or not
                        list_last.append(float(tmp_val)*scale)
                        print(list_last)

        elif(key==list_KEY[-1]): #end
            #print("result")
            lastVal = 0
            if(key_type==1):
                lastVal = lastVal_ON

            elif(key_type==2):
                lastVal = lastVal_hour

            elif(key_type==3):
                lastVal = lastVal_IGNITE

            if((lastVal*scale) in list_last): # choose lastVal as output
                result = lastVal
            else: # update lastVal and choose mode value
                for i in list_last:
                    counts = np.bincount(list_last)
                    result = (np.argmax(counts)/scale)
                    lastVal = (result/scale)

            # reset
            #key_type = 0
            list_last = []

    return result

# define ModbusTCP with multi-threaded processing
class ModbusTCP:
    def __init__(self, client='127.0.0.1', port=502, slaveID=1, size=8):
        self.client = client
        self.slaveID = slaveID
        self.port = port
        self.size = size
        self.val = 0

        # create modbus server
        self.SERVER = modbus_tcp.TcpServer(address = '127.0.0.1', port = int(self.port))

        # server start
        self.SERVER.start()

        # print("slaveID: {}, size: {}".format(self.slaveID, self.size))
        # build slave
        self.slave = self.SERVER.add_slave(int(self.slaveID))
        self.slave.add_block('A', cst.HOLDING_REGISTERS, 0, self.size) # adres:0, len:8

        self.stopped = True

        # reference to the thread for reading next available frame from input stream
        self.t = Thread(target=self.holdingRegister, args=())
        self.t.daemon = True # daemon threads keep running in the background while the program is executing

    # start
    def start(self):
        self.stopped = False
        self.t.start()

    # update
    def holdingRegister(self):
        while True:
            if(self.stopped):
                break
            else:
                self.slave.set_values('A', 0, int(self.val))
            time.sleep(0.5)

        self.SERVER.stop()

    def update(self, value):
        self.val = value

    def stop(self):
        self.stopped = True

# initializing and starting multi-threaded webcam capture input stream
# webcam_stream = WebcamStream(stream_id=0) #  stream_id = 0 is for primary camera
# webcam_stream.start()

class CustomData:
    ''' Example class for user data passed to the on new image callback function
    '''
    def __init__(self, newImageReceived, image):
        self.newImageReceived = newImageReceived
        self.image = image
        self.busy = False

CD = CustomData(False, None)

def on_new_image(tis, userdata):
    '''
    Callback function, which will be called by the TIS class
    :param tis: the camera TIS class, that calls this callback
    :param userdata: This is a class with user data, filled by this call.
    :return:
    '''
    # Avoid being called, while the callback is busy
    if userdata.busy is True:
        return

    userdata.busy = True
    userdata.newImageReceived = True
    userdata.image = tis.Get_image()
    userdata.busy = False

Tis = TIS.TIS()

# The following line opens and configures the video capture device.
Tis.openDevice("16024087", 640, 480, "30/1",TIS.SinkFormats.BGRA, False)

# The next line is for selecting a device, video format and frame rate.
# if not Tis.selectDevice():
#     quit(0)

#Tis.List_Properties()
Tis.Set_Image_Callback(on_new_image, CD)

Tis.Set_Property("TriggerMode", "On")

Tis.Start_pipeline()

# Remove comment below in oder to get a propety list.
# Tis.List_Properties()

# In case a color camera is used, the white balance automatic must be
# disabled, because this does not work good in trigger mode
try:
    Tis.Set_Property("BalanceWhiteAuto", "Continuous")
    # Tis.Set_Property("BalanceWhiteAuto", "Off")
    # Tis.Set_Property("BalanceWhiteRed", 1.2)
    # Tis.Set_Property("BalanceWhiteGreen", 1.0)
    # Tis.Set_Property("BalanceWhiteBlue", 1.4)

except Exception as error:
    print(error)

try:
    # Query the gain auto and current value :
    print("GainAuto : %s " % Tis.Get_Property("GainAuto"))
    print("Gain : %d" % Tis.Get_Property("Gain"))

    # Check, whether gain auto is enabled. If so, disable it.
    if Tis.Get_Property("GainAuto"):
        Tis.Set_Property("GainAuto", "Off")
        print("Gain Auto now : %s " % Tis.Get_Property("GainAuto"))

    Tis.Set_Property("Gain", 0)

    # Now do the same with exposure. Disable automatic if it was enabled
    # then set an exposure time.
    if Tis.Get_Property("ExposureAuto") :
        Tis.Set_Property("ExposureAuto", "Off")
        print("Exposure Auto now : %s " % Tis.Get_Property("ExposureAuto"))

    Tis.Set_Property("ExposureTime", 24000)

except Exception as error:
    print(error)
    quit()

error = 0
print('Press Esc to stop')
lastkey = 0
# cv2.namedWindow('Window',cv2.WINDOW_NORMAL)

modbus = ModbusTCP(client=mod_client, port=mod_port, slaveID=slaveID, size=size)
modbus.start()

if(debug):
    # processing frames in input stream
    num_frames_processed = 0
    start = time.time()

try:
    while lastkey != 27 and error < 5:
        # time.sleep(1)
        Tis.execute_command("TriggerSoftware") # Send a software trigger

        # Wait for a new image. Use 10 tries.
        tries = 10
        while CD.newImageReceived is False and tries > 0:
            time.sleep(0.1)
            tries -= 1

        # Check, whether there is a new image and handle it.
        if CD.newImageReceived is True:
            CD.newImageReceived = False
            frame = CD.image

    # while True :
    #     if webcam_stream.stopped is True :
    #         break
    #     else :
    #         frame = webcam_stream.read()

            if(image_debug):
                frame = cv2.imread(image_file)
                # gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
                # gray = cv2.bilateralFilter(gray, 13, 15, 15)

            # image pre-processing
            proced_img = img_preprocessing(frame, angle)

            # plt.imshow(proced_img)

            # OCR
            if(azure_ocr_en):
                ocr_text = azure_ocr.azure_ocr(proced_img)
            else: # local ocr
                ocr_type = "--psm {}".format(win_psm)
                ocr_text = pytesseract.image_to_string(proced_img, config=ocr_type)

            print(ocr_text.strip())

            # modbus pre-processing
            tmpVal = modbus_preprocessing(ocr_text)
            print("ocr_value: {}".format(tmpVal))

            # modbus
            if(debug):
                ocr_val = random.uniform(3.5,4.5)*100
            # print(ocr_text)

            #modbus.update(ocr_val)

            if(debug):
            # adding a delay for simulating time taken for processing a frame
            # delay = 0.03 # delay value in seconds. so, delay=1 is equivalent to 1 second
            # time.sleep(delay)
                num_frames_processed += 1

                if(wind_show):
                    cv2.imshow('frame' , proced_img)
                    key = cv2.waitKey(1)
                    if key == ord('q'):
                        break

except KeyboardInterrupt:
    print("User Terminated..!")
    # webcam_stream.stop()
    # modbus.stop()

finally:
    if(debug):
        end = time.time()

    # terminate process
    print("Stop the threading..!")
    # webcam_stream.stop()
    Tis.Stop_pipeline()
    modbus.stop()

    if(debug):
        # printing time elapsed and fps
        elapsed = end-start
        fps = num_frames_processed/elapsed
        print("FPS: {} , Elapsed Time: {} , Frames Processed: {}".format(fps, elapsed, num_frames_processed))

        if(wind_show):
            # closing all windows
            cv2.destroyAllWindows()
        #cv2.destroyAllWindows()


"""
The most important thing an ORM gives us is persistence ignorance: the idea that our fancy domain model doesn’t
need to know anything about how data is loaded or persisted. This helps keep the domain clean of direct
dependencies on particular database technologies.

Here, the usual way, defined in the SQLAlchemy tutorial, for associating model classes to a declarative_base() function
has not been used. Because the target architecture is clearly to apply the DIP : avoid model properties to depend on
storage concerns.
https://docs.sqlalchemy.org/en/14/orm/tutorial.html#declare-a-mapping

Instead, a mapper is used to make the model aware of the infrastructure layer only when start_mappers() is invoked.
https://docs.sqlalchemy.org/en/14/orm/mapping_api.html?#sqlalchemy.orm.mapper
"""

from sqlalchemy import MetaData, Table, Column, Integer, String, ForeignKey
from sqlalchemy.orm import mapper, relationship

from src.staffoptimizer.domain import model

metadata = MetaData()

staffoptimizers = Table(
    "staffoptimizers",
    metadata,
    Column("run_id", String(255)),
)
tasks = Table(
    "tasks",
    metadata,
    Column("id", Integer, primary_key=True, autoincrement=True),
    Column("run_id", ForeignKey("staffoptimizer_runs.id")),
    Column("reference", String(255)),
    Column("status", String(255)),
    Column("team", String(255)),
)
users = Table(
    "users",
    metadata,
    Column("id", Integer, primary_key=True, autoincrement=True),
    Column("name", String(255)),
    Column("editor_id", String(255)),
    Column("csid", String(255)),
    Column("esid", String(255)),
)
assignments = Table(
    "assignments",
    metadata,
    Column("id", Integer, primary_key=True, autoincrement=True),
    Column("task_id", ForeignKey("tasks.id")),
    Column("user_id", ForeignKey("users.id")),
)

def start_mappers():
    users_mapper = mapper(model.User, users)
    tasks_mapper = mapper(
        model.Task,
        tasks,
        properties={
            "_assignments": relationship(
                users_mapper,
                secondary=assignments,
                collection_class=set,
            )
        },
    )
    mapper(
        model.StaffOptimizer,
        staffoptimizers,
        properties={"tasks": relationship(tasks_mapper)},
    )


import numpy as np
import math
import cmath
from scipy.optimize import minimize
import scipy.fftpack
from sklearn.linear_model import LinearRegression
from scipy.optimize import minimize_scalar, curve_fit, leastsq
import os

# Running tensorflow in compat mode
import tensorflow
if int(tensorflow.__version__[0]) > 1:
    os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
    import tensorflow.compat.v1 as tf
    tf.disable_v2_behavior()
    tf.logging.set_verbosity(tf.logging.ERROR)
else:
    os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
    import tensorflow as tf

# Set number of threads of tensorflow
tf_num_threads=1
tf.config.threading.set_inter_op_parallelism_threads(tf_num_threads)

# other optimization imports
from . import fft_estimators, linear_estimators

# plotting
import matplotlib.pyplot as plt

class InitialEstimator(object):
    """
    Sets up the Initial Estimator with different solver options and loss functions

    InitialEstimator uses the rabi oscillations for estimation

    TODO: Add options for different parameterizations and loss functions
    """

    def __init__(self, syn_data, env_cr=None, A_cr=None,
                 FLAG_mle_solver=False, solver_options=None, do_plot=False,
                 FLAG_fft_window=False, type_fft_window=None, FLAG_fft_high_resoultion=False):

        # Regarding the use of FFT
        self.FLAG_fft_window = FLAG_fft_window
        self.type_fft_window = type_fft_window
        self.FLAG_fft_high_resolution = FLAG_fft_high_resoultion

        # Regarding the use of data
        samples = np.array(syn_data['samples']).astype(np.float32)
        mvec = np.array(syn_data['mvec']).astype(np.int16)
        uvec = np.array(syn_data['uvec']).astype(np.int16)
        tvec = np.array(syn_data['tvec']).astype(np.float32)

        # This is really dirty but I think this is the easiest way to get this working for the moment
        self.samples_config = {0: [], 1: [], 2: [], 3: [], 4: [], 5: []}
        self.mvec_config = {0: [], 1: [], 2: [], 3: [], 4: [], 5: []}
        self.uvec_config = {0: [], 1: [], 2: [], 3: [], 4: [], 5: []}
        self.tvec_config = {0: [], 1: [], 2: [], 3: [], 4: [], 5: []}

        for ind_sample in range(len(mvec)):
            ind_config = int(2*mvec[ind_sample] + uvec[ind_sample])

            self.samples_config[ind_config].extend([samples[ind_sample]])
            self.mvec_config[ind_config].extend([mvec[ind_sample]])
            self.uvec_config[ind_config].extend([uvec[ind_sample]])
            self.tvec_config[ind_config].extend([tvec[ind_sample]])

        for ind_config in range(6):
            self.samples_config[ind_config] = np.array(self.samples_config[ind_config]).astype(np.float32)
            self.mvec_config[ind_config] = np.array(self.mvec_config[ind_config]).astype(np.int16)
            self.uvec_config[ind_config] = np.array(self.uvec_config[ind_config]).astype(np.int16)
            self.tvec_config[ind_config] = np.array(self.tvec_config[ind_config]).astype(np.float32)

        self.FLAG_classification = syn_data['FLAG_classification']

        # Extract relevant parameters
        time_stamps = A_cr.tset * A_cr.xi_t
        freq_convert = A_cr.freq_convert

        # Calculate the rabi oscillations from the given data
        pvec_data = A_cr.rabi_data(env_cr)
        nshots_rabi_data, n0shots_rabi_data = A_cr.nshots_rabi_data(env_cr)

        # Get the sample frequency
        if 'sample_freq' in syn_data.keys():
            self.sample_freq = syn_data['sample_freq']
        else:
            first_time_stamp = np.amin(time_stamps)
            last_time_stamp = np.amax(time_stamps)

            delta_time = (last_time_stamp - first_time_stamp) / (len(time_stamps) - 1)
            self.sample_freq = 1.0 / delta_time

        self.time_stamps = time_stamps
        self.pvec_data = pvec_data
        self.nshots_rabi_data = nshots_rabi_data
        self.n0shots_rabi_data = n0shots_rabi_data
        self.freq_convert = freq_convert
        self.do_plot = do_plot

        if FLAG_mle_solver:
            # Setup solver options
            default_solver_options = {'nepochs': 200, 'neval_period': 10,
                                      'learning_rate': 0.001, 'optimizer': 'adam',
                                      'mini_batch_size': 100}

            if solver_options is None:
                self.mle_nepochs = default_solver_options['nepochs']
                self.mle_neval_period = default_solver_options['neval_period']
                self.mle_learning_rate = default_solver_options['learning_rate']
                self.mle_optimizer = default_solver_options['optimizer']
                self.mle_mini_batch_size = default_solver_options['mini_batch_size']
            else:
                if 'nepochs' in solver_options.keys():
                    self.mle_nepochs = solver_options['nepochs']
                else:
                    self.mle_nepochs = default_solver_options['nepochs']

                if 'neval_period' in solver_options.keys():
                    self.mle_neval_period = solver_options['neval_period']
                else:
                    self.mle_neval_period = default_solver_options['neval_period']

                if 'learning_rate' in solver_options.keys():
                    self.mle_learning_rate = solver_options['learning_rate']
                else:
                    self.mle_learning_rate = default_solver_options['learning_rate']

                if 'optimizer' in solver_options.keys():
                    self.mle_optimizer = solver_options['optimizer']
                else:
                    self.mle_optimizer = default_solver_options['optimizer']

                if 'mini_batch_size' in solver_options.keys():
                    self.mle_mini_batch_size = solver_options['mini_batch_size']
                else:
                    self.mle_mini_batch_size = default_solver_options['mini_batch_size']

    def rabi_amplitude_mle_estimator(self, init_param, rabi_freq, ind_config, verbose=True, do_plot=False):
        '''
        MLE Estimator for rabi amplitudes A, B, C given value of rabi frequency
        '''
        if verbose:
            print("Initial Estimation: Model training for %s epochs, with evaluation every %s steps" % (
            self.mle_nepochs, self.mle_neval_period))

        # Extract solver options
        nepochs = self.mle_nepochs
        neval_period = self.mle_neval_period
        learning_rate = self.mle_learning_rate
        mini_batch_size = self.mle_mini_batch_size

        xy = np.stack([self.tvec_config[ind_config], self.samples_config[ind_config]], axis=1)
        batch_size = len(self.tvec_config[ind_config])  # MLE Batch Size

        n_mini_batches = int(xy.shape[0] / mini_batch_size)

        # graph input
        T = tf.placeholder(tf.float32, name="X_time_points")
        Y = tf.placeholder(tf.float32, name="Y_samples")  # y = 0 or 1
        XI = tf.placeholder(tf.float32, name="XI_param")  # y = 0 or 1
        RABI_FREQ = tf.placeholder(tf.float32, name="rabi_freq")  # y = 0 or 1

        # model variables
        theta_nd = tf.Variable(init_param, name="theta_nd", dtype=tf.float32)

        if verbose:
            print("Input data xy shape=%s" % str(xy.shape))
            print("Each epoch has %d mini_batches of size %s" % (n_mini_batches, mini_batch_size))
            print("Input data shapes samples=%s, tpts=%s" % (self.samples_config[ind_config].shape, self.tvec_config[ind_config].shape))
            print("Input data placeholders X=%s, Y=%s, XI=%s" % (T, Y, XI))
            print("initial values: theta=%s" % (init_param,))

        # model to predict sample probability (prob(k=0))
        Y_sp = tf_model_free_probabilities(theta_nd, RABI_FREQ, T)

        # loss function (binary cross-entropy)
        loss = -1 * tf.log((1 - Y) * Y_sp + Y * (1 - Y_sp) + 1.0e-10)  # log likelihood
        loss = tf.reduce_mean(loss)  # take mean over batch

        # optimizer
        if self.mle_optimizer == "gd":
            gdo = tf.train.GradientDescentOptimizer(learning_rate=learning_rate)
            optimizer_op = gdo.minimize(loss)
            op_name = "GD"
        elif self.mle_optimizer == "adagrad":
            optimizer_op = tf.train.AdagradOptimizer(learning_rate).minimize(loss)
            op_name = "Adagrad"
        else:
            optimizer_op = tf.train.AdamOptimizer(learning_rate=learning_rate).minimize(loss)
            op_name = "Adam"

        losses = []  # record losses at each epoch step
        steps = []

        # run MLE
        if verbose:
            print("Using %s optimmizer, learning rate=%s" % (op_name, learning_rate))
            print("Running MLE over %s datapoints with %s epochs" % (batch_size, nepochs))

        config = tf.ConfigProto(
            device_count={'GPU': 0}
        )

        with tf.Session(config=config) as sess:
            sess.run(tf.global_variables_initializer())
            for k in range(nepochs):
                np.random.shuffle(xy)  # random in-place permutation of first dimension
                for n in range(n_mini_batches):
                    n0 = n * mini_batch_size
                    sess.run(optimizer_op, feed_dict={T: xy[n0:n0 + mini_batch_size, 0],
                                                      Y: xy[n0:n0 + mini_batch_size, 1],
                                                      RABI_FREQ: rabi_freq})

                if not (k % neval_period):
                    results = sess.run([loss, theta_nd],
                                       feed_dict={T: self.tvec_config[ind_config],
                                                  Y: self.samples_config[ind_config],
                                                  RABI_FREQ: rabi_freq})
                    if verbose:
                        print("    Epoch %s: loss=%s, theta=%s" % tuple([k] + results))
                    losses.append(results[0])
                    steps.append(k)
                    if np.isnan(results[0]):
                        raise Exception("loss is NaN, quitting!")

            results = sess.run([loss, theta_nd],
                               feed_dict={T: self.tvec_config[ind_config],
                                          Y: self.samples_config[ind_config],
                                          RABI_FREQ: rabi_freq})

        m_loss, m_theta = results

        if do_plot:
            plt.figure(2, figsize=(10, 6))
            plt.plot(steps, losses, 'go')
            plt.plot(steps, losses)
            plt.grid(True)
            plt.xlabel("Epoch step number")
            plt.ylabel("Loss (negative log likelihood)")
            plt.suptitle("Tensorflow MLE on dataset with %s samples using %s optimizer" %
                         (samples.size, op_name))

        ds = {'loss': losses, 'steps': steps, 'results': results}

        return ds

    def normal_equations(self, omega, f):
        R00 = np.sum([ np.cos(omega*t)**2 for t in self.time_stamps ])
        R01 = np.sum([ np.cos(omega*t)*np.sin(omega*t) for t in self.time_stamps ])
        R02 = np.sum([np.cos(omega * t) for t in self.time_stamps])
        R11 = np.sum([np.sin(omega * t) ** 2 for t in self.time_stamps])
        R12 = np.sum([np.sin(omega * t) for t in self.time_stamps])
        R22 = np.sum([1 for t in self.time_stamps])

        R = np.array([[R00, R01, R02], [R01, R11, R12], [R02, R12, R22]])

        x = np.linalg.solve(R, f)

        # Ref: https://numpy.org/doc/stable/reference/generated/numpy.vdot.html
        return np.vdot(x,f)

    def normal_equations_uncertainty(self, omega, f, var_x):
        cos_omega_t = np.array([np.cos(omega*t) for t in self.time_stamps])
        sin_omega_t = np.array([np.sin(omega*t) for t in self.time_stamps])
        temp_t = np.array(1/var_x)

        R00 = np.dot(cos_omega_t**2,temp_t)
        R01 = np.dot(cos_omega_t * sin_omega_t, temp_t)
        R02 = np.dot(cos_omega_t, temp_t)
        R11 = np.dot(sin_omega_t**2, temp_t)
        R12 = np.dot(sin_omega_t, temp_t)
        R22 = np.sum(temp_t)

        R = np.array([[R00, R01, R02], [R01, R11, R12], [R02, R12, R22]])

        x = np.linalg.solve(R, f)

        # Ref: https://numpy.org/doc/stable/reference/generated/numpy.vdot.html
        return np.vdot(x,f)

    def plot_spectrum(self, omega_array, power_spectrum, rabi_chisq, omega_truth=None):
        plt.figure(2, figsize=(10, 7))
        freq_plot = np.copy(omega_array)
        rabi_power_plot = np.copy(power_spectrum)
        rabi_chisq_plot = np.copy(rabi_chisq)
        xlabel_plot = 'Frequency'

        plt.subplot(311)
        plt.title(r'Blue = Control in $|0\rangle$,  Red = Control in $|1\rangle$')
        plt.plot(freq_plot, rabi_power_plot[0, :], 'bo-')
        plt.plot(freq_plot, rabi_power_plot[1, :], 'r^-')

        if omega_truth is not None:
            plt.axvline(x=omega_truth[0], linestyle='--', c='b')
            plt.axvline(x=omega_truth[1], linestyle='--', c='r')
        # plt.yscale('log')
        plt.ylabel(r'$\langle X \rangle$')

        plt.subplot(312)
        plt.plot(freq_plot, rabi_power_plot[2, :], 'bo-')
        plt.plot(freq_plot, rabi_power_plot[3, :], 'r^-')

        if omega_truth is not None:
            plt.axvline(x=omega_truth[0], linestyle='--', c='b')
            plt.axvline(x=omega_truth[1], linestyle='--', c='r')
        # plt.yscale('log')
        plt.ylabel(r'$\langle Y \rangle$')

        plt.subplot(313)
        plt.plot(freq_plot, rabi_power_plot[4, :], 'bo-')
        plt.plot(freq_plot, rabi_power_plot[5, :], 'r^-')

        if omega_truth is not None:
            plt.axvline(x=omega_truth[0], linestyle='--', c='b')
            plt.axvline(x=omega_truth[1], linestyle='--', c='r')
        # plt.yscale('log')
        plt.ylabel(r'$\langle Z \rangle$')

        plt.xlabel(xlabel_plot)

        plt.show()

        # Plot chisq
        plt.figure(3, figsize=(10, 3))
        # plt.title('Blue = Control in |0> max @ '
        #          + str(np.argmax(rabi_chisq[0,:int((1 + rabi_chisq.shape[0]) / 2)]))
        #          + ',  Red = Control in |1> max @ '
        #          + str(np.argmax(rabi_chisq[1,:int((1 + rabi_chisq.shape[1]) / 2)])))

        plt.plot(freq_plot, rabi_chisq_plot[0, :], 'bo-')
        plt.plot(freq_plot, rabi_chisq_plot[1, :], 'r^-')

        if omega_truth is not None:
            plt.axvline(x=omega_truth[0], linestyle='--', c='b')
            plt.axvline(x=omega_truth[1], linestyle='--', c='r')
        plt.ylabel('Sum Magnitude Squared')
        plt.xlabel(xlabel_plot)

        plt.show()

    def corrected_power_spectrum(self, rabi_data=None, prior_freq=[0,1e8], do_plot=False, omega_truth=None):
        '''
        Solve the normal equations as laid out in Ed's slides

        TODO: Implement for the rotating exponential
        '''
        if rabi_data is None:
            rabi_data = self.pvec_data

        # Get the fourier transform of the rabi_data
        rabi_fourier, _, _ = fft_estimators.fft_rabi_data(rabi_data, 6)

        freq_temp = self.freq_convert*np.arange(rabi_data.shape[1])
        ind_cut = np.where(freq_temp > prior_freq[1])[0][0]

        rabi_f0 = rabi_fourier[:,0]

        if prior_freq is not None:
            ind_cut = np.where(freq_temp > prior_freq[1])[0][0]
            rabi_fourier = rabi_fourier[:,1:ind_cut]
            omega_array = freq_temp[1:ind_cut]
        else:
            omega_array = freq_temp[1:ind_cut]

        power_spectrum = np.ndarray(shape=rabi_fourier.shape, dtype=float)

        counter_omega = 0
        for omega in omega_array:
            for ind_config in range(6):
                f_temp = np.array([np.real(rabi_fourier[ind_config, counter_omega]),
                                   -np.imag(rabi_fourier[ind_config, counter_omega]),
                                   rabi_f0[ind_config]])

                power_spectrum[ind_config, counter_omega] = self.normal_equations(omega, f_temp)

            counter_omega += 1

        # Combine power series to form chi-squared specra for frequency identification
        rabi_chisq = np.zeros(shape=(2, rabi_fourier.shape[1]))
        rabi_chisq[0, :] = power_spectrum[0, :] + power_spectrum[2, :] + power_spectrum[4, :]
        rabi_chisq[1, :] = power_spectrum[1, :] + power_spectrum[3, :] + power_spectrum[5, :]

        if do_plot:
            self.plot_spectrum(omega_array, power_spectrum, rabi_chisq, omega_truth=omega_truth)

        return power_spectrum, rabi_chisq, omega_array

    def corrected_power_spectrum_uncertainty(self, rabi_data=None, prior_freq=[0,1e8], do_plot=False, omega_truth=None):
        '''
        Solve the normal equations as laid out in Ed's slides

        TODO: Implement for the rotating exponential
        '''
        if rabi_data is None:
            rabi_data = self.pvec_data

        # Get the fourier transform of the rabi_data
        alpha_rabi_data = self.n0shots_rabi_data + 1 / 2
        beta_rabi_data = (self.nshots_rabi_data - self.n0shots_rabi_data) + 1 / 2

        var_x = alpha_rabi_data*beta_rabi_data/ ( ((self.nshots_rabi_data + 1)**2) * (self.nshots_rabi_data + 2) )

        rabi_fourier, _, _ = fft_estimators.fft_rabi_data(rabi_data/var_x, 6)

        freq_temp = self.freq_convert*np.arange(rabi_data.shape[1])
        ind_cut = np.where(freq_temp > prior_freq[1])[0][0]

        rabi_f0 = rabi_fourier[:,0]

        if prior_freq is not None:
            ind_cut = np.where(freq_temp > prior_freq[1])[0][0]
            rabi_fourier = rabi_fourier[:,1:ind_cut]
            omega_array = freq_temp[1:ind_cut]
        else:
            omega_array = freq_temp[1:ind_cut]

        power_spectrum = np.ndarray(shape=rabi_fourier.shape, dtype=float)

        counter_omega = 0
        for omega in omega_array:
            for ind_config in range(6):
                f_temp = np.array([np.real(rabi_fourier[ind_config, counter_omega]),
                                   -np.imag(rabi_fourier[ind_config, counter_omega]),
                                   rabi_f0[ind_config]])

                power_spectrum[ind_config, counter_omega] = self.normal_equations_uncertainty(omega, f_temp, var_x[ind_config,:])

            counter_omega += 1

        # Combine power series to form chi-squared specra for frequency identification
        rabi_chisq = np.zeros(shape=(2, rabi_fourier.shape[1]))
        rabi_chisq[0, :] = power_spectrum[0, :] + power_spectrum[2, :] + power_spectrum[4, :]
        rabi_chisq[1, :] = power_spectrum[1, :] + power_spectrum[3, :] + power_spectrum[5, :]

        if do_plot:
            self.plot_spectrum(omega_array, power_spectrum, rabi_chisq, omega_truth=omega_truth)

        return power_spectrum, rabi_chisq, omega_array

    def corrected_power_spectrum_halfway(self, rabi_data=None, prior_freq=[0, 1e8], do_plot=False, omega_truth=None):
        '''
        Solve the normal equations as laid out in Ed's slides

        TODO: Implement for the rotating exponential
        '''
        if rabi_data is None:
            rabi_data = self.pvec_data

        n_configs, n_t = rabi_data.shape

        #rotating_exponential = np.array([np.exp(-1j * 2 * np.pi * self.sample_freq / (2 * n_t) * k) for k in range(n_t)])
        rotating_exponential = np.array([np.exp(-1j * 2 * np.pi * 1 / (2 * n_t) * k) for k in range(n_t)])
        #rotating_exponential = np.array([np.exp(-1j * 2 * np.pi * sample_freq / (2 * n_t) * _t) for _t in tvec])

        rabi_data_rot = np.zeros(shape=rabi_data.shape, dtype=complex)
        for ind_config in range(n_configs):
            rabi_data_rot[ind_config, :] = rabi_data[ind_config, :] * rotating_exponential

        # Get the fourier transform of the rabi_data
        rabi_fourier, _, _ = fft_rabi_data(rabi_data_rot, 6)

        freq_temp = self.freq_convert*(np.arange(rabi_data.shape[1]) + 1/2)
        ind_cut = np.where(freq_temp > prior_freq[1])[0][0]

        rabi_f0 = rabi_fourier[:,0]

        if prior_freq is not None:
            ind_cut = np.where(freq_temp > prior_freq[1])[0][0]
            rabi_fourier = rabi_fourier[:,0:ind_cut]
            omega_array = freq_temp[0:ind_cut]
        else:
            omega_array = freq_temp

        power_spectrum = np.ndarray(shape=rabi_fourier.shape, dtype=float)

        counter_omega = 0
        for omega in omega_array:
            for ind_config in range(6):
                f_temp = np.array([np.real(rabi_fourier[ind_config, counter_omega]),
                                   -np.imag(rabi_fourier[ind_config, counter_omega]),
                                   rabi_f0[ind_config]])

                power_spectrum[ind_config, counter_omega] = self.normal_equations(omega, f_temp)

            counter_omega += 1

        # Combine power series to form chi-squared specra for frequency identification
        rabi_chisq = np.zeros(shape=(2, rabi_fourier.shape[1]))
        rabi_chisq[0, :] = power_spectrum[0, :] + power_spectrum[2, :] + power_spectrum[4, :]
        rabi_chisq[1, :] = power_spectrum[1, :] + power_spectrum[3, :] + power_spectrum[5, :]

        if do_plot:
            self.plot_spectrum(omega_array, power_spectrum, rabi_chisq)

        return power_spectrum, rabi_chisq, omega_array

    def corrected_power_spectrum_halfway_uncertainty(self, rabi_data=None, prior_freq=[0,1e8], do_plot=False, omega_truth=None):
        '''
        Solve the normal equations as laid out in Ed's slides

        TODO: Implement for the rotating exponential
        '''
        if rabi_data is None:
            rabi_data = self.pvec_data

        # Get the fourier transform of the rabi_data
        alpha_rabi_data = self.n0shots_rabi_data + 1 / 2
        beta_rabi_data = (self.nshots_rabi_data - self.n0shots_rabi_data) + 1 / 2

        var_x = alpha_rabi_data*beta_rabi_data/ ( ((self.nshots_rabi_data + 1)**2) * (self.nshots_rabi_data + 2) )

        n_configs, n_t = rabi_data.shape

        # rotating_exponential = np.array([np.exp(-1j * 2 * np.pi * self.sample_freq / (2 * n_t) * k) for k in range(n_t)])
        rotating_exponential = np.array([np.exp(-1j * 2 * np.pi * 1 / (2 * n_t) * k) for k in range(n_t)])
        # rotating_exponential = np.array([np.exp(-1j * 2 * np.pi * sample_freq / (2 * n_t) * _t) for _t in tvec])

        rabi_data_rot = np.zeros(shape=rabi_data.shape, dtype=complex)
        for ind_config in range(n_configs):
            rabi_data_rot[ind_config, :] = rabi_data[ind_config, :] * rotating_exponential

        # Get the fourier transform of the rabi_data
        rabi_fourier, _, _ = fft_rabi_data(rabi_data_rot/var_x, 6)

        freq_temp = self.freq_convert * (np.arange(rabi_data.shape[1]) + 1 / 2)
        ind_cut = np.where(freq_temp > prior_freq[1])[0][0]

        rabi_f0 = rabi_fourier[:, 0]

        if prior_freq is not None:
            ind_cut = np.where(freq_temp > prior_freq[1])[0][0]
            rabi_fourier = rabi_fourier[:, 0:ind_cut]
            omega_array = freq_temp[0:ind_cut]
        else:
            omega_array = freq_temp

        power_spectrum = np.ndarray(shape=rabi_fourier.shape, dtype=float)

        counter_omega = 0
        for omega in omega_array:
            for ind_config in range(6):
                f_temp = np.array([np.real(rabi_fourier[ind_config, counter_omega]),
                                   -np.imag(rabi_fourier[ind_config, counter_omega]),
                                   rabi_f0[ind_config]])

                power_spectrum[ind_config, counter_omega] = self.normal_equations_uncertainty(omega, f_temp, var_x[ind_config,:])

            counter_omega += 1

        # Combine power series to form chi-squared specra for frequency identification
        rabi_chisq = np.zeros(shape=(2, rabi_fourier.shape[1]))
        rabi_chisq[0, :] = power_spectrum[0, :] + power_spectrum[2, :] + power_spectrum[4, :]
        rabi_chisq[1, :] = power_spectrum[1, :] + power_spectrum[3, :] + power_spectrum[5, :]

        if do_plot:
            self.plot_spectrum(omega_array, power_spectrum, rabi_chisq, omega_truth=omega_truth)

        return power_spectrum, rabi_chisq, omega_array

    def plot_corrected_power_spectrum(self, prior_freq=[0,5e7], omega_truth=None):
        # Normal FFT + Higher Resolution FFT
        rabi_power0, rabi_chisq0, omega_array0 = self.corrected_power_spectrum(prior_freq=prior_freq)
        rabi_power1, rabi_chisq1, omega_array1 = self.corrected_power_spectrum_halfway(prior_freq=prior_freq)

        # Sorting the collected arrays
        omega_array = np.concatenate((omega_array0, omega_array1))
        rabi_power = np.concatenate((rabi_power0, rabi_power1), axis=1)
        rabi_chisq = np.concatenate((rabi_chisq0, rabi_chisq1), axis=1)

        omega_sorted = np.sort(omega_array)
        sort_omega_index = np.argsort(omega_array)
        rabi_power_sorted = rabi_power[:, sort_omega_index]
        rabi_chisq_sorted = rabi_chisq[:, sort_omega_index]

        # Plot the Power Spectrum
        plt.figure(2, figsize=(10, 7))
        freq_plot = np.copy(omega_sorted)
        rabi_power_plot = np.copy(rabi_power_sorted)
        rabi_chisq_plot = np.copy(rabi_chisq_sorted)
        xlabel_plot = 'Frequency'

        if freq_plot[0] == 0:
            ind_start = 1
        else:
            ind_start = 0

        rabi_freq = np.array([
            freq_plot[np.argmax(rabi_chisq_plot[0, ind_start:int((rabi_chisq_plot.shape[1] - 1) / 2)]) + ind_start],
            freq_plot[np.argmax(rabi_chisq_plot[1, ind_start:int((rabi_chisq_plot.shape[1] - 1) / 2)]) + ind_start]])

        plt.subplot(311)
        plt.title(r'Blue = Control in $|0\rangle$,  Red = Control in $|1\rangle$')
        plt.semilogx(freq_plot, rabi_power_plot[0, :], 'bo-')
        plt.semilogx(freq_plot, rabi_power_plot[1, :], 'r^-')

        if omega_truth is not None:
            plt.axvline(x=omega_truth[0], linestyle='--', c='b')
            plt.axvline(x=omega_truth[1], linestyle='--', c='r')
        # plt.yscale('log')
        plt.ylabel(r'$\langle X \rangle$')

        plt.subplot(312)
        plt.semilogx(freq_plot, rabi_power_plot[2, :], 'bo-')
        plt.semilogx(freq_plot, rabi_power_plot[3, :], 'r^-')

        if omega_truth is not None:
            plt.axvline(x=omega_truth[0], linestyle='--', c='b')
            plt.axvline(x=omega_truth[1], linestyle='--', c='r')
        # plt.yscale('log')
        plt.ylabel(r'$\langle Y \rangle$')

        plt.subplot(313)
        plt.semilogx(freq_plot, rabi_power_plot[4, :], 'bo-')
        plt.semilogx(freq_plot, rabi_power_plot[5, :], 'r^-')

        if omega_truth is not None:
            plt.axvline(x=omega_truth[0], linestyle='--', c='b')
            plt.axvline(x=omega_truth[1], linestyle='--', c='r')
        # plt.yscale('log')
        plt.ylabel(r'$\langle Z \rangle$')

        plt.xlabel(xlabel_plot)

        plt.show()

        # Plot chisq
        plt.figure(3, figsize=(10, 3))
        # plt.title('Blue = Control in |0> max @ '
        #          + str(np.argmax(rabi_chisq[0,:int((1 + rabi_chisq.shape[0]) / 2)]))
        #          + ',  Red = Control in |1> max @ '
        #          + str(np.argmax(rabi_chisq[1,:int((1 + rabi_chisq.shape[1]) / 2)])))

        plt.semilogx(freq_plot, rabi_chisq_plot[0, :], 'bo-')
        plt.semilogx(freq_plot, rabi_chisq_plot[1, :], 'r^-')

        if omega_truth is not None:
            plt.axvline(x=omega_truth[0], linestyle='--', c='b')
            plt.axvline(x=omega_truth[1], linestyle='--', c='r')
        plt.ylabel('Sum Magnitude Squared')
        plt.xlabel(xlabel_plot)

        plt.show()

        print(rabi_freq)
        if omega_truth is not None:
            print([np.abs(rabi_freq[0] - omega_truth[0]) / omega_truth[0],
                   np.abs(rabi_freq[1] - omega_truth[1]) / omega_truth[1]])

    def plot_corrected_power_spectrum_uncertainty(self, prior_freq=[0,1e8], omega_truth=None):
        # Normal FFT + Higher Resolution FFT
        rabi_power0, rabi_chisq0, omega_array0 = self.corrected_power_spectrum_uncertainty()
        rabi_power1, rabi_chisq1, omega_array1 = self.corrected_power_spectrum_halfway_uncertainty()

        # Sorting the collected arrays
        omega_array = np.concatenate((omega_array0, omega_array1))
        rabi_power = np.concatenate((rabi_power0, rabi_power1), axis=1)
        rabi_chisq = np.concatenate((rabi_chisq0, rabi_chisq1), axis=1)

        omega_sorted = np.sort(omega_array)
        sort_omega_index = np.argsort(omega_array)
        rabi_power_sorted = rabi_power[:, sort_omega_index]
        rabi_chisq_sorted = rabi_chisq[:, sort_omega_index]

        # Plot the Power Spectrum
        plt.figure(2, figsize=(10, 7))
        freq_plot = np.copy(omega_sorted)
        rabi_power_plot = np.copy(rabi_power_sorted)
        rabi_chisq_plot = np.copy(rabi_chisq_sorted)
        xlabel_plot = 'Frequency'

        plt.subplot(311)
        plt.title(r'Blue = Control in $|0\rangle$,  Red = Control in $|1\rangle$')
        plt.plot(freq_plot, rabi_power_plot[0, :], 'bo-')
        plt.plot(freq_plot, rabi_power_plot[1, :], 'r^-')

        if omega_truth is not None:
            plt.axvline(x=omega_truth[0], linestyle='--', c='b')
            plt.axvline(x=omega_truth[1], linestyle='--', c='r')
        # plt.yscale('log')
        plt.ylabel(r'$\langle X \rangle$')

        plt.subplot(312)
        plt.plot(freq_plot, rabi_power_plot[2, :], 'bo-')
        plt.plot(freq_plot, rabi_power_plot[3, :], 'r^-')

        if omega_truth is not None:
            plt.axvline(x=omega_truth[0], linestyle='--', c='b')
            plt.axvline(x=omega_truth[1], linestyle='--', c='r')
        # plt.yscale('log')
        plt.ylabel(r'$\langle Y \rangle$')

        plt.subplot(313)
        plt.plot(freq_plot, rabi_power_plot[4, :], 'bo-')
        plt.plot(freq_plot, rabi_power_plot[5, :], 'r^-')

        if omega_truth is not None:
            plt.axvline(x=omega_truth[0], linestyle='--', c='b')
            plt.axvline(x=omega_truth[1], linestyle='--', c='r')
        # plt.yscale('log')
        plt.ylabel(r'$\langle Z \rangle$')

        plt.xlabel(xlabel_plot)

        plt.show()

        # Plot chisq
        plt.figure(3, figsize=(10, 3))
        # plt.title('Blue = Control in |0> max @ '
        #          + str(np.argmax(rabi_chisq[0,:int((1 + rabi_chisq.shape[0]) / 2)]))
        #          + ',  Red = Control in |1> max @ '
        #          + str(np.argmax(rabi_chisq[1,:int((1 + rabi_chisq.shape[1]) / 2)])))

        plt.plot(freq_plot, rabi_chisq_plot[0, :], 'bo-')
        plt.plot(freq_plot, rabi_chisq_plot[1, :], 'r^-')

        if omega_truth is not None:
            plt.axvline(x=omega_truth[0], linestyle='--', c='b')
            plt.axvline(x=omega_truth[1], linestyle='--', c='r')
        plt.ylabel('Sum Magnitude Squared')
        plt.xlabel(xlabel_plot)

        plt.show()

    def power_spectrum_linear_regression(self, omega_array, rabi_data=None, prior_freq=None):
        '''
        Assuming omega_array is sorted
        '''
        if rabi_data is None:
            rabi_data = self.pvec_data

        if prior_freq is not None:
            ind_cut = np.where(omega_array > prior_freq[1])[0][0]
            omega_vec = omega_array[0:ind_cut]
        else:
            omega_vec = np.copy(omega_array)

        # Get the power computed from the linear regression fits
        rabi_freqfit = np.zeros(shape=(2,len(omega_vec)))
        for k in range(len(omega_vec)):
            omega_temp = omega_vec[k]
            rabi_freqfit[0, k] = power_wrt_omega(omega_temp, 0, rabi_data, self.time_stamps)
            rabi_freqfit[1, k] = power_wrt_omega(omega_temp, 1, rabi_data, self.time_stamps)

        rabi_freq = np.zeros(2)

        if omega_vec[0] == 0:
            ind_start = 1
        else:
            ind_start = 0

        if prior_freq is None:
            rabi_freq[0] = omega_vec[np.argmax(rabi_freqfit[0, ind_start:int((rabi_freqfit.shape[1] - 1) / 2)]) + ind_start]
            rabi_freq[1] = omega_vec[np.argmax(rabi_freqfit[1, ind_start:int((rabi_freqfit.shape[1] - 1) / 2)]) + ind_start]
        else:
            rabi_freq[0] = omega_vec[np.argmax(rabi_freqfit[0, ind_start:]) + ind_start]
            rabi_freq[1] = omega_vec[np.argmax(rabi_freqfit[1, ind_start:]) + ind_start]

        return omega_vec, rabi_freqfit, rabi_freq

    def plot_power_spectrum_linear_regression(self, rabi_data=None,
                                              prior_freq=[0,5e8], omega_truth=None, do_plot=False, freq_resolution=4):
        if rabi_data is None:
            rabi_data = self.pvec_data

        n_t = len(self.time_stamps)
        omega_array = (self.freq_convert/freq_resolution)*np.arange(freq_resolution*n_t)

        omega_vec, rabi_chisq, rabi_freq = self.power_spectrum_linear_regression(omega_array, rabi_data=rabi_data,
                                                                                prior_freq=prior_freq)

        if do_plot:
            # Plot the Power Spectrum
            plt.figure(2, figsize=(10, 7))
            freq_plot = np.copy(omega_vec)
            rabi_chisq_plot = np.copy(rabi_chisq)
            xlabel_plot = 'Frequency'

            # Plot chisq
            plt.semilogx(freq_plot, rabi_chisq_plot[0, :], 'bo-', label=r'Control in $|0 \rangle$')
            plt.semilogx(freq_plot, rabi_chisq_plot[1, :], 'r^-', label=r'Control in $|1 \rangle$')

            plt.axvline(x=rabi_freq[0], linestyle='--', c='b')
            plt.axvline(x=rabi_freq[1], linestyle='--', c='r')

            if omega_truth is not None:
                plt.axvline(x=omega_truth[0], linestyle='-', c='b')
                plt.axvline(x=omega_truth[1], linestyle='-', c='r')

            plt.ylabel('Power')
            plt.xlabel(xlabel_plot)
            plt.legend(loc='best')

            plt.show()

            print(rabi_freq)
            if omega_truth is not None:
                print([np.abs(rabi_freq[0] - omega_truth[0]) / omega_truth[0],
                       np.abs(rabi_freq[1] - omega_truth[1]) / omega_truth[1]])

        return {'omega_array': omega_vec, 'rabi_chisq': rabi_chisq, 'rabi_freq': rabi_freq}

    def plot_corrected_power_spectrum_check(self, prior_freq=[0,1e8], omega_truth=None):
        # Get Normal FFT + Higher Resolution FFT
        rabi_power0, rabi_chisq0, omega_array0 = self.corrected_power_spectrum(prior_freq=prior_freq)
        rabi_power1, rabi_chisq1, omega_array1 = self.corrected_power_spectrum_halfway(prior_freq=prior_freq)

        # Sorting the collected arrays
        omega_array = np.concatenate((omega_array0, omega_array1))
        rabi_power = np.concatenate((rabi_power0, rabi_power1), axis=1)
        rabi_chisq = np.concatenate((rabi_chisq0, rabi_chisq1), axis=1)

        omega_sorted = np.sort(omega_array)
        sort_omega_index = np.argsort(omega_array)
        rabi_power_sorted = rabi_power[:, sort_omega_index]
        rabi_chisq_sorted = rabi_chisq[:, sort_omega_index]

        # Get the power computed from the linear regression fits
        rabi_freqfit = np.zeros(shape=rabi_chisq_sorted.shape)
        for k in range(len(omega_sorted)):
            omega_temp = omega_sorted[k]
            rabi_freqfit[0, k] = power_wrt_omega(omega_temp, 0, self.pvec_data, self.time_stamps)
            rabi_freqfit[1, k] = power_wrt_omega(omega_temp, 1, self.pvec_data, self.time_stamps)

        # Plot the Power Spectrum
        plt.figure(2, figsize=(10, 10))
        freq_plot = np.copy(omega_sorted)
        rabi_power_plot = np.copy(rabi_power_sorted)
        rabi_chisq_plot = np.copy(rabi_chisq_sorted)
        xlabel_plot = 'Frequency'

        plt.subplot(211)
        # Plot chisq
        plt.semilogx(freq_plot, rabi_chisq_plot[0, :], 'bo-')
        plt.semilogx(freq_plot, rabi_chisq_plot[1, :], 'r^-')

        if omega_truth is not None:
            plt.axvline(x=omega_truth[0], linestyle='--', c='b')
            plt.axvline(x=omega_truth[1], linestyle='--', c='r')
        plt.ylabel('Normal Equations')
        plt.xlabel(xlabel_plot)

        plt.subplot(212)
        # Plot chisq
        plt.semilogx(freq_plot, rabi_freqfit[0, :], 'bo-')
        plt.semilogx(freq_plot, rabi_freqfit[1, :], 'r^-')

        if omega_truth is not None:
            plt.axvline(x=omega_truth[0], linestyle='--', c='b')
            plt.axvline(x=omega_truth[1], linestyle='--', c='r')
        plt.ylabel('power wrt omega')
        plt.xlabel(xlabel_plot)

        plt.show()

        if freq_plot[0] == 0:
            ind_start = 1
        else:
            ind_start = 0

        rabi_freq = np.array([
            freq_plot[np.argmax(rabi_chisq_plot[0, ind_start:int((rabi_chisq_plot.shape[1] - 1) / 2)]) + ind_start],
            freq_plot[np.argmax(rabi_chisq_plot[1, ind_start:int((rabi_chisq_plot.shape[1] - 1) / 2)]) + ind_start]])

        print(rabi_freq)
        if omega_truth is not None:
            print([np.abs(rabi_freq[0] - omega_truth[0]) / omega_truth[0],
                   np.abs(rabi_freq[1] - omega_truth[1]) / omega_truth[1]])

    def point_estimator_linear_regression(self, rabi_data=None, freq_resolution=4):
        """
        rabi_data: If you want to calculate the point estimate of rabi oscillations for something besides the default
        """
        if rabi_data is None:
            rabi_data = self.pvec_data

        # Take fourier transforms of the rabi oscillations
        results_freq = self.plot_power_spectrum_linear_regression(rabi_data=rabi_data,
                                                                  prior_freq=[0,5e8], do_plot=False,
                                                                  freq_resolution=freq_resolution)

        rabi_freq = results_freq['rabi_freq']

        # Get the coarse estimates of rabi_gpt
        rabi_freq, rabi_gpt = self.model_free_coarse_hamiltonian_parameters_estimator(rabi_freq)

        # Refine the estimate
        rabi_refined_gpt = self.model_free_fine_hamiltonian_parameters_estimator(rabi_freq, rabi_gpt)

        param_array = np.array([rabi_freq[0] / 2, 2 * (rabi_refined_gpt[0, 2] - np.pi / 4), rabi_refined_gpt[0, 1],
                                rabi_freq[1] / 2, 2 * (rabi_refined_gpt[1, 2] - np.pi / 4), rabi_refined_gpt[1, 1]])

        return param_array

    def uncertainty_estimator(self, n_MC=100, do_plot=False):
        """
        n_MC: number of MC samples to consider from the sampling noise to be added to the pvec_data
        """
        p1_data = (1-self.pvec_data)/2

        param_vec = np.zeros(shape=(n_MC,6), dtype=float)

        # One should be as usual
        param_vec[0, :] = self.point_estimator()

        for ind_MC in range(1,n_MC):
            # Generate noisy samples
            noise_temp = np.random.binomial(self.nshots_rabi_data, p1_data)
            rabi_data_temp = 1.0 - 2.0*(noise_temp/self.nshots_rabi_data)

            # compute the point estimator
            param_vec[ind_MC,:] = self.point_estimator(rabi_data=rabi_data_temp)

        if do_plot:
            # Plot the histogram over omegas
            fig = plt.figure(0, figsize=(12, 4))
            ax1 = fig.add_subplot(121)
            ax1.hist(param_vec[:,0]/(1e6), bins=int(n_MC / 4), c='b')
            ax1.set(xlabel=r'$\omega_0 (\times 10^6)$', ylabel='Count')

            ax2 = fig.add_subplot(122)
            ax2.hist(param_vec[:,3]/(1e6), bins=int(n_MC / 4), c='r')
            ax2.set(xlabel=r'$\omega_1 (\times 10^6)$', ylabel='Count')

        return param_vec

    def uncertainty_estimator_linear_regression(self, n_MC=100, freq_resolution=4, do_plot=False):
        """
        n_MC: number of MC samples to consider from the sampling noise to be added to the pvec_data
        """
        p1_data = (1-self.pvec_data)/2

        param_vec = np.zeros(shape=(n_MC,6), dtype=float)
        rabi_freq_vec = np.zeros(shape=(n_MC, 2), dtype=float)

        # Create array of rabi oscillations
        rabi_data_vec = []

        # One should be as usual
        rabi_data_vec.append(self.pvec_data)
        results_freq_temp = self.plot_power_spectrum_linear_regression(rabi_data=self.pvec_data,
                                                                       prior_freq=[0, 5e8], do_plot=False,
                                                                       freq_resolution=freq_resolution)
        rabi_freq_vec[0, :] = results_freq_temp['rabi_freq']

        for ind_MC in range(1,n_MC):
            # Generate noisy samples
            noise_temp = np.random.binomial(self.nshots_rabi_data, p1_data)
            rabi_data_temp = 1.0 - 2.0*(noise_temp/self.nshots_rabi_data)

            rabi_data_vec.append(rabi_data_temp)

            # compute the point estimator
            results_freq_temp = self.plot_power_spectrum_linear_regression(rabi_data=rabi_data_temp,
                                                                           prior_freq=[0, 5e8], do_plot=False,
                                                                           freq_resolution=freq_resolution)
            rabi_freq_vec[ind_MC,:] = results_freq_temp['rabi_freq']

            if do_plot:
                if np.mod(ind_MC, 5) == 0:
                    print('%d' % ind_MC)

        # Fit log-normal distribution to rabi_freq_vec
        mu0 = np.mean(np.log(rabi_freq_vec[:, 0] / 1e6))
        mu1 = np.mean(np.log(rabi_freq_vec[:, 1] / 1e6))
        sigma0 = np.std(np.log(rabi_freq_vec[:, 0] / 1e6))
        sigma1 = np.std(np.log(rabi_freq_vec[:, 1] / 1e6))

        # Sample from the above log-normal distributions
        rabi_freq_vec_samples = np.zeros(shape=(n_MC,2))
        rabi_freq_vec_samples[:,0] = 1e6*np.random.lognormal(mu0, sigma0, n_MC)
        rabi_freq_vec_samples[:,1] = 1e6*np.random.lognormal(mu1, sigma1, n_MC)

        # Carry out linear regression for each sample
        for ind_MC in range(n_MC):
            rabi_freq = rabi_freq_vec_samples[ind_MC,:]
            rabi_data_temp = rabi_data_vec[ind_MC]

            # Get the coarse estimates of rabi_gpt
            rabi_freq, rabi_gpt = self.model_free_coarse_hamiltonian_parameters_estimator(rabi_freq,
                                                                                rabi_data=rabi_data_temp)

            # Refine the estimate
            rabi_refined_gpt = self.model_free_fine_hamiltonian_parameters_estimator(rabi_freq, rabi_gpt,
                                                                                rabi_data=rabi_data_temp)

            param_vec[ind_MC,:] = np.array([rabi_freq[0] / 2, 2 * (rabi_refined_gpt[0, 2] - np.pi / 4), rabi_refined_gpt[0, 1],
                                            rabi_freq[1] / 2, 2 * (rabi_refined_gpt[1, 2] - np.pi / 4), rabi_refined_gpt[1, 1]])

            if do_plot:
                if np.mod(ind_MC, 5) == 0:
                    print('%d\n' % ind_MC)

        if do_plot:
            # Plot the histogram over omegas
            fig = plt.figure(0, figsize=(14, 4))
            ax1 = fig.add_subplot(121)
            ax1.hist(param_vec[:,0]/(1e6), color='skyblue')
            ax1.set(xlabel=r'$\omega_0 (\times 10^6)$', ylabel='Count')

            ax2 = fig.add_subplot(122)
            ax2.hist(param_vec[:,3]/(1e6), color='r')
            ax2.set(xlabel=r'$\omega_1 (\times 10^6)$', ylabel='Count')

        return param_vec

    def uncertainty_estimator_bayesian(self, n_MC=100, do_plot=False, param_truth=None, freq_resolution=8):
        """
        n_MC: number of MC samples to consider from the sampling noise to be added to the pvec_data
        """
        param_vec = np.zeros(shape=(n_MC,6), dtype=float)

        # Set up the parameters of the beta distribution
        alpha_rabi_data = self.n0shots_rabi_data + 1/2
        beta_rabi_data = (self.nshots_rabi_data - self.n0shots_rabi_data) + 1/2

        # One should be as usual
        if self.FLAG_fft_high_resolution:
            param_vec[0, :] = self.point_estimator_linear_regression(freq_resolution=freq_resolution)
        else:
            param_vec[0, :] = self.point_estimator()

        for ind_MC in range(1,n_MC):
            # Generate noisy samples
            noise_temp = np.random.beta(alpha_rabi_data, beta_rabi_data)

            rabi_data_temp = 2.0*noise_temp - 1

            # compute the point estimator
            if self.FLAG_fft_high_resolution:
                param_vec[ind_MC, :] = self.point_estimator_linear_regression(rabi_data=rabi_data_temp,freq_resolution=freq_resolution)
            else:
                param_vec[ind_MC,:] = self.point_estimator(rabi_data=rabi_data_temp)

        if do_plot:
            type_distribution_fit = 'log_normal'

            if type_distribution_fit is 'log_normal':
                mu0 = np.mean(np.log(param_vec[:, 0] / 1e6))
                mu1 = np.mean(np.log(param_vec[:, 3] / 1e6))
                sigma0 = np.std(np.log(param_vec[:, 0] / 1e6))
                sigma1 = np.std(np.log(param_vec[:, 3] / 1e6))

                x0 = np.linspace(scipy.stats.lognorm.ppf(0.001, sigma0, scale=np.exp(mu0)), scipy.stats.lognorm.ppf(0.999, sigma0, scale=np.exp(mu0)), 1000)
                x1 = np.linspace(scipy.stats.lognorm.ppf(0.001, sigma1, scale=np.exp(mu0)), scipy.stats.lognorm.ppf(0.999, sigma1, scale=np.exp(mu1)), 1000)

                y0 = scipy.stats.lognorm.pdf(x0, sigma0, scale=np.exp(mu0))
                y1 = scipy.stats.lognorm.pdf(x1, sigma1, scale=np.exp(mu1))

                print(np.exp(mu0))
                print(sigma0)
                print(np.exp(mu1))
                print(sigma1)
            else:
                mu0 = np.mean(param_vec[:, 0] / 1e6)
                mu1 = np.mean(param_vec[:, 3] / 1e6)
                sigma0 = np.std(param_vec[:, 0] / 1e6)
                sigma1 = np.std(param_vec[:, 3] / 1e6)

                x0 = np.linspace(mu0 - 3 * sigma0, mu0 + 3 * sigma0, 100)
                x1 = np.linspace(mu1 - 3 * sigma1, mu1 + 3 * sigma1, 100)

                y0 = scipy.stats.norm.pdf(x0, mu0, sigma0)
                y1 = scipy.stats.norm.pdf(x1, mu1, sigma1)

                print(mu0)
                print(sigma0)
                print(mu1)
                print(sigma1)

            fig = plt.figure(0, figsize=(14, 4))
            ax1 = fig.add_subplot(121)
            ax1.hist(param_vec[:,0]/(1e6), bins=int(n_MC / 4), color='b')
            ax1.set_xscale('log')
            ax1.plot(x0, y0, 'k')

            if param_truth is not None:
                plt.axvline(x=param_truth[0]/1e6, linestyle='--', c='k')

            ax1.set(xlabel=r'$\omega_0 (\times 10^6)$', ylabel='Count')

            ax2 = fig.add_subplot(122)
            ax2.hist(param_vec[:,3]/(1e6), bins=int(n_MC / 4), color='r')
            ax2.set_xscale('log')
            ax2.plot(x1, y1, 'k')

            if param_truth is not None:
                plt.axvline(x=param_truth[1]/1e6, linestyle='--', c='k')

            ax2.set(xlabel=r'$\omega_1 (\times 10^6)$', ylabel='Count')

            fig = plt.figure(1, figsize=(16, 4))
            ax1 = fig.add_subplot(121)
            ax1.plot(x0, y0, 'k')
            ax1.set_xscale('log')

            if param_truth is not None:
                plt.axvline(x=param_truth[0] / 1e6, linestyle='--', c='k')

            ax1.set(xlabel=r'$\omega_0 (\times 10^6)$')

            ax2 = fig.add_subplot(122)
            ax2.plot(x1, y1, 'k')
            ax2.set_xscale('log')

            if param_truth is not None:
                plt.axvline(x=param_truth[1] / 1e6, linestyle='--', c='k')

            ax2.set(xlabel=r'$\omega_1 (\times 10^6)$')

        return param_vec

    def param_distribution(self):
        '''
        Get the log-normal distribution for omega's or the main different values of omega
        '''
    def line_search_estimator(self, prep, FLAG_mle_solve=False, verbose=False, FLAG_debug=False, do_plot=False):
        # Consider a range of frequencies
        N_t = len(self.time_stamps)
        set_frequencies = np.array([ind*(self.sample_freq/(2*N_t)) for ind in range(N_t + 1)])
        n_freqs = len(set_frequencies)

        # For each value of frequency, get an estimate of the rabi amplitudes A, B, C based on regression
        set_rabi_abc = []
        for freq in set_frequencies:
            _, rabi_abc_temp = self.model_free_coarse_rabi_amplitude_estimator(freq, prep, FLAG_refine_rabi_freq=False)
            set_rabi_abc.extend([rabi_abc_temp])

        # Use these as seed values to MLE estimator to get finer estimates of A, B and C
        config_array = np.array([0, 2, 4]) + prep

        if FLAG_mle_solve:
            for ind_freq in range(n_freqs):
                if FLAG_debug:
                    print(ind_freq)

                counter = 0
                for ind_config in config_array:
                    ds = self.rabi_amplitude_mle_estimator(set_rabi_abc[ind_freq][counter,:],
                                                           set_frequencies[ind_freq],
                                                           ind_config, verbose=verbose, do_plot=do_plot)

                    set_rabi_abc[ind_freq][counter,:] = ds['results'][1]

        # Obtain a probability distribution over set of discrete frequencies
        p_freq = np.zeros(n_freqs)
        for ind_freq in range(n_freqs):
            log_temp_freq = 1
            counter = 0
            for ind_config in config_array:
                p0_temp = np_model_free_probabilities(set_rabi_abc[ind_freq][counter,:],
                                                      set_frequencies[ind_freq],
                                                      self.tvec_config[ind_config])

                Y_temp = self.samples_config[ind_config]
                log_p_samples_temp = np.sum(np.log( (1-Y_temp)*p0_temp + Y_temp*(1 - p0_temp) + 1e-8 ))

                log_temp_freq += log_p_samples_temp
                counter += 1
                print('ind_freq: %d, counter: %d, config: %d, log_p_samples: %f' %(ind_freq, counter, ind_config, log_p_samples_temp))

                # import pdb; pdb.set_trace()
            p_freq[ind_freq] = np.exp(log_temp_freq)

        p_freq = p_freq/(np.sum(p_freq) + 1e-8)

        # Sample from this set and run the rabi amplitude estimator for each value

        return set_frequencies, p_freq, set_rabi_abc

class PointEstimator(object):
    """
    Use rabi oscillations for estimation of Hamiltonian parameters
    1. Primarily uses functions from "fft_estimators" and "linear_estimators"
    2. Doesn't incorporate uncertainty

    TODO: Add options for different parameterizations
    """

    def __init__(self, A_cr, do_plot=False,
                 FLAG_fft_window=False, type_fft_window=None,
                 FLAG_fft_high_resoultion=False,
                 FLAG_verbose=False):

        # Regarding the use of FFT
        self.FLAG_fft_window = FLAG_fft_window
        self.type_fft_window = type_fft_window
        self.FLAG_fft_high_resolution = FLAG_fft_high_resoultion

        # Extract relevant parameters
        time_stamps = A_cr.tset * A_cr.xi_t
        freq_convert = A_cr.freq_convert

        # Calculate the rabi oscillations from the given data
        pvec_data = A_cr.rabi_data()

        self.time_stamps = time_stamps
        self.pvec_data = pvec_data
        self.sample_freq = A_cr.sample_freq
        self.freq_convert = freq_convert

        # Other options
        self.do_plot = do_plot
        self.FLAG_verbose = FLAG_verbose

    def model_free_fine_rabi_frequency_estimator(self, rabi_freq, prep, rabi_data=None):
        """
        Given a value of omega for a particular prep op in A cos(wt) + B sin(wt) + C, refine this based on regression
        """
        # Calculate and store the initial frequency fit considering this value of omega for each of X, Y, Z meas ops

        if rabi_data is None:
            rabi_data = self.pvec_data

        rabi_prefit = np.ndarray(shape=(3, rabi_data.shape[1]), dtype=float)

        config_array = np.array([0, 2, 4]) + prep
        counter = 0
        for ind_config in config_array:
            cs_array = fft_estimators.cos_sin_features(rabi_freq, self.time_stamps)
            reg = LinearRegression().fit(cs_array, rabi_data[ind_config, :])
            rabi_prefit[counter] = reg.predict(cs_array)
            counter += 1

        # Refine the initial frequency estimates by fitting to the Rabi Oscillations
        ## INTRODUCE TRY AND CATCH WITH 2*freq_convert below
        bracketing_factors = np.arange(1,4)
        n_fails = 0
        while True:
            try:
                bracketing_half_width = bracketing_factors[n_fails]*self.freq_convert
                res = minimize_scalar(fft_estimators.sumsq_wrt_omega,
                                      bracket=(rabi_freq - bracketing_half_width,
                                               rabi_freq,
                                               rabi_freq + bracketing_half_width),
                                      args=(prep, rabi_data, self.time_stamps))
                break
            except ValueError:
                n_fails += 1

                if n_fails > 2:
                    print('Max no. of trials reached')
                    return None
                else:
                    print('ValueError, trying again with different backeting interval')

        return res.x

    def model_free_coarse_rabi_amplitude_estimator(self, rabi_freq, prep, rabi_data=None, FLAG_refine_rabi_freq=False):
        """
        Estimate the amplitudes i.e., A, B and C given particular values of omega in
        A cos(wt) + B sin(wt) + C

        Inputs:
        rabi_freq - scalar quantity here
        prep - index of the state preparation operator where U0 = II and U1 = XI, takes values 0 or 1
        """
        # Refine initial frequency estimates
        if rabi_data is None:
            rabi_data = self.pvec_data

        if FLAG_refine_rabi_freq:
            rabi_freq_refined = self.model_free_fine_rabi_frequency_estimator(rabi_freq, prep, rabi_data=rabi_data)

            # Indicative of error!
            if rabi_freq_refined is not None:
                rabi_freq = rabi_freq_refined

        config_array = np.array([0,2,4]) + prep

        # Calculate and store the fitted curves and corresponding Fourier coefficients
        rabi_cos_sin = np.ndarray(shape=(3,rabi_data.shape[1]), dtype=float)
        rabi_residuals = np.ndarray(shape=(3,rabi_data.shape[1]), dtype=float)
        rabi_abc = np.ndarray(shape=(3, 3), dtype=float)

        counter = 0
        for ind_config in config_array:
            cs_array = fft_estimators.cos_sin_features(rabi_freq, self.time_stamps)
            reg = LinearRegression().fit(cs_array, rabi_data[ind_config, :])
            rabi_abc[counter, 0] = reg.coef_[0]
            rabi_abc[counter, 1] = reg.coef_[1]
            rabi_abc[counter, 2] = reg.intercept_
            rabi_cos_sin[counter] = reg.predict(cs_array)
            rabi_residuals[counter] = rabi_data[ind_config, :] - rabi_cos_sin[counter]
            counter += 1

        return rabi_freq, rabi_abc

    def convert_rabi_amplitudes_to_hamiltonian_parameters(self, rabi_abc):
        """
        Given amplitudes A, B, C of rabi oscillations, obtain the hamiltonian parameters (as in the spectral param)

        rabi_abc is an array of amplitudes A, B, C for the different meas ops of X, Y, Z but only for a given prep op
        """
        # Calculate initial guesses for gamma, phi, theta
        rabi_gpt = np.ndarray(shape=(3,), dtype=float)

        g0 = -cmath.phase(complex(rabi_abc[2, 0], rabi_abc[2, 1]))
        expg0 = cmath.exp(complex(0, g0))
        x0 = expg0 * complex(rabi_abc[0, 0], rabi_abc[0, 1])
        y0 = expg0 * complex(rabi_abc[1, 0], rabi_abc[1, 1])
        z0 = expg0 * complex(rabi_abc[2, 0], rabi_abc[2, 1])
        rabi_gpt[0] = g0
        rabi_gpt[1] = cmath.phase(complex(-y0.imag, x0.imag))

        if z0.real > 1.0:
            rabi_gpt[2] = 0.25 * np.pi
        else:
            rabi_gpt[2] = (0.25 * np.pi) + 0.5 * np.sign(x0.real * y0.imag) * math.acos(z0.real)

        return rabi_gpt

    def model_free_coarse_hamiltonian_parameters_estimator(self, rabi_freq_array, rabi_data=None):
        if rabi_data is None:
            rabi_data = self.pvec_data

        # Get rabi amplitude estimates for each prep op
        rabi_freq_U0, rabi_abc_U0 = self.model_free_coarse_rabi_amplitude_estimator(rabi_freq_array[0], 0,
                                                                                rabi_data=rabi_data,
                                                                                FLAG_refine_rabi_freq=False)

        rabi_freq_U1, rabi_abc_U1 = self.model_free_coarse_rabi_amplitude_estimator(rabi_freq_array[1], 1,
                                                                                rabi_data=rabi_data,
                                                                                FLAG_refine_rabi_freq=False)

        # Convert to hamiltonian parameters (spectral param)
        rabi_gpt = np.ndarray(shape=(2, 3), dtype=float)
        rabi_gpt[0,:] = self.convert_rabi_amplitudes_to_hamiltonian_parameters(rabi_abc_U0)
        rabi_gpt[1,:] = self.convert_rabi_amplitudes_to_hamiltonian_parameters(rabi_abc_U1)

        # Return refined values of rabi_freq
        rabi_freq = np.array([rabi_freq_U0, rabi_freq_U1])

        return rabi_freq, rabi_gpt

    def model_free_fine_hamiltonian_parameters_estimator(self, rabi_freq, rabi_gpt, rabi_data=None):
        if rabi_data is None:
            rabi_data = self.pvec_data

        rabi_refined_gpt = np.ndarray(shape=(2, 3), dtype=float)

        res = minimize(linear_estimators.sumsq_wrt_gpt, rabi_gpt[0, :],
                       args=(0, rabi_data, rabi_freq, self.time_stamps),
                       method="trust-constr",
                       jac=False,  # grad_sumsq_wrt_gpt,
                       bounds=((-math.pi, math.pi),
                               (-math.pi, math.pi),
                               (0, 0.5 * math.pi))
                       )

        rabi_refined_gpt[0, :] = res.x

        res = minimize(linear_estimators.sumsq_wrt_gpt, rabi_gpt[1, :],
                       args=(1, rabi_data, rabi_freq, self.time_stamps),
                       method="trust-constr",
                       jac=False,  # grad_sumsq_wrt_gpt,
                       bounds=((-math.pi, math.pi),
                               (-math.pi, math.pi),
                               (0, 0.5 * math.pi))
                       )
        rabi_refined_gpt[1, :] = res.x

        return rabi_refined_gpt

    def baseline(self, rabi_data=None):
        """
        Baseline solver -- Doesn't refine the estimate of frequencies after FFT

        So basically FFT + Regression (as advertised in the paper)

        Inputs:
            rabi_data: rabi oscillations if not the default
        """
        if rabi_data is None:
            rabi_data = self.pvec_data

        # Take fourier transforms of the rabi oscillations
        results_freq = fft_estimators.calc_freq(rabi_data, self.freq_convert, sample_freq=self.sample_freq,
                                                do_plot=self.do_plot,
                                                FLAG_window=self.FLAG_fft_window, type_window=self.type_fft_window,
                                                FLAG_high_resolution=self.FLAG_fft_high_resolution)

        rabi_freq = results_freq['rabi_freq']

        if self.FLAG_verbose:
            print('FFT Frequency Estimates')
            print(rabi_freq)

        # Get the coarse estimates of rabi_gpt
        rabi_freq, rabi_gpt = self.model_free_coarse_hamiltonian_parameters_estimator(rabi_freq)

        # Refine the estimate
        rabi_refined_gpt = self.model_free_fine_hamiltonian_parameters_estimator(rabi_freq, rabi_gpt)

        if self.FLAG_verbose:
            print('Linear Regression Estimates')
            print(rabi_freq)

        param_array = np.array([rabi_freq[0] / 2, 2 * (rabi_refined_gpt[0, 2] - np.pi / 4), rabi_refined_gpt[0, 1],
                                rabi_freq[1] / 2, 2 * (rabi_refined_gpt[1, 2] - np.pi / 4), rabi_refined_gpt[1, 1]])

        return param_array

    def point_estimator(self, rabi_data=None):
        """
        Baseline solver and a bit more -- frequencies after FFT are refined using a regression step

        Inputs:
            rabi_data: rabi oscillations if not the default
        """
        if rabi_data is None:
            rabi_data = self.pvec_data

        # Take fourier transforms of the rabi oscillations
        results_freq = fft_estimators.calc_freq(rabi_data, self.freq_convert, sample_freq=self.sample_freq,
                                                do_plot=self.do_plot,
                                                FLAG_window=self.FLAG_fft_window, type_window=self.type_fft_window,
                                                FLAG_high_resolution=self.FLAG_fft_high_resolution)

        rabi_freq = results_freq['rabi_freq']

        if self.FLAG_verbose:
            print('FFT Frequency Estimates')
            print(rabi_freq)

        # Refine the estimate of frequencies
        for prep_ind in [0, 1]:
            rabi_freq[prep_ind] = self.model_free_fine_rabi_frequency_estimator(rabi_freq[prep_ind], prep_ind)

        if self.FLAG_verbose:
            print('FFT Frequency Estimates')
            print(rabi_freq)

        # Get the coarse estimates of rabi_gpt
        rabi_freq, rabi_gpt = self.model_free_coarse_hamiltonian_parameters_estimator(rabi_freq)

        # Refine the estimate
        rabi_refined_gpt = self.model_free_fine_hamiltonian_parameters_estimator(rabi_freq, rabi_gpt)

        if self.FLAG_verbose:
            print('Linear Regression Estimates')
            print(rabi_freq)

        param_array = np.array([rabi_freq[0] / 2, 2 * (rabi_refined_gpt[0, 2] - np.pi / 4), rabi_refined_gpt[0, 1],
                                rabi_freq[1] / 2, 2 * (rabi_refined_gpt[1, 2] - np.pi / 4), rabi_refined_gpt[1, 1]])

        return param_array


import io
import csv
import logging
import requests
import glob, os
import numpy as np
import pandas as pd
from pathlib import Path
from astropy.table import Table
from astropy import units as u
from astropy.time import Time
from astropy.coordinates import SkyCoord
import matplotlib.pyplot as plt

# ["g", "r", "i", "z"], value=[1, 2, 3, 4]
dic_filters = {1: "g", 2: "r", 3: "i", 4: "z"}

def i_minus_g(fid, mag):
    """ Compute i-g based on vectors of filters and magnitudes
    """

    if list(set(fid)) == [1, 3]:
        # last measurement
        last_fid = fid[-1]

        # last measurement with different filter
        # could take the mean
        index_other = np.where(np.array(fid) != last_fid)[0][-1]

        sign = np.diff([fid[index_other], last_fid])[0]
        mag = [mag[index_other], mag[-1]]

        return -1 * sign * np.diff(mag)[0]

    else:
        return np.nan

def extract_delta_color(pdf: pd.DataFrame, smooth_by: float):
    """ Extract last g-r and delta mag for each object

    Modified from Fink https://github.com/astrolabsoftware/fink-science-portal/apps/utils.py
    Adding a grouping for a given mjd interval

    Parameters
    ----------
    pdf: pandas DataFrame
        DataFrame containing magnitude, mjd, filters
        Filter band, as integer. g: 1, r: 2, i: 3, z:4
    smooth_by: float
        MJD delta to smooth computation

    Returns
    ----------

    """

    # dmag, rate for each filter
    dic_dmag = dict.fromkeys([1, 2, 3, 4])
    dic_dmag_mjd = dict.fromkeys([1, 2, 3, 4])
    dic_rate = dict.fromkeys([1, 2, 3, 4])
    for fil in pdf["filter"].unique():
        subpdf = pdf[pdf["filter"] == fil]
        subpdf = subpdf.sort_values("mjd", ascending=False)

        # dmag smoothed
        min_mjd = float(int(subpdf.mjd.min()))
        max_mjd = float(int(subpdf.mjd.max()))
        bins = np.arange(min_mjd, max_mjd + 1, smooth_by)
        df_grouped = subpdf.groupby(np.digitize(subpdf.mjd, bins)).median()
        mag_grouped = df_grouped["magnitude"]
        mjd_grouped = df_grouped["mjd"]
        if len(mag_grouped) > 1:
            # only compute if more than one observation
            dmag_ = mag_grouped.diff(periods=1).values[1:]
            dmag_mjd_ = mjd_grouped.values[1:].astype(int)
            dic_dmag[fil] = dmag_
            dic_dmag_mjd[fil] = dmag_mjd_

        # Rate by day
        bins = np.arange(min_mjd, max_mjd + 1, 1)
        df_grouped_byday = subpdf.groupby(np.digitize(subpdf.mjd, bins)).median()
        mag_grouped_byday = df_grouped_byday["magnitude"]
        rate_ = mag_grouped_byday.diff(periods=1).values
        if len(mag_grouped_byday) > 1:
            dic_rate[fil] = rate_[1:]

    # for color (can be modified to take median mag)
    # group by night
    gpdf = pdf.groupby("mjd_int")[["filter", "mjd", "magnitude"]].agg(list)
    # take only nights with at least measurements on 2 different filters
    mask = gpdf["filter"].apply(lambda x: (len(x) > 1) & (np.sum(x) / len(x) != x[0]))
    gpdf_night = gpdf[mask]

    # compute i-g for those nights
    color_tmp = [
        i_minus_g(i, j)
        for i, j in zip(gpdf_night["filter"].values, gpdf_night["magnitude"].values)
    ]
    mask = np.argwhere(~np.isnan(color_tmp)).flatten()
    if len(mask) > 1:
        color_mjd = gpdf_night.index.values[mask]
        color = [color_tmp[k] for k in mask]
    else:
        color_mjd = []
        color = []

    return (dic_dmag, dic_dmag_mjd, dic_rate, color, color_mjd)

if __name__ == "__main__":
    inpath = "./S82sub8_59.12"
    list_files = glob.glob(f"{inpath}/*.forced.difflc.txt")

    # If daily cadence then = rate
    smooth_by = 1

    for fname in list_files[:1]:
        # read file and convert to pandas
        df_tmp = Table.read(fname, format="ascii").to_pandas()

        # variable reformatting
        t = Time(df_tmp["dateobs"].to_list(), format="isot", scale="utc")
        df_tmp["mjd"] = t.mjd.astype(float)
        df_tmp["mjd_int"] = t.mjd.astype(int)
        df_tmp["magnitude"] = df_tmp["m"].values
        df_tmp["filter"] = df_tmp["filt"].replace(
            to_replace=["g", "r", "i", "z"], value=[1, 2, 3, 4]
        )
        df_tmp = df_tmp[df_tmp["magnitude"] != "-"]
        df_tmp["magnitude"] = df_tmp["magnitude"].astype(np.float).copy()

        dic_dmag, dic_dmag_mjd, dic_rate, color, color_mjd = extract_delta_color(
            df_tmp, smooth_by=smooth_by
        )

        import ipdb

        ipdb.set_trace()

        # # Plot to verify processing
        # from matplotlib import gridspec

        # fig = plt.figure(figsize=(14, 14))
        # gs = gridspec.GridSpec(2, 1, height_ratios=[2, 1])
        # ax0 = plt.subplot(gs[0, 0])
        # ax1 = plt.subplot(gs[1, 0])
        # x_arr = np.arange(df_tmp.mjd.min(), df_tmp.mjd.max() + 1, 1)
        # ax1.plot(x_arr, np.zeros(len(x_arr)), color="grey", linestyle="-")

        # for fil in df_tmp["filter"].unique():
        #     sel = df_tmp[df_tmp["filter"] == fil]
        #     if len(sel) > 1:
        #         ax0.errorbar(
        #             sel.mjd,
        #             sel.magnitude,
        #             yerr=sel.dm.astype(float),
        #             fmt="o",
        #             label=dic_filters[fil],
        #         )
        #         if fil in dic_dmag_mjd.keys():
        #             ax1.scatter(dic_dmag_mjd[fil], dic_dmag[fil])
        #             if len(color) > 0:
        #                 ax1.scatter(color_mjd, color, color="black")
        # idx = Path(fname).stem.replace(".forced.difflc", "")
        # ax0.legend()
        # plt.savefig(f"{idx}.png")


"""Test initialization of tedee."""
from http import HTTPStatus
from typing import Any
from unittest.mock import MagicMock
from urllib.parse import urlparse

from pytedee_async.exception import (
    TedeeAuthException,
    TedeeClientException,
    TedeeWebhookException,
)
import pytest
from syrupy import SnapshotAssertion

from homeassistant.components.webhook import async_generate_url
from homeassistant.config_entries import ConfigEntryState
from homeassistant.const import EVENT_HOMEASSISTANT_STOP
from homeassistant.core import HomeAssistant
from homeassistant.helpers import device_registry as dr

from .conftest import WEBHOOK_ID

from tests.common import MockConfigEntry
from tests.typing import ClientSessionGenerator

async def test_load_unload_config_entry(
    hass: HomeAssistant,
    mock_config_entry: MockConfigEntry,
    mock_tedee: MagicMock,
) -> None:
    """Test loading and unloading the integration."""
    mock_config_entry.add_to_hass(hass)
    await hass.config_entries.async_setup(mock_config_entry.entry_id)
    await hass.async_block_till_done()

    assert mock_config_entry.state is ConfigEntryState.LOADED

    await hass.config_entries.async_unload(mock_config_entry.entry_id)
    await hass.async_block_till_done()

    assert mock_config_entry.state is ConfigEntryState.NOT_LOADED

@pytest.mark.parametrize(
    "side_effect", [TedeeClientException(""), TedeeAuthException("")]
)
async def test_config_entry_not_ready(
    hass: HomeAssistant,
    mock_config_entry: MockConfigEntry,
    mock_tedee: MagicMock,
    side_effect: Exception,
) -> None:
    """Test the Tedee configuration entry not ready."""
    mock_tedee.get_locks.side_effect = side_effect

    mock_config_entry.add_to_hass(hass)
    await hass.config_entries.async_setup(mock_config_entry.entry_id)
    await hass.async_block_till_done()

    assert len(mock_tedee.get_locks.mock_calls) == 1
    assert mock_config_entry.state is ConfigEntryState.SETUP_RETRY

async def test_cleanup_on_shutdown(
    hass: HomeAssistant,
    mock_config_entry: MockConfigEntry,
    mock_tedee: MagicMock,
) -> None:
    """Test the webhook is cleaned up on shutdown."""
    mock_config_entry.add_to_hass(hass)
    await hass.config_entries.async_setup(mock_config_entry.entry_id)
    await hass.async_block_till_done()

    assert mock_config_entry.state is ConfigEntryState.LOADED

    hass.bus.async_fire(EVENT_HOMEASSISTANT_STOP)
    await hass.async_block_till_done()
    mock_tedee.delete_webhook.assert_called_once()

async def test_webhook_cleanup_errors(
    hass: HomeAssistant,
    mock_config_entry: MockConfigEntry,
    mock_tedee: MagicMock,
    caplog: pytest.LogCaptureFixture,
) -> None:
    """Test the webhook is cleaned up on shutdown."""
    mock_config_entry.add_to_hass(hass)
    await hass.config_entries.async_setup(mock_config_entry.entry_id)
    await hass.async_block_till_done()

    assert mock_config_entry.state is ConfigEntryState.LOADED

    mock_tedee.delete_webhook.side_effect = TedeeWebhookException("")

    hass.bus.async_fire(EVENT_HOMEASSISTANT_STOP)
    await hass.async_block_till_done()
    mock_tedee.delete_webhook.assert_called_once()
    assert "Failed to unregister Tedee webhook from bridge" in caplog.text

async def test_webhook_registration_errors(
    hass: HomeAssistant,
    mock_config_entry: MockConfigEntry,
    mock_tedee: MagicMock,
    caplog: pytest.LogCaptureFixture,
) -> None:
    """Test the webhook is cleaned up on shutdown."""
    mock_tedee.register_webhook.side_effect = TedeeWebhookException("")
    mock_config_entry.add_to_hass(hass)
    await hass.config_entries.async_setup(mock_config_entry.entry_id)
    await hass.async_block_till_done()

    assert mock_config_entry.state is ConfigEntryState.LOADED

    mock_tedee.register_webhook.assert_called_once()
    assert "Failed to register Tedee webhook from bridge" in caplog.text

async def test_bridge_device(
    hass: HomeAssistant,
    mock_config_entry: MockConfigEntry,
    mock_tedee: MagicMock,
    device_registry: dr.DeviceRegistry,
    snapshot: SnapshotAssertion,
) -> None:
    """Ensure the bridge device is registered."""
    mock_config_entry.add_to_hass(hass)
    await hass.config_entries.async_setup(mock_config_entry.entry_id)
    await hass.async_block_till_done()

    device = device_registry.async_get_device(
        {(mock_config_entry.domain, mock_tedee.get_local_bridge.return_value.serial)}
    )
    assert device
    assert device == snapshot

@pytest.mark.parametrize(
    ("body", "expected_code", "side_effect"),
    [
        ({"hello": "world"}, HTTPStatus.OK, None),  # Success
        (None, HTTPStatus.BAD_REQUEST, None),  # Missing data
        ({}, HTTPStatus.BAD_REQUEST, TedeeWebhookException),  # Error
    ],
)
async def test_webhook_post(
    hass: HomeAssistant,
    mock_config_entry: MockConfigEntry,
    mock_tedee: MagicMock,
    hass_client_no_auth: ClientSessionGenerator,
    body: dict[str, Any],
    expected_code: HTTPStatus,
    side_effect: Exception,
) -> None:
    """Test webhook callback."""

    mock_config_entry.add_to_hass(hass)
    await hass.config_entries.async_setup(mock_config_entry.entry_id)
    await hass.async_block_till_done()

    client = await hass_client_no_auth()
    webhook_url = async_generate_url(hass, WEBHOOK_ID)
    mock_tedee.parse_webhook_message.side_effect = side_effect
    resp = await client.post(urlparse(webhook_url).path, json=body)

    # Wait for remaining tasks to complete.
    await hass.async_block_till_done()

    assert resp.status == expected_code


from tortoise.query_utils import Prefetch
from tortoise import Tortoise, fields, run_async

from app.core.init_app import TORTOISE_ORM
from app.applications.users.models import User, Group

from app.applications.users.schemas import BaseProperties

import logging

# logging.basicConfig(level=logging.DEBUG)
from aiocache import caches
from aiocache.serializers import PickleSerializer

async def run():
    await Tortoise.init(config=TORTOISE_ORM)
    cache_config = {
        "default": {
            "cache": "aiocache.RedisCache",
            "endpoint": "localhost",
            "port": 6379,
            "serializer": {
                "class": PickleSerializer
            }
        }
    }
    caches.set_config(cache_config)

    cache = caches.get("default")
    cache_key = "users"
    # 尝试从缓存中获取数据
    cached_data = await cache.get(cache_key)
    if cached_data:
        print("cached", cached_data)
        return cached_data

    # 如果缓存中没有数据，则从数据库中获取数据
    users = await User.all().values()

    # 将数据存入缓存
    await cache.set(cache_key, users, ttl=3600)  # 设置缓存过期时间为1小时

    # a = await Group.filter(name='admin').count()

    # if a == 0:
    #     g = await Group.create(name='admin')
    # else:
    #     g = await Group.get(name='admin').prefetch_related('users')

    # print(await g.fetch_related("users"))
    # await g.add_permission("test")
    # await g.save()
    # print(await g.to_dict())

    # admin_group = await Group.get(name='admin').prefetch_related('users').values('permissions', 'users__id')
    # print(admin_group)

    # user = await User.get(username='root@admin.com').prefetch_related('groups')
    # print(user.permissions)
    # print([g for g in user.groups])

    # print(permissions, user_ids)
    # users = await User.all()
    # for user in users:
    #     print(user)
    #     print("Groups:", await user.groups.all())
    #     if user.is_superuser:
    #         await user.groups.add(g)
    #         await user.save()

if __name__ == "__main__":
    run_async(run())


##### INIT API

from yikemng.api.api_setup import app

##### INIT MANAGER

import yikemng.bot_manager

##### Start server

from uvicorn.config import Config
from uvicorn.server import Server

if __name__ == '__main__':
    config = Config(
        'yikemng.__main__:app',
        port=8000,
        log_level='debug',
        reload=False,
        host='0.0.0.0',
        forwarded_allow_ips='*'
        )
    server = Server(config)
    server.run()