from PyQt5.QtWidgets import QMainWindow, QApplication, QMessageBox, QScrollArea, QButtonGroup, QComboBox, QFrame, \
    QCheckBox, QFileDialog, QGridLayout, QTextEdit, QLineEdit, QVBoxLayout, QWidget, QLabel, QHBoxLayout, QPushButton, \
    QRadioButton, QSpacerItem, QSizePolicy
from PyQt5.QtCore import Qt, QSize, QRect
from PyQt5.QtGui import QIcon, QDoubleValidator
from sklearn.preprocessing import MinMaxScaler

from celldetective.gui import Styles
from celldetective.gui.gui_utils import center_window, FeatureChoice, ListWidget, QHSeperationLine, FigureCanvas, \
    GeometryChoice, OperationChoice
from superqt import QLabeledSlider
from superqt.fonticon import icon
from fonticon_mdi6 import MDI6
from celldetective.utils import extract_experiment_channels, get_software_location, _extract_labels_from_config
from celldetective.io import interpret_tracking_configuration, load_frames, auto_load_number_of_frames, \
    load_experiment_tables, get_experiment_antibodies, get_experiment_cell_types, get_experiment_concentrations, \
    get_positions_in_well, get_experiment_wells
from celldetective.measure import compute_haralick_features, contour_of_instance_segmentation
from celldetective.signals import columnwise_mean, mean_signal
import numpy as np
from tifffile import imread
import json
from shutil import copyfile
import os
import matplotlib.pyplot as plt

plt.rcParams['svg.fonttype'] = 'none'
from mpl_toolkits.axes_grid1 import make_axes_locatable
from glob import glob
from natsort import natsorted
from tifffile import imread
from pathlib import Path, PurePath
import gc
import pandas as pd
from tqdm import tqdm
from lifelines import KaplanMeierFitter
from matplotlib.cm import viridis, tab10
import math
import seaborn as sns


class ConfigMeasurementsPlot(QWidget,Styles):
    """
    UI to set survival instructions.

    """

    def __init__(self, parent_window=None):

        super().__init__()
        self.parent_window = parent_window
        self.setWindowTitle("Configure signal plot")
        self.setWindowIcon(QIcon(os.sep.join(['celldetective', 'icons', 'mexican-hat.png'])))
        self.exp_dir = self.parent_window.exp_dir
        self.soft_path = get_software_location()
        self.exp_config = self.exp_dir + "config.ini"
        self.wells = np.array(self.parent_window.parent_window.wells, dtype=str)
        self.well_labels = _extract_labels_from_config(self.exp_config, len(self.wells))
        self.FrameToMin = self.parent_window.parent_window.FrameToMin
        self.float_validator = QDoubleValidator()
        self.target_class = [0, 1]
        self.show_ci = True
        self.show_cell_lines = False
        self.ax2 = None
        self.auto_close = False
        self.palette = sns.cubehelix_palette()
        # sns.color_palette("cubehelix", as_cmap=True)
        # sns.color_palette("ch:s=-.2,r=.6", as_cmap=True)

        print('Parent wells: ', self.wells)

        self.well_option = self.parent_window.parent_window.well_list.currentIndex()
        self.position_option = self.parent_window.parent_window.position_list.currentIndex()
        self.interpret_pos_location()
        # self.load_available_tables()
        # self.config_path = self.exp_dir + self.config_name

        self.screen_height = self.parent_window.parent_window.parent_window.screen_height
        center_window(self)
        # self.setMinimumHeight(int(0.8*self.screen_height))
        # self.setMaximumHeight(int(0.8*self.screen_height))
        self.populate_widget()
        # self.load_previous_measurement_instructions()
        if self.auto_close:
            self.close()
        self.hue_colors = {0: self.palette[0], 1: self.palette[3]}

    def interpret_pos_location(self):

        """
        Read the well/position selection from the control panel to decide which data to load
        Set position_indices to None if all positions must be taken

        """

        if self.well_option == len(self.wells):
            self.well_indices = np.arange(len(self.wells))
        else:
            self.well_indices = np.array([self.well_option], dtype=int)

        if self.position_option == 0:
            self.position_indices = None
        else:
            self.position_indices = np.array([self.position_option], dtype=int)

    def populate_widget(self):

        """
        Create the multibox design.

        """

        # Create button widget and layout
        main_layout = QVBoxLayout()
        self.setLayout(main_layout)
        main_layout.setContentsMargins(30, 30, 30, 30)
        panel_title = QLabel('Options')
        panel_title.setStyleSheet("""
			font-weight: bold;
			padding: 0px;
			""")
        main_layout.addWidget(panel_title, alignment=Qt.AlignCenter)

        labels = [QLabel('population: '), QLabel('class: '), QLabel('group: ')]  # , QLabel('time of\ninterest: ')]
        self.cb_options = [['targets', 'effectors'], ['class'], ['group']]  # , ['t0','first detection']]
        self.cbs = [QComboBox() for i in range(len(labels))]
        self.cbs[0].currentIndexChanged.connect(self.set_classes_and_times)

        choice_layout = QVBoxLayout()
        choice_layout.setContentsMargins(20, 20, 20, 20)
        for i in range(len(labels)):
            hbox = QHBoxLayout()
            hbox.addWidget(labels[i], 33)
            hbox.addWidget(self.cbs[i], 66)
            self.cbs[i].addItems(self.cb_options[i])
            if i == 1:
                vbox = QVBoxLayout()
                self.check_class = QCheckBox("plot by class")
                vbox.addWidget(self.check_class)
                vbox.addLayout(hbox)
                choice_layout.addLayout(vbox)
            if i == 2:
                vbox = QVBoxLayout()
                self.check_group = QCheckBox("plot by group")
                vbox.addWidget(self.check_group)
                vbox.addLayout(hbox)
                choice_layout.addLayout(vbox)
            else:
                choice_layout.addLayout(hbox)

        self.cbs[0].setCurrentIndex(1)
        self.cbs[0].setCurrentIndex(0)
        self.cbs[1].setEnabled(False)
        self.check_class.toggled.connect(self.class_enable)
        self.cbs[2].setEnabled(False)
        self.check_group.toggled.connect(self.group_enable)

        # self.abs_time_checkbox = QCheckBox('absolute time')
        # self.frame_slider = QLabeledSlider()
        # self.frame_slider.setSingleStep(1)
        # self.frame_slider.setOrientation(1)
        # self.frame_slider.setRange(0,self.parent.parent.len_movie)
        # self.frame_slider.setValue(0)
        # self.frame_slider.setEnabled(False)
        # slider_hbox = QHBoxLayout()
        # slider_hbox.addWidget(self.abs_time_checkbox, 33)
        # slider_hbox.addWidget(self.frame_slider, 66)
        # choice_layout.addLayout(slider_hbox)
        main_layout.addLayout(choice_layout)

        # self.abs_time_checkbox.stateChanged.connect(self.switch_ref_time_mode)

        # time_calib_layout = QHBoxLayout()
        # time_calib_layout.setContentsMargins(20,20,20,20)
        # time_calib_layout.addWidget(QLabel('time calibration\n(frame to min)'), 33)
        # self.time_calibration_le = QLineEdit(str(self.FrameToMin).replace('.',','))
        # self.time_calibration_le.setValidator(self.float_validator)
        # time_calib_layout.addWidget(self.time_calibration_le, 66)
        # #time_calib_layout.addWidget(QLabel(' min'))
        # main_layout.addLayout(time_calib_layout)

        self.submit_btn = QPushButton('Submit')
        self.submit_btn.setStyleSheet(self.button_style_sheet)
        self.submit_btn.clicked.connect(self.process_signal)
        main_layout.addWidget(self.submit_btn)

    # self.populate_left_panel()
    # grid.addLayout(self.left_side, 0, 0, 1, 1)

    # self.setCentralWidget(self.scroll_area)
    # self.show()
    def class_enable(self):
        if self.check_class.isChecked():
            self.cbs[1].setEnabled(True)
            self.check_group.setChecked(False)
        else:
            self.cbs[1].setEnabled(False)

    def group_enable(self):
        if self.check_group.isChecked():
            self.cbs[2].setEnabled(True)
            self.check_class.setChecked(False)
        else:
            self.cbs[2].setEnabled(False)

    def set_classes_and_times(self):
        ext = 'csv'
        # Look for all classes and times
        tables = glob(self.exp_dir + os.sep.join(['W*', '*', 'output', 'tables', f'trajectories_*{ext}']))
        self.all_columns = []
        for tab in tables:
            cols = pd.read_csv(tab, nrows=1).columns.tolist()
            self.all_columns.extend(cols)
        self.all_columns = np.unique(self.all_columns)
        class_idx = np.array([s.startswith('class_') for s in self.all_columns])
        group_idx = np.array([s.startswith('group_') for s in self.all_columns])

        # time_idx = np.array([s.startswith('t_') for s in self.all_columns])

        try:
            class_columns = list(self.all_columns[class_idx])
            group_columns = list(self.all_columns[group_idx])
        # time_columns = list(self.all_columns[time_idx])
        except:
            print('columns not found')
            self.auto_close = True
            return None

        # self.cbs[2].clear()
        # self.cbs[2].addItems(np.unique(self.cb_options[2]+time_columns))

        self.cbs[1].clear()
        self.cbs[1].addItems(np.unique(self.cb_options[1] + class_columns))
        self.cbs[2].clear()
        self.cbs[2].addItems(np.unique(group_columns))

    def ask_for_feature(self):

        cols = np.array(list(self.df.columns))
        is_number = np.vectorize(lambda x: np.issubdtype(x, np.number))
        feats = cols[is_number(self.df.dtypes)]

        self.feature_choice_widget = QWidget()
        self.feature_choice_widget.setWindowTitle("Select numeric feature")
        layout = QVBoxLayout()
        self.feature_choice_widget.setLayout(layout)
        self.feature_cb = QComboBox()
        self.feature_cb.addItems(feats)
        hbox = QHBoxLayout()
        hbox.addWidget(QLabel('feature: '), 33)
        hbox.addWidget(self.feature_cb, 66)
        layout.addLayout(hbox)

        self.set_feature_btn = QPushButton('set')
        self.set_feature_btn.clicked.connect(self.compute_signals)
        layout.addWidget(self.set_feature_btn)
        self.feature_choice_widget.show()
        center_window(self.feature_choice_widget)

    def ask_for_features(self):

        cols = np.array(list(self.df.columns))
        is_number = np.vectorize(lambda x: np.issubdtype(x, np.number))
        feats = cols[is_number(self.df.dtypes)]

        self.feature_choice_widget = QWidget()
        self.feature_choice_widget.setWindowTitle("Select numeric feature")
        layout = QVBoxLayout()
        self.feature_choice_widget.setLayout(layout)
        self.feature_cb = QComboBox()
        self.feature_cb.addItems(feats)
        hbox = QHBoxLayout()
        hbox.addWidget(QLabel('feature: '), 33)
        hbox.addWidget(self.feature_cb, 66)
        # hbox.addWidget((QLabel('Plot two features')))
        layout.addLayout(hbox)
        self.checkBox_feature = QCheckBox(self.feature_choice_widget)
        self.checkBox_feature.setGeometry(QRect(170, 120, 81, 20))
        # self.checkBox_feature.setText('Plot two features')
        # #hbox.addWidget(self.checkBox_feature)
        # layout.addWidget(self.checkBox_feature, alignment=Qt.AlignCenter)
        # self.checkBox_feature.stateChanged.connect(self.enable_second_feature)

        # layout = QVBoxLayout()
        # self.feature_two_choice_widget.setLayout(layout)
        # self.feature_two_cb = QComboBox()
        # self.feature_two_cb.addItems(feats)
        # self.feature_two_cb.setEnabled(False)
        # hbox_two = QHBoxLayout()
        # hbox_two.addWidget(QLabel('feature two: '), 33)
        # hbox_two.addWidget(self.feature_two_cb, 66)
        # layout.addLayout(hbox_two)

        self.set_feature_btn = QPushButton('set')
        self.set_feature_btn.clicked.connect(self.compute_signals)
        layout.addWidget(self.set_feature_btn)
        self.feature_choice_widget.show()
        center_window(self.feature_choice_widget)

    def enable_second_feature(self):
        if self.checkBox_feature.isChecked():
            self.feature_two_cb.setEnabled(True)
        else:
            self.feature_two_cb.setEnabled(False)

    def compute_signals(self):

        if self.df is not None:
            self.feature_selected = self.feature_cb.currentText()
            self.feature_choice_widget.close()
            self.compute_signal_functions()
            # prepare survival

            # plot survival
            self.survivalWidget = QWidget()
            self.scroll = QScrollArea()
            self.survivalWidget.setMinimumHeight(int(0.8 * self.screen_height))
            self.survivalWidget.setWindowTitle('signals')
            self.plotvbox = QVBoxLayout(self.survivalWidget)
            self.plotvbox.setContentsMargins(30, 30, 30, 30)
            self.survival_title = QLabel('Signal function')
            self.survival_title.setStyleSheet("""
				font-weight: bold;
				padding: 0px;
				""")
            self.plotvbox.addWidget(self.survival_title, alignment=Qt.AlignCenter)

            plot_buttons_hbox = QHBoxLayout()
            plot_buttons_hbox.addWidget(QLabel(''), 80, alignment=Qt.AlignLeft)

            self.legend_btn = QPushButton('')
            self.legend_btn.setIcon(icon(MDI6.text_box, color="black"))
            self.legend_btn.setStyleSheet(self.button_select_all)
            self.legend_btn.setToolTip('Show or hide the legend')
            self.legend_visible = True
            self.legend_btn.clicked.connect(self.show_hide_legend)
            plot_buttons_hbox.addWidget(self.legend_btn, 5, alignment=Qt.AlignRight)

            self.log_btn = QPushButton('')
            self.log_btn.setIcon(icon(MDI6.math_log, color="black"))
            self.log_btn.setStyleSheet(self.button_select_all)
            self.log_btn.clicked.connect(self.switch_to_log)
            self.log_btn.setToolTip('Enable or disable log scale')
            plot_buttons_hbox.addWidget(self.log_btn, 5, alignment=Qt.AlignRight)

            # self.ci_btn = QPushButton('')
            # self.ci_btn.setIcon(icon(MDI6.arrow_expand_horizontal,color="blue"))
            # self.ci_btn.setStyleSheet(self.parent.parent.parent.button_select_all)
            # self.ci_btn.clicked.connect(self.switch_ci)
            # self.ci_btn.setToolTip('Show or hide confidence intervals.')
            # plot_buttons_hbox.addWidget(self.ci_btn, 5, alignment=Qt.AlignRight)

            self.cell_lines_btn = QPushButton('')
            self.cell_lines_btn.setIcon(icon(MDI6.view_headline, color="black"))
            self.cell_lines_btn.setStyleSheet(self.button_select_all)
            self.cell_lines_btn.clicked.connect(self.switch_cell_lines)
            self.cell_lines_btn.setToolTip('Show or hide individual cell signals.')
            plot_buttons_hbox.addWidget(self.cell_lines_btn, 5, alignment=Qt.AlignRight)

            self.fig, self.ax = plt.subplots(1, 1, figsize=(10, 10))  # ,figsize=(10,10)
            # self.setMinimumHeight(100)
            self.survival_window = FigureCanvas(self.fig, title="Survival")
            self.survival_window.setContentsMargins(0, 0, 0, 0)
            if self.df is not None:
                self.initialize_axis()
            plt.tight_layout()

            self.fig.set_facecolor('none')  # or 'None'
            self.fig.canvas.setStyleSheet("background-color: transparent;")
            self.survival_window.canvas.draw()

            # self.survival_window.layout.addWidget(QLabel('WHAAAAATTT???'))

            self.plot_options = [QRadioButton() for i in range(3)]
            self.radio_labels = ['well', 'pos', 'both']
            radio_hbox = QHBoxLayout()
            radio_hbox.setContentsMargins(30,30,30,30)
            self.plot_btn_group = QButtonGroup()
            for i in range(3):
            	self.plot_options[i].setText(self.radio_labels[i])
            	#self.plot_options[i].toggled.connect(self.plot_survivals)
            	self.plot_btn_group.addButton(self.plot_options[i])
            	radio_hbox.addWidget(self.plot_options[i], 33, alignment=Qt.AlignCenter)
            self.plot_btn_group.buttonClicked[int].connect(self.plot_survivals)

            if self.position_indices is not None:
                if len(self.well_indices) > 1 and len(self.position_indices) == 1:
                    self.plot_btn_group.buttons()[0].click()
                    for i in [1, 2]:
                        self.plot_options[i].setEnabled(False)
                elif len(self.well_indices) > 1:
                    self.plot_btn_group.buttons()[0].click()
                elif len(self.well_indices) == 1 and len(self.position_indices) == 1:
                    self.plot_btn_group.buttons()[1].click()
                    for i in [0, 2]:
                        self.plot_options[i].setEnabled(False)
            # else:
            # 	if len(self.well_indices)>1:
            # 		self.plot_btn_group.buttons()[0].click()
            # 	elif len(self.well_indices)==1:
            # 		self.plot_btn_group.buttons()[2].click()

            # elif len(self.well_indices)>1:
            # 	self.plot_btn_group.buttons()[0].click()
            # else:
            # 	self.plot_btn_group.buttons()[1].click()

            if self.position_indices is not None:
                for i in [0, 2]:
                    self.plot_options[i].setEnabled(False)

            # self.plot_options[0].setChecked(True)
            # self.plotvbox.addLayout(radio_hbox)
            self.plotvbox.addLayout(plot_buttons_hbox)
            self.plotvbox.addWidget(self.survival_window)
            # self.class_selection_lbl = QLabel('Select class')
            # self.class_selection_lbl.setStyleSheet("""
            # 	font-weight: bold;
            # 	padding: 0px;
            # 	""")
            # self.plotvbox.addWidget(self.class_selection_lbl, alignment=Qt.AlignCenter)
            # class_selection_hbox = QHBoxLayout()
            # class_selection_hbox.setContentsMargins(30,30,30,30)
            # self.all_btn = QRadioButton('*')
            # self.all_btn.setChecked(True)
            # self.event_btn = QRadioButton('event')
            # self.no_event_btn = QRadioButton('no event')
            # self.class_btn_group = QButtonGroup()
            # for btn in [self.all_btn, self.event_btn, self.no_event_btn]:
            # 	self.class_btn_group.addButton(btn)
            #
            # self.class_btn_group.buttonClicked[int].connect(self.set_class_to_plot)
            #
            # class_selection_hbox.addWidget(self.all_btn, 33, alignment=Qt.AlignLeft)
            # class_selection_hbox.addWidget(self.event_btn, 33, alignment=Qt.AlignCenter)
            # class_selection_hbox.addWidget(self.no_event_btn, 33, alignment=Qt.AlignRight)
            # self.plotvbox.addLayout(class_selection_hbox)

            self.select_pos_label = QLabel('Select positions')
            self.select_pos_label.setStyleSheet("""
				font-weight: bold;
				padding: 0px;
				""")
            self.plotvbox.addWidget(self.select_pos_label, alignment=Qt.AlignCenter)
            #
            # self.select_option = [QRadioButton() for i in range(2)]
            # self.select_label = ['name', 'spatial']
            # select_hbox = QHBoxLayout()
            # select_hbox.setContentsMargins(30,30,30,30)
            # self.select_btn_group = QButtonGroup()
            # for i in range(2):
            # 	self.select_option[i].setText(self.select_label[i])
            # 	#self.select_option[i].toggled.connect(self.switch_selection_mode)
            # 	self.select_btn_group.addButton(self.select_option[i])
            # 	select_hbox.addWidget(self.select_option[i],33, alignment=Qt.AlignCenter)
            # self.select_btn_group.buttonClicked[int].connect(self.switch_selection_mode)
            # self.plotvbox.addLayout(select_hbox)

            self.look_for_metadata()
            if self.metadata_found:
                self.fig_scatter, self.ax_scatter = plt.subplots(1, 1, figsize=(4, 3))
                self.position_scatter = FigureCanvas(self.fig_scatter)
                self.load_coordinates()
                self.plot_spatial_location()
                # self.plot_positions()
                self.ax_scatter.spines['top'].set_visible(False)
                self.ax_scatter.spines['right'].set_visible(False)
                self.ax_scatter.set_aspect('equal')
                self.ax_scatter.set_xticks([])
                self.ax_scatter.set_yticks([])
                plt.tight_layout()

                self.fig_scatter.set_facecolor('none')  # or 'None'
                self.fig_scatter.canvas.setStyleSheet("background-color: transparent;")
                self.plotvbox.addWidget(self.position_scatter)

            self.generate_pos_selection_widget()

            # if self.df is not None and len(self.ks_estimators_per_position)>0:
            # 	self.plot_survivals()
            # self.select_btn_group.buttons()[0].click()
            self.scroll.setVerticalScrollBarPolicy(Qt.ScrollBarAlwaysOn)
            self.scroll.setHorizontalScrollBarPolicy(Qt.ScrollBarAlwaysOff)
            # self.scroll.setWidgetResizable(True)
            self.scroll.setWidget(self.survivalWidget)

            self.scroll.setMinimumHeight(int(0.8 * self.screen_height))
            self.survivalWidget.setSizePolicy(QSizePolicy.Expanding, QSizePolicy.Expanding)
            self.scroll.show()

    def process_signal(self):

        print('you clicked!!')
        # self.FrameToMin = float(self.time_calibration_le.text().replace(',','.'))
        print(self.FrameToMin, 'set')

        # read instructions from combobox options
        self.load_available_tables()
        if self.df is not None:
            self.ask_for_features()
        else:
            return None

    def generate_pos_selection_widget(self):

        self.well_names = self.df['well_name'].unique()
        self.pos_names = self.df_pos_info[
            'pos_name'].unique()  # pd.DataFrame(self.ks_estimators_per_position)['position_name'].unique()
        print(f'POSITION NAMES: ', self.pos_names)
        self.usable_well_labels = []
        for name in self.well_names:
            for lbl in self.well_labels:
                if name + ':' in lbl:
                    self.usable_well_labels.append(lbl)

        self.line_choice_widget = QWidget()
        self.line_check_vbox = QVBoxLayout()
        self.line_choice_widget.setLayout(self.line_check_vbox)
        if len(self.well_indices) > 1:
            self.well_display_options = [QCheckBox(self.usable_well_labels[i]) for i in
                                         range(len(self.usable_well_labels))]
            for i in range(len(self.well_names)):
                self.line_check_vbox.addWidget(self.well_display_options[i], alignment=Qt.AlignLeft)
                self.well_display_options[i].setChecked(True)
                self.well_display_options[i].toggled.connect(self.select_survival_lines)
        else:
            self.pos_display_options = [QCheckBox(self.pos_names[i]) for i in range(len(self.pos_names))]
            for i in range(len(self.pos_names)):
                self.line_check_vbox.addWidget(self.pos_display_options[i], alignment=Qt.AlignLeft)
                self.pos_display_options[i].setChecked(True)
                self.pos_display_options[i].toggled.connect(self.select_survival_lines)

        self.plotvbox.addWidget(self.line_choice_widget, alignment=Qt.AlignCenter)

    def load_available_tables(self):

        """
        Load the tables of the selected wells/positions from the control Panel for the population of interest

        """

        self.well_option = self.parent_window.parent_window.well_list.currentIndex()
        if self.well_option == len(self.wells):
            wo = '*'
        else:
            wo = self.well_option
        self.position_option = self.parent_window.parent_window.position_list.currentIndex()
        if self.position_option == 0:
            po = '*'
        else:
            po = self.position_option - 1

        self.df, self.df_pos_info = load_experiment_tables(self.exp_dir, well_option=wo, position_option=po,
                                                           population=self.cbs[0].currentText(), return_pos_info=True)

        if self.df is None:

            print('No table could be found...')
            msgBox = QMessageBox()
            msgBox.setIcon(QMessageBox.Warning)
            msgBox.setText("No table could be found to compute survival...")
            msgBox.setWindowTitle("Warning")
            msgBox.setStandardButtons(QMessageBox.Ok)
            returnValue = msgBox.exec()
            if returnValue == QMessageBox.Ok:
                self.close()
                return None
            else:
                self.close()
                return None
        else:
            self.df_well_info = self.df_pos_info.loc[:,
                                ['well_path', 'well_index', 'well_name', 'well_number', 'well_alias']].drop_duplicates()

    def compute_signal_functions(self):
        wells = get_experiment_wells(self.exp_dir)
        antibodies = get_experiment_antibodies(self.exp_dir)
        cell_types = get_experiment_cell_types(self.exp_dir)
        concentrations = get_experiment_concentrations(self.exp_dir)
        well_name = [' '.join([s1, '+', s2, s3, 'pM']) for s1, s2, s3 in zip(cell_types, antibodies, concentrations)]

        data = []
        full_data = []
        print(self.well_option)
        if self.well_option > 1:
            self.plot_mode = 'wells'
            for z, well in enumerate(wells):  # loop over wells
                # print(well)
                if z not in self.well_indices:
                    pass
                else:
                    positions = get_positions_in_well(well)
                    for ind, pos in enumerate(positions):  # loop over positions
                        if self.position_indices is not None:
                            if ind + 1 in self.position_indices:
                                print(f'Processing position {pos}...')
                                tab_tc = pos + os.sep.join(
                                    ['output', 'tables', f'trajectories_{self.cbs[0].currentText()}.csv'])
                                if not os.path.exists(tab_tc):
                                    return None
                                else:
                                    data = pd.read_csv(tab_tc)
                                    data['well_name'] = well_name[z]
                                full_data.append(data)

                        else:
                            print(f'Processing position {pos}...')
                            tab_tc = pos + os.sep.join(
                                ['output', 'tables', f'trajectories_{self.cbs[0].currentText()}.csv'])
                            if not os.path.exists(tab_tc):
                                return None
                            else:
                                data = pd.read_csv(tab_tc)
                                data['well_name'] = well_name[z]
                            full_data.append(data)
        else:
            self.plot_mode = 'positions'
            positions = get_positions_in_well(wells[self.well_indices[0]])
            for ind, pos in enumerate(positions):  # loop over positions
                pos_name = pos.split(os.sep)[-2]
                if self.position_indices is not None:
                    if ind + 1 in self.position_indices:
                        print(f'Processing position {pos}...')
                        tab_tc = pos + os.sep.join(
                            ['output', 'tables', f'trajectories_{self.cbs[0].currentText()}.csv'])
                        if not os.path.exists(tab_tc):
                            return None
                        else:
                            data = pd.read_csv(tab_tc)
                            data['position'] = pos_name
                        full_data.append(data)
                else:
                    print(f'Processing position {pos}...')
                    tab_tc = pos + os.sep.join(['output', 'tables', f'trajectories_{self.cbs[0].currentText()}.csv'])
                    if not os.path.exists(tab_tc):
                        return None
                    else:
                        data = pd.read_csv(tab_tc)
                        data['position'] = pos_name
                    full_data.append(data)
        self.plot_data = pd.concat(full_data, ignore_index=True)

    def generate_synchronized_matrix(self, well_group, feature_selected, cclass, max_time):

        if isinstance(cclass, int):
            cclass = [cclass]

        n_cells = len(well_group.groupby(['position', 'TRACK_ID']))
        depth = int(2 * max_time + 3)
        matrix = np.zeros((n_cells, depth))
        matrix[:, :] = np.nan
        mapping = np.arange(-max_time - 1, max_time + 2)
        cid = 0
        for block, movie_group in well_group.groupby('position'):
            for tid, track_group in movie_group.loc[movie_group[self.cbs[1].currentText()].isin(cclass)].groupby(
                    'TRACK_ID'):
                try:
                    timeline = track_group['FRAME'].to_numpy().astype(int)
                    feature = track_group[feature_selected].to_numpy()
                    if self.checkBox_feature.isChecked():
                        second_feature = track_group[self.second_feature_selected].to_numpy()
                    if self.cbs[2].currentText().startswith('t') and not self.abs_time_checkbox.isChecked():
                        t0 = math.floor(track_group[self.cbs[2].currentText()].to_numpy()[0])
                        timeline -= t0
                    elif self.cbs[2].currentText() == 'first detection' and not self.abs_time_checkbox.isChecked():

                        if 'area' in list(track_group.columns):
                            print('area in list')
                            feat = track_group['area'].values
                        else:
                            feat = feature

                        first_detection = timeline[feat == feat][0]
                        timeline -= first_detection
                        print(first_detection, timeline)

                    elif self.abs_time_checkbox.isChecked():
                        timeline -= int(self.frame_slider.value())

                    loc_t = [np.where(mapping == t)[0][0] for t in timeline]
                    matrix[cid, loc_t] = feature
                    if second_feature:
                        matrix[cid, loc_t + 1] = second_feature
                    print(timeline, loc_t)

                    cid += 1
                except:
                    pass
        return matrix

    def col_mean(self, matrix):

        mean_line = np.zeros(matrix.shape[1])
        mean_line[:] = np.nan
        std_line = np.copy(mean_line)

        for k in range(matrix.shape[1]):
            values = matrix[:, k]
            # values = values[values!=0]
            if len(values[values == values]) > 2:
                mean_line[k] = np.nanmean(values)
                std_line[k] = np.nanstd(values)

        return mean_line, std_line

    def initialize_axis(self):
        plt.rcParams['svg.fonttype'] = 'none'
        plt.rcParams["font.family"] = "sans-serif"
        SMALL_SIZE = 5
        MEDIUM_SIZE = 6
        BIGGER_SIZE = 7

        plt.rc('font', size=SMALL_SIZE)  # controls default text sizes
        plt.rc('axes', titlesize=SMALL_SIZE)  # fontsize of the axes title
        plt.rc('axes', labelsize=MEDIUM_SIZE)  # fontsize of the x and y labels
        plt.rc('xtick', labelsize=SMALL_SIZE)  # fontsize of the tick labels
        plt.rc('ytick', labelsize=SMALL_SIZE)  # fontsize of the tick labels
        plt.rc('legend', fontsize=SMALL_SIZE)  # legend fontsize
        plt.rc('figure', titlesize=BIGGER_SIZE)  # fontsize of the figure title

        fig, ax = plt.subplots(1, 1, figsize=(10, 5))
        if self.well_option > 1:
            if self.check_class.isChecked():
                sns.boxenplot(ax=self.ax, data=self.plot_data, y=self.feature_selected, x='well_name',
                              hue=self.cbs[1].currentText(),
                              dodge=True, palette=self.palette)
                if self.show_cell_lines:
                    sns.stripplot(ax=self.ax, data=self.plot_data, y=self.feature_selected, x='well_name',
                                  hue=self.cbs[1].currentText(),
                                  dodge=True, alpha=0.3, linewidth=0.5, size=3, palette=self.palette)
            elif self.check_group.isChecked():
                sns.boxenplot(ax=self.ax, data=self.plot_data, y=self.feature_selected, x='well_name',
                              hue=self.cbs[2].currentText(),
                              dodge=True, palette=self.hue_colors)
                if self.show_cell_lines:
                    sns.stripplot(ax=self.ax, data=self.plot_data, y=self.feature_selected, x='well_name',
                                  hue=self.cbs[2].currentText(),
                                  dodge=True, alpha=0.3, linewidth=0.5, size=3, palette=self.hue_colors)
            else:
                sns.boxenplot(ax=self.ax, data=self.plot_data, y=self.feature_selected, x='well_name',
                              color=self.palette[0],
                              dodge=True)
                if self.show_cell_lines:
                    sns.stripplot(ax=self.ax, data=self.plot_data, y=self.feature_selected, x='well_name',
                                  color=self.palette[0],
                                  dodge=True, alpha=0.3, linewidth=0.5, size=3)
        else:
            if self.check_class.isChecked():
                sns.boxenplot(ax=self.ax, data=self.plot_data, y=self.feature_selected, x='position',
                              hue=self.cbs[1].currentText(),
                              dodge=True, palette=self.palette)
                if self.show_cell_lines:
                    sns.stripplot(ax=self.ax, data=self.plot_data, y=self.feature_selected, x='position',
                                  hue=self.cbs[1].currentText(),
                                  dodge=True, alpha=0.3, linewidth=0.5, size=3, palette=self.palette)
            elif self.check_group.isChecked():
                sns.boxenplot(ax=self.ax, data=self.plot_data, y=self.feature_selected, x='position',
                              hue=self.cbs[2].currentText(),
                              dodge=True, palette=self.hue_colors)
                if self.show_cell_lines:
                    sns.stripplot(ax=self.ax, data=self.plot_data, y=self.feature_selected, x='position',
                                  hue=self.cbs[2].currentText(),
                                  dodge=True, alpha=0.3, linewidth=0.5, size=3, palette=self.hue_colors)
            else:
                sns.boxenplot(ax=self.ax, data=self.plot_data, y=self.feature_selected, x='position',
                              color=self.palette[0],
                              dodge=True)
                if self.show_cell_lines:
                    sns.stripplot(ax=self.ax, data=self.plot_data, y=self.feature_selected, x='position',
                                  color=self.palette[0],
                                  dodge=True, alpha=0.3, linewidth=0.5, size=3)
        self.ax.spines['top'].set_visible(False)
        self.ax.spines['right'].set_visible(False)
        self.ax.set_ylabel(self.feature_selected)
        # ax.set_ylim(0.95,1.3)
        plt.tight_layout()

    def plot_survivals(self, id):
        self.ax.clear()
        plt.rcParams['svg.fonttype'] = 'none'
        plt.rcParams["font.family"] = "sans-serif"
        SMALL_SIZE = 5
        MEDIUM_SIZE = 6
        BIGGER_SIZE = 7

        plt.rc('font', size=SMALL_SIZE)  # controls default text sizes
        plt.rc('axes', titlesize=SMALL_SIZE)  # fontsize of the axes title
        plt.rc('axes', labelsize=MEDIUM_SIZE)  # fontsize of the x and y labels
        plt.rc('xtick', labelsize=SMALL_SIZE)  # fontsize of the tick labels
        plt.rc('ytick', labelsize=SMALL_SIZE)  # fontsize of the tick labels
        plt.rc('legend', fontsize=SMALL_SIZE)  # legend fontsize
        plt.rc('figure', titlesize=BIGGER_SIZE)  # fontsize of the figure title

        fig, ax = plt.subplots(1, 1, figsize=(10, 5))
        if self.well_option > 1:
            if self.check_class.isChecked():
                sns.boxenplot(ax=self.ax, data=self.plot_data, y=self.feature_selected, x='well_name',
                              hue=self.cbs[1].currentText(),
                              dodge=True, palette=self.palette)
                if self.show_cell_lines:
                    sns.stripplot(ax=self.ax, data=self.plot_data, y=self.feature_selected, x='well_name',
                                  hue=self.cbs[1].currentText(),
                                  dodge=True, alpha=0.3, linewidth=0.5, size=3, palette=self.palette)
            elif self.check_group.isChecked():
                sns.boxenplot(ax=self.ax, data=self.plot_data, y=self.feature_selected, x='well_name',
                              hue=self.cbs[2].currentText(),
                              dodge=True, palette=self.hue_colors)
                if self.show_cell_lines:
                    sns.stripplot(ax=self.ax, data=self.plot_data, y=self.feature_selected, x='well_name',
                                  hue=self.cbs[2].currentText(),
                                  dodge=True, alpha=0.3, linewidth=0.5, size=3, palette=self.hue_colors)
            else:
                sns.boxenplot(ax=self.ax, data=self.plot_data, y=self.feature_selected, x='well_name',
                              color=self.palette[0],
                              dodge=True)
                if self.show_cell_lines:
                    sns.stripplot(ax=self.ax, data=self.plot_data, y=self.feature_selected, x='well_name',
                                  color=self.palette[0],
                                  dodge=True, alpha=0.3, linewidth=0.5, size=3)
        else:
            if self.check_class.isChecked():
                sns.boxenplot(ax=self.ax, data=self.plot_data, y=self.feature_selected, x='position',
                              hue=self.cbs[1].currentText(),
                              dodge=True, palette=self.palette)
                if self.show_cell_lines:
                    sns.stripplot(ax=self.ax, data=self.plot_data, y=self.feature_selected, x='position',
                                  hue=self.cbs[1].currentText(),
                                  dodge=True, alpha=0.3, linewidth=0.5, size=3, palette=self.palette)
            elif self.check_group.isChecked():
                sns.boxenplot(ax=self.ax, data=self.plot_data, y=self.feature_selected, x='position',
                              hue=self.cbs[2].currentText(),
                              dodge=True, palette=self.hue_colors)
                if self.show_cell_lines:
                    sns.stripplot(ax=self.ax, data=self.plot_data, y=self.feature_selected, x='position',
                                  hue=self.cbs[2].currentText(),
                                  dodge=True, alpha=0.3, linewidth=0.5, size=3, palette=self.hue_colors)
            else:
                sns.boxenplot(ax=self.ax, data=self.plot_data, y=self.feature_selected, x='position',
                              color=self.palette[0],
                              dodge=True)
                if self.show_cell_lines:
                    sns.stripplot(ax=self.ax, data=self.plot_data, y=self.feature_selected, x='position',
                                  color=self.palette[0],
                                  dodge=True, alpha=0.3, linewidth=0.5, size=3)
        ax.spines['top'].set_visible(False)
        ax.spines['right'].set_visible(False)
        ax.set_ylabel(self.feature_selected)
        plt.tight_layout()
        self.survival_window.setMinimumHeight(int(0.5 * self.screen_height))
        self.survival_window.setMinimumWidth(int(0.8 * self.survivalWidget.width()))
        self.survival_window.setSizePolicy(QSizePolicy.Expanding, QSizePolicy.Expanding)

        self.survival_window.canvas.draw()

    def plot_line(self, line, color, label, mean_signal, ci_option=True, cell_lines_option=False, alpha_ci=0.5,
                  alpha_cell_lines=0.5, std_signal=None, matrix=None):
        try:
            if 'second' in str(mean_signal):
                self.ax2.plot(line['timeline'] * self.FrameToMin, line[mean_signal], color=color, label=label)
            else:
                self.ax.plot(line['timeline'] * self.FrameToMin, line[mean_signal], color=color, label=label)
            if ci_option and std_signal is not None:
                if 'second' in str(mean_signal):
                    self.ax2.fill_between(line['timeline'] * self.FrameToMin,
                                          [a - b for a, b in zip(line[mean_signal], line[std_signal])],
                                          [a + b for a, b in zip(line[mean_signal], line[std_signal])],
                                          color=color,
                                          alpha=alpha_ci,
                                          )
                else:
                    self.ax.fill_between(line['timeline'] * self.FrameToMin,
                                         [a - b for a, b in zip(line[mean_signal], line[std_signal])],
                                         [a + b for a, b in zip(line[mean_signal], line[std_signal])],
                                         color=color,
                                         alpha=alpha_ci,
                                         )
            if cell_lines_option and matrix is not None:
                print(mean_signal)
                mat = line[matrix]
                if 'second' in str(mean_signal):
                    for i in range(mat.shape[0]):
                        self.ax2.plot(line['timeline'] * self.FrameToMin, mat[i, :], color=color,
                                      alpha=alpha_cell_lines)
                else:
                    for i in range(mat.shape[0]):
                        self.ax.plot(line['timeline'] * self.FrameToMin, mat[i, :], color=color, alpha=alpha_cell_lines)

        except Exception as e:
            print(f'Exception {e}')

    def switch_to_log(self):

        """
        Switch threshold histogram to log scale. Auto adjust.
        """

        if self.ax.get_yscale() == 'linear':
            self.ax.set_yscale('log')
        # self.ax.set_ylim(0.01,1.05)
        else:
            self.ax.set_yscale('linear')
        # self.ax.set_ylim(0.01,1.05)

        # self.ax.autoscale()
        self.survival_window.canvas.draw_idle()

    def show_hide_legend(self):
        if self.legend_visible:
            self.ax.legend().set_visible(False)
            self.legend_visible = False
            self.legend_btn.setIcon(icon(MDI6.text_box_outline, color="black"))
        else:
            self.ax.legend().set_visible(True)
            self.legend_visible = True
            self.legend_btn.setIcon(icon(MDI6.text_box, color="black"))

        self.survival_window.canvas.draw_idle()

    def look_for_metadata(self):

        self.metadata_found = False
        self.metafiles = glob(self.exp_dir + os.sep.join([f'W*', '*', 'movie', '*metadata.txt'])) \
                         + glob(self.exp_dir + os.sep.join([f'W*', '*', '*metadata.txt'])) \
                         + glob(self.exp_dir + os.sep.join([f'W*', '*metadata.txt'])) \
                         + glob(self.exp_dir + '*metadata.txt')
        print(f'Found {len(self.metafiles)} metadata files...')
        if len(self.metafiles) > 0:
            self.metadata_found = True

    def switch_selection_mode(self, id):
        print(f'button {id} was clicked')
        for i in range(2):
            if self.select_option[i].isChecked():
                self.selection_mode = self.select_label[i]
        if self.selection_mode == 'name':
            if len(self.metafiles) > 0:
                self.position_scatter.hide()
            self.line_choice_widget.show()
        else:
            if len(self.metafiles) > 0:
                self.position_scatter.show()
            self.line_choice_widget.hide()

    def load_coordinates(self):

        """
        Read metadata and try to extract position coordinates
        """

        self.no_meta = False
        try:
            with open(self.metafiles[0], 'r') as f:
                data = json.load(f)
                positions = data['Summary']['InitialPositionList']
        except Exception as e:
            print(f'Trouble loading metadata: error {e}...')
            return None

        for k in range(len(positions)):
            pos_label = positions[k]['Label']
            try:
                coords = positions[k]['DeviceCoordinatesUm']['XYStage']
            except:
                try:
                    coords = positions[k]['DeviceCoordinatesUm']['PIXYStage']
                except:
                    self.no_meta = True

            if not self.no_meta:
                files = self.df_pos_info['stack_path'].values
                pos_loc = [pos_label in f for f in files]
                self.df_pos_info.loc[pos_loc, 'x'] = coords[0]
                self.df_pos_info.loc[pos_loc, 'y'] = coords[1]
                self.df_pos_info.loc[pos_loc, 'metadata_tag'] = pos_label

    def update_annot(self, ind):

        pos = self.sc.get_offsets()[ind["ind"][0]]
        self.annot.xy = pos
        text = self.scat_labels[ind["ind"][0]]
        self.annot.set_text(text)
        self.annot.get_bbox_patch().set_facecolor('k')
        self.annot.get_bbox_patch().set_alpha(0.4)

    def hover(self, event):
        vis = self.annot.get_visible()
        if event.inaxes == self.ax_scatter:
            cont, ind = self.sc.contains(event)
            if cont:
                self.update_annot(ind)
                self.annot.set_visible(True)
                self.fig_scatter.canvas.draw_idle()
            else:
                if vis:
                    self.annot.set_visible(False)
                    self.fig_scatter.canvas.draw_idle()

    def unselect_position(self, event):
        print('unselecting position')
        self.survival_window.canvas.clear()
        ind = event.ind  # index of selected position
        well_idx = self.df_pos_info.iloc[ind]['well_index'].values[0]
        selectedPos = self.df_pos_info.iloc[ind]['pos_path'].values[0]
        currentSelState = self.df_pos_info.iloc[ind]['select'].values[0]
        if self.plot_options[0].isChecked() or self.plot_options[2].isChecked():
            self.df_pos_info.loc[self.df_pos_info['well_index'] == well_idx, 'select'] = not currentSelState
            self.df_well_info.loc[self.df_well_info['well_index'] == well_idx, 'select'] = not currentSelState
            if len(self.well_indices) > 1:
                self.well_display_options[well_idx].setChecked(not currentSelState)
            else:
                for p in self.pos_display_options:
                    p.setChecked(not currentSelState)
        else:
            self.df_pos_info.loc[self.df_pos_info['pos_path'] == selectedPos, 'select'] = not currentSelState
            if len(self.well_indices) <= 1:
                self.pos_display_options[ind[0]].setChecked(not currentSelState)

        self.sc.set_color(self.select_color(self.df_pos_info["select"].values))
        self.position_scatter.canvas.draw_idle()
        self.plot_survivals(0)

    def select_survival_lines(self):
        if self.plot_mode == 'wells':
            selected_wells = []
            for i in range(len(self.well_display_options)):
                if self.well_display_options[i].isChecked():
                    selected_wells.append(i)
                self.df_well_info.loc[self.df_well_info['well_index'] == i, 'select'] = self.well_display_options[
                    i].isChecked()
                self.df_pos_info.loc[self.df_pos_info['well_index'] == i, 'select'] = self.well_display_options[
                    i].isChecked()
            if len(selected_wells) == 0:
                print('No wells selected')
                self.ax.clear()
            else:
                self.ax.clear()
                self.well_indices = selected_wells
                self.compute_signal_functions()
                self.plot_survivals(0)
        else:
            selected_pos = []
            for i in range(len(self.pos_display_options)):
                if self.pos_display_options[i].isChecked():
                    selected_pos.append(i + 1)
                self.df_pos_info.loc[self.df_pos_info['pos_index'] == i, 'select'] = self.pos_display_options[
                    i].isChecked()
            if len(selected_pos) == 0:
                self.ax.clear()
            else:
                self.position_indices = selected_pos
                self.ax.clear()
                self.compute_signal_functions()
                self.plot_survivals(0)

        if len(self.metafiles) > 0:
            self.sc.set_color(self.select_color(self.df_pos_info["select"].values))
            self.position_scatter.canvas.draw_idle()

    def select_color(self, selection):
        colors = [tab10(0) if s else tab10(0.1) for s in selection]
        return colors

    def plot_spatial_location(self):

        try:
            self.sc = self.ax_scatter.scatter(self.df_pos_info["x"].values, self.df_pos_info["y"].values, picker=True,
                                              pickradius=1, color=self.select_color(self.df_pos_info["select"].values))
            self.scat_labels = self.df_pos_info['metadata_tag'].values
            self.ax_scatter.invert_xaxis()
            self.annot = self.ax_scatter.annotate("", xy=(0, 0), xytext=(10, 10), textcoords="offset points",
                                                  bbox=dict(boxstyle="round", fc="w"),
                                                  arrowprops=dict(arrowstyle="->"))
            self.annot.set_visible(False)
            self.fig_scatter.canvas.mpl_connect("motion_notify_event", self.hover)
            self.fig_scatter.canvas.mpl_connect("pick_event", self.unselect_position)
        except Exception as e:
            pass

    def switch_ref_time_mode(self):
        if self.abs_time_checkbox.isChecked():
            self.frame_slider.setEnabled(True)
            self.cbs[-1].setEnabled(False)
        else:
            self.frame_slider.setEnabled(False)
            self.cbs[-1].setEnabled(True)

    def switch_ci(self):

        if self.show_ci:
            self.ci_btn.setIcon(icon(MDI6.arrow_expand_horizontal, color="black"))
        else:
            self.ci_btn.setIcon(icon(MDI6.arrow_expand_horizontal, color="blue"))
        self.show_ci = not self.show_ci
        self.plot_survivals(0)

    def switch_cell_lines(self):

        if self.show_cell_lines:
            self.cell_lines_btn.setIcon(icon(MDI6.view_headline, color="black"))
        else:
            self.cell_lines_btn.setIcon(icon(MDI6.view_headline, color="blue"))
        self.show_cell_lines = not self.show_cell_lines
        self.plot_survivals(0)

    def set_class_to_plot(self):

        if self.all_btn.isChecked():
            self.target_class = [0, 1]
        elif self.event_btn.isChecked():
            self.target_class = [0]
        else:
            self.target_class = [1]

        self.plot_survivals(0)
