# -*- coding: utf-8 -*-
"""
Created on Thu Mar 11 14:58:18 2021

@author: Nick Sauerwein
"""
import lyse
import numpy as np

from HelperClasses.Plotting import WidgetFakeColorPlot, MainPlotPanel, Widget1DLivingPlot
from HelperClasses import DyTransition, Camera, AbsorptionImaging

import cProfile
import pstats

profile = cProfile.Profile()
profile.enable()

figure_manager = lyse.figure_manager

h5_paths = lyse.h5_paths()

print(h5_paths)

if lyse.spinning_top:
    # If so, use the filepath of the current shot
    h5_path = lyse.path
else:
    # If not, get the filepath of the last shot of the lyse DataFrame
    h5_path = lyse.h5_paths().iloc[-1]

if len(h5_paths):
    last_globals = run = lyse.Run(h5_paths.iloc[-1]).get_globals()

if not hasattr(figure_manager, 'ap'):
    figure_manager.MainPlotPanel = MainPlotPanel.MainPlotPanel(h5_paths)

# Add absorption imaging
plot_name = '3D-MOT Absorption Imaging - With Atoms'
if not plot_name in figure_manager.MainPlotPanel.plots:
    current_plotting = WidgetFakeColorPlot.WidgetFakeColorPlot(plot_name + '1', ('in_situ_absorption', 'atoms'),
                                                               data_group='MOT_3D_Camera',
                                                               Roi=False, colorbarText=False)
    figure_manager.MainPlotPanel.add_plot_dock(plot_name + '1', current_plotting)

# plot_name = '3D-MOT Absorption Imaging - Without Atoms'
# if not plot_name in figure_manager.MainPlotPanel.plots:
#     current_plotting = WidgetFakeColorPlot.WidgetFakeColorPlot(plot_name + '1', ('in_situ_absorption', 'background'),
#                                                                data_group='MOT_3D_Camera',
#                                                                Roi=False, colorbarText=False)
#     figure_manager.MainPlotPanel.add_plot_dock(plot_name + '1', current_plotting)
#
# plot_name = '3D-MOT Absorption Imaging - Dark'
# if not plot_name in figure_manager.MainPlotPanel.plots:
#     current_plotting = WidgetFakeColorPlot.WidgetFakeColorPlot(plot_name + '1', ('in_situ_absorption', 'dark'),
#                                                                data_group='MOT_3D_Camera',
#                                                                Roi=False, colorbarText=False)
#     figure_manager.MainPlotPanel.add_plot_dock(plot_name + '1', current_plotting)

# Add plotting for absorption imaging
plot_name = '3D-MOT Absorption Imaging'
if not plot_name in figure_manager.MainPlotPanel.plots:
    current_plotting = WidgetFakeColorPlot.WidgetFakeColorPlot(plot_name, ('absorption_imaging', 'absorption_imaging'))

    figure_manager.MainPlotPanel.add_plot_dock(plot_name, current_plotting)

    absorption_imaging_transition = DyTransition.creat_Dy421()
    mot_3D_camera = Camera.c11440_36u(absorption_imaging_transition['wavelength'])
    absorption_image = AbsorptionImaging.absorption_imaging(None, 'MOT_3D_Camera', 'in_situ_absorption',
                                                                absorption_imaging_transition, mot_3D_camera, 0, 0)
    figure_manager.MainPlotPanel.analyse_panel.add_absorption_imaging(absorption_image, current_plotting)

# Add realtime plotting for atom number
# plot_name = '3D-MOT Atom Number'
# if not plot_name in figure_manager.MainPlotPanel.plots:
#     current_plotting = Widget1DLivingPlot.Widget1DLivingPlot(plot_name, y_data_path=('absorption_imaging', 'atom_number'), y_in_results=True)
#
#     figure_manager.MainPlotPanel.add_plot_dock(plot_name, current_plotting)

figure_manager.MainPlotPanel.update_h5_paths(h5_paths)
figure_manager.MainPlotPanel.refresh(h5_path)

profile.disable()
ps = pstats.Stats(profile)
ps.sort_stats('cumtime')
ps.print_stats(10)