import json
import os
from pathlib import Path
from typing import List, Union

import pyqtgraph as pg
from qgis.core import QgsApplication, QgsProject, QgsVectorLayer
from qgis.PyQt import uic
from qgis.PyQt.QtCore import QCoreApplication, Qt
from qgis.PyQt.QtGui import QIcon
from qgis.PyQt.QtWidgets import (
    QAbstractItemView,
    QAction,
    QCheckBox,
    QFileDialog,
    QInputDialog,
    QMenu,
    QMessageBox,
    QProgressDialog,
    QRadioButton,
    QToolBar,
    QWidget,
)

from openlog.__about__ import DIR_PLUGIN_ROOT
from openlog.core.trace_symbology import TraceSymbology
from openlog.datamodel.assay.generic_assay import AssayDefinition, AssayDomainType
from openlog.datamodel.connection.geotic.geotic_connection import GeoticConnection
from openlog.datamodel.connection.openlog_connection import OpenLogConnection
from openlog.datamodel.connection.spatialite.spatialite_connection import (
    SpatialiteConnection,
)
from openlog.datamodel.connection.xplordb.xplordb_connection import XplordbConnection
from openlog.gui.assay_visualization import config_loader
from openlog.gui.assay_visualization.config.assay_column_visualization_config import (
    AssayColumnVisualizationConfig,
)
from openlog.gui.assay_visualization.config.assay_visualization_config import (
    AssayVisualizationConfig,
    AssayVisualizationConfigList,
)
from openlog.gui.assay_visualization.config.assay_visualization_config_tree_model import (
    AssayVisualizationConfigProxyModel,
    AssayVisualizationConfigTreeModel,
)
from openlog.gui.assay_visualization.config_factory import (
    AssayColumnVisualizationConfigFactory,
)
from openlog.gui.assay_visualization.discrete.config import (
    DiscreteAssayColumnVisualizationConfig,
    DiscreteCategoricalAssayColumnVisualizationConfig,
)
from openlog.gui.assay_visualization.extended.config import (
    ExtendedCategoricalAssayColumnVisualizationConfig,
    ExtendedNumericalAssayColumnVisualizationConfig,
)
from openlog.gui.assay_visualization.plot_grid import AssaySort, PlotGrid
from openlog.gui.assay_visualization.plot_item_factory import AssayPlotItemFactory
from openlog.gui.assay_visualization.stacked.stacked_config import (
    StackedConfiguration,
    StackedConfigurationList,
)
from openlog.toolbelt import PlgLogger
from openlog.toolbelt.checkable_combobox_input_dialog import (
    CheckableComboBoxInputDialog,
)
from openlog.toolbelt.preferences import PlgOptionsManager


class AssayWidget(QWidget):
    def __init__(
        self,
        domain: AssayDomainType,
        parent: QWidget = None,
        openlog_connection: OpenLogConnection = None,
    ):
        """
        Widget to display assay from an OpenLogConnection

        Args:
            parent: parent QWidget
            domain: AssayDomainType assay domain type
            openlog_connection: OpenLogConnection used to get assays
        """
        super().__init__(parent)
        self.domain = domain
        self._openlog_connection = openlog_connection
        self.log = PlgLogger().log

        uic.loadUi(os.path.join(os.path.dirname(__file__), "assay_widget.ui"), self)

        self._visualization_config = AssayVisualizationConfigList()
        self._stacked_config = StackedConfigurationList()

        self._settings_default_configs = AssayVisualizationConfigList()
        self._settings_default_config_file = None

        # Load/Save configuration actions
        self.load_conf_action = QAction(
            QIcon(":images/themes/default/mActionFileOpen.svg"),
            self.tr("Load..."),
            self,
        )
        self.load_conf_action.triggered.connect(self._load_configuration)

        self.save_conf_action = QAction(
            QIcon(":images/themes/default/mActionFileSaveAs.svg"),
            self.tr("Save As..."),
            self,
        )
        self.save_conf_action.triggered.connect(self._save_current_configuration)

        self.add_conf_action = QAction(
            QIcon(":images/themes/default/mActionFileSave.svg"),
            self.tr("Merge to default"),
            self,
        )
        self.add_conf_action.triggered.connect(self._add_config_to_default)

        # Add/Remove assay actions
        self.add_assay_action = QAction(
            QIcon(str(DIR_PLUGIN_ROOT / "resources" / "images" / "add_assay.svg")),
            self.tr("Add"),
            self,
        )
        self.add_assay_action.triggered.connect(self._add_new_assay)

        self.remove_assay_action = QAction(
            QIcon(QgsApplication.iconPath("mActionRemove.svg")),
            self.tr("Remove"),
            self,
        )
        self.remove_assay_action.triggered.connect(self._remove_selected_rows)

        self.remove_assay_action.setEnabled(False)

        self.split_traces = QAction(
            QIcon(
                str(DIR_PLUGIN_ROOT / "resources" / "images" / "trace_symbology.svg")
            ),
            self.tr("Display assay on canvas"),
            self,
        )
        self.split_traces.triggered.connect(self._display_splitted)
        self.split_traces.setEnabled(False)

        self.sort_by_collar_button = QRadioButton(self.tr("Sort by collar"))
        self.sort_by_collar_button.setIcon(
            QIcon(str(DIR_PLUGIN_ROOT / "resources" / "images" / "sort_collar.svg"))
        )
        self.sort_by_collar_button.setChecked(True)
        self.sort_by_collar_button.clicked.connect(self._assay_sort_updated)

        self.sort_by_assay_button = QRadioButton(self.tr("Sort by source"))
        self.sort_by_assay_button.setIcon(
            QIcon(str(DIR_PLUGIN_ROOT / "resources" / "images" / "sort_assay.svg"))
        )
        self.sort_by_assay_button.clicked.connect(self._assay_sort_updated)

        self.flat_grid_checkbox = QCheckBox(self.tr("Transpose"))
        self.flat_grid_checkbox.setIcon(
            QIcon(str(DIR_PLUGIN_ROOT / "resources" / "images" / "flat_grid.svg"))
        )
        self.flat_grid_checkbox.clicked.connect(
            lambda checked: self._plot_grid.enable_flat_grid(checked)
        )

        # Plot tree widget
        self.plot_tree_widget.setEditTriggers(QAbstractItemView.NoEditTriggers)
        self.plot_tree_widget.setSortingEnabled(True)
        self.plot_tree_widget.setSelectionMode(QAbstractItemView.ExtendedSelection)

        self._plot_grid = PlotGrid(
            domain,
            self._visualization_config,
            self._stacked_config,
            self._openlog_connection,
            self,
        )

        self.toolbar = QToolBar(self)
        self.toolbar.addAction(self.add_assay_action)
        self.toolbar.addAction(self.remove_assay_action)
        self.toolbar.addAction(self.split_traces)
        self.toolbar.addSeparator()
        self.toolbar.addWidget(self.sort_by_assay_button)
        self.toolbar.addWidget(self.sort_by_collar_button)
        self.toolbar.addWidget(self.flat_grid_checkbox)

        self.toolbar.addSeparator()
        self._plot_grid.add_action_to_toolbar(self.toolbar)

        # Add menu for symbology load/save

        self.toolbar.addSeparator()
        symbology_menu = QMenu(self)
        symbology_menu.setTitle(self.tr("Symbology"))
        symbology_menu.menuAction().setToolTip(self.tr("Symbology"))
        symbology_menu.addAction(self.load_conf_action)
        symbology_menu.addAction(self.save_conf_action)
        symbology_menu.addAction(self.add_conf_action)

        self.toolbar.addAction(symbology_menu.menuAction())

        self.toolbar_hboxlayout.addWidget(self.toolbar)

        self.plot_grid_scrollarea.setWidget(self._plot_grid)

        # Visualization configuration
        self._visualization_config_model = AssayVisualizationConfigTreeModel(
            self._visualization_config,
            self._stacked_config,
            self._openlog_connection,
            self,
        )
        self._proxy_mdl = AssayVisualizationConfigProxyModel(self)
        self._proxy_mdl.setSourceModel(self._visualization_config_model)
        self.plot_tree_widget.setModel(self._proxy_mdl)
        self.plot_tree_widget.selectionModel().selectionChanged.connect(
            self._tree_selection_changed
        )
        self.plot_tree_widget.header().sortIndicatorChanged.connect(
            self._sort_order_updated
        )

        self._check_state_connect = None
        self._connect_check_state_signal()

        self.plot_selection_splitter.setStretchFactor(0, 0)
        self.plot_selection_splitter.setStretchFactor(1, 1)

        self.assay_config_widget.apply_to_all.connect(self._apply_to_all)
        self.assay_config_widget.apply_to_visible.connect(self._apply_to_visible)
        self.assay_config_widget.apply_to_selected.connect(self._apply_to_selected)

        self.stacked_config_creation_widget.add_stacked_config.connect(
            self._add_current_stacked_config
        )

        self._collar_layer_selection_connect = None
        self.set_openlog_connection(self._openlog_connection)

    def _can_split_geometries(self) -> bool:
        """
        Trace geometries could be available in :
        - XplordbConnection
        - SpatialiteConnection

        """
        result = False
        if isinstance(
            self._openlog_connection,
            (XplordbConnection, SpatialiteConnection, GeoticConnection),
        ):
            result = (
                self._openlog_connection.get_layers_iface().is_trace_geometry_exist()
            )

        return result

    def set_openlog_connection(self, openlog_connection: OpenLogConnection):
        """
        Define used OpenLogConnection

        Args:
            openlog_connection: OpenLogConnection
        """
        self._clear_config()
        self._openlog_connection = openlog_connection
        self._visualization_config_model.openlog_connection = openlog_connection
        self.add_assay_action.setEnabled(self._openlog_connection is not None)
        self._plot_grid.set_openlog_connection(openlog_connection)

        self.assay_config_widget.setVisible(False)
        self.stacked_config_creation_widget.setVisible(False)
        self.stacked_config_widget.setVisible(False)

        self.load_conf_action.setEnabled(self._openlog_connection is not None)
        self.save_conf_action.setEnabled(self._openlog_connection is not None)
        self.add_conf_action.setEnabled(self._openlog_connection is not None)

        self._connect_collar_selection_signal()

    def _connect_check_state_signal(self):
        """
        Connect check state signal for automatic limit and sync plot definition

        """
        self._check_state_connect = (
            self._visualization_config_model.itemCheckStateChanged.connect(
                self._check_state_changed
            )
        )

    def _disconnect_check_state_signal(self):
        """
        Disconnect check state signal for automatic limit and sync plot definition

        """
        if self._check_state_connect:
            self._visualization_config_model.itemCheckStateChanged.disconnect(
                self._check_state_connect
            )
            self._check_state_connect = None

    def _get_collar_layer(self) -> Union[QgsVectorLayer, None]:
        """
        Get collar QgsVectorLayer from current connection

        Returns: Union[QgsVectorLayer, None] collar layer

        """
        if self._openlog_connection is not None:
            return self._openlog_connection.get_layers_iface().get_collar_layer()
        else:
            return None

    def _connect_collar_selection_signal(self) -> None:
        """
        Connect collar selection signal for feature synchronization

        """
        # For now collar selection is done by collar feature synchronization
        collar_layer = self._get_collar_layer()
        if collar_layer is not None:
            self._collar_layer_selection_connect = (
                collar_layer.selectionChanged.connect(self._collar_feature_selected)
            )

    def _disconnect_collar_selection_signal(self) -> None:
        """
        Disconnect collar selection signal for feature synchronization

        """
        # For now collar selection is done by collar feature synchronization
        collar_layer = self._get_collar_layer()
        if self._collar_layer_selection_connect and collar_layer is not None:
            collar_layer.selectionChanged.disconnect(
                self._collar_layer_selection_connect
            )
            self._collar_layer_selection_connect = None

    def _clear_config(self) -> None:
        """
        Clear current visualization configuration.
        If plot is associated to configuration, plot is also deleted

        """
        self._visualization_config.clear()
        self._visualization_config_model.refresh()
        for assay in self._visualization_config_model.get_available_assay():
            self._visualization_config_model.remove_assay(assay)
        for collar in self._visualization_config_model.get_available_collar():
            self._visualization_config_model.remove_collar(collar)

    def _add_new_assay(self) -> None:
        """
        Ask user for assay selection

        """
        assay_definitions = (
            self._openlog_connection.get_assay_iface().get_all_available_assay_definitions()
        )

        # Display only not displayed assay
        selected_hole_id = self._get_selected_hole_id()

        assay_list = [
            assay
            for assay in assay_definitions
            if not self._visualization_config.check_if_all_assay_column_are_available(
                assay
            )
            and assay.domain == self.domain
            and self._openlog_connection.get_assay_iface().is_assay_available_for_collars(
                selected_hole_id, assay.variable
            )
        ]

        assay_name_list = [assay.variable for assay in assay_list]
        assay_display_name_list = [assay.display_name for assay in assay_list]

        dialog = CheckableComboBoxInputDialog(
            assay_name_list,
            assay_display_name_list,
            self,
            self.tr("Downhole data selection"),
            self.tr("Select downhole data"),
        )
        dialog.clear_selection()
        if dialog.exec():

            # Disconnect check state signal to avoid multiple automatic limits calculation
            self._disconnect_check_state_signal()

            assay_list = dialog.get_selected_values()
            self._add_selected_assay_and_collar(
                assay_list, selected_hole_id, assay_definitions
            )

            # Restore check state signal connection for automatic limits calculation
            self._connect_check_state_signal()
            self._check_state_changed()

    def _add_selected_assay_and_collar(
        self,
        assay_list: List[str],
        selected_hole_id: List[str],
        assay_definitions: List[AssayDefinition],
    ) -> None:
        """
        Add selected assay and collar to current visualization widget :
        - create default configuration for each assay
        - add configuration to

        Args:
            assay_list: (List[str]) selected assay names
            selected_hole_id: (List[str]) selected hole id
            assay_definitions: (List[AssayDefinition]) available assay definitions
        """

        # Define imported column for each assay

        assay_column_dict = {}
        for assay_name in assay_list:
            assay_def = [
                definition
                for definition in assay_definitions
                if definition.variable == assay_name
            ][0]

            plotable_columns = list(assay_def.get_plottable_columns().keys())

            # Remove already added columns
            added_column = self._visualization_config.get_assay_added_columns(
                assay_name
            )
            columns = [col for col in plotable_columns if col not in added_column]

            if len(plotable_columns) > 1:
                dialog = CheckableComboBoxInputDialog(
                    columns,
                    columns,
                    self,
                    self.tr("Select colum for {0}").format(assay_def.display_name),
                    assay_def.display_name,
                )
                dialog.clear_selection()
                if not dialog.exec():
                    return
                else:
                    assay_column_dict[assay_name] = dialog.get_selected_values()
            else:
                assay_column_dict[assay_name] = columns

        # Add configuration for each selected collar
        hole_len = len(selected_hole_id)
        assay_len = len(assay_list)
        if assay_len:
            if not hole_len:
                hole_len = 1

            progress_dialog = QProgressDialog(
                self.tr("Importing collar data"),
                self.tr("Cancel"),
                0,
                hole_len * assay_len,
                self,
            )
            progress_dialog.setWindowModality(Qt.WindowModal)
            progress_dialog.show()

            i = 0
            for i, assay_name in enumerate(assay_list):
                assay_def = [
                    definition
                    for definition in assay_definitions
                    if definition.variable == assay_name
                ][0]
                display_name = assay_def.display_name
                title = (
                    self.tr("Time domain downhole data : {0}").format(display_name)
                    if self.domain == AssayDomainType.TIME
                    else self.tr("Depth domain downhole data : {0}").format(
                        display_name
                    )
                )
                progress_dialog.setWindowTitle(title)
                progress_dialog.setValue(i * hole_len)
                QCoreApplication.processEvents()

                # Get selected columns
                selected_columns = assay_column_dict[assay_name]

                # Add assay to tree model
                self._visualization_config_model.add_assay(assay_def)

                # Update default configuration
                progress_dialog.setLabelText(self.tr("Creating default configuration"))
                QCoreApplication.processEvents()

                self._update_default_assay_visualization_config(
                    assay_name, selected_columns
                )
                for j, hole_id in enumerate(selected_hole_id):
                    progress_dialog.setValue(i * hole_len + j)
                    progress_dialog.setLabelText(
                        self.tr("Importing {0} downhole data").format(
                            self._openlog_connection.get_read_iface().get_collar_display_name(
                                hole_id
                            )
                        )
                    )
                    self._add_visualization_config(
                        hole_id, assay_name, selected_columns
                    )

                    if progress_dialog.wasCanceled():
                        break
                    QCoreApplication.processEvents()

                if progress_dialog.wasCanceled():
                    break
                QCoreApplication.processEvents()

            progress_dialog.hide()
            progress_dialog.deleteLater()

            # Refresh displayed visualization
            self._visualization_config_model.refresh()

            # Refresh grid to respect order
            self._plot_grid.refresh_grid()

            # Simulate new selection
            self._tree_selection_changed()

    def _read_settings_default_assay_visualization_config(
        self, force_update: bool = False
    ) -> None:
        # Read default config from configuration file
        default_visualization_config_file = (
            PlgOptionsManager.get_plg_settings().default_visualization_config_file
        )

        if (
            default_visualization_config_file
            and Path(default_visualization_config_file).exists()
            and Path(default_visualization_config_file).is_file()
        ):
            # Read file only if updated from previous read file
            if (
                default_visualization_config_file != self._settings_default_config_file
                or force_update
            ):
                self._settings_default_config_file = default_visualization_config_file
                with open(Path(default_visualization_config_file), "r") as f:
                    self._settings_default_configs, _ = config_loader.from_json(
                        json.load(f), self._openlog_connection, self.domain
                    )
        else:
            self._settings_default_configs = AssayVisualizationConfigList()
            self._settings_default_config_file = None

    def _update_default_assay_visualization_config(
        self, assay_name: str, selected_columns: List[str]
    ) -> None:

        # Read default config from settings
        self._read_settings_default_assay_visualization_config()

        try:
            assay_visualization_config = (
                self._settings_default_configs.get_default_configuration(assay_name)
            )
            self._visualization_config.append_to_list(assay_visualization_config)
        except AssayVisualizationConfigList.UnavailableDefaultConfig:
            pass

        # Get assay from connection
        assay = self._openlog_connection.get_assay_iface().get_assay("", assay_name)

        try:
            assay_visualization_config = (
                self._visualization_config.get_default_configuration(assay_name)
            )
        except AssayVisualizationConfigList.UnavailableDefaultConfig:
            # Create assay visualization configuration
            assay_visualization_config = AssayVisualizationConfig(
                hole_id="", assay=assay_name
            )
            assay_visualization_config.is_visible = True
            assay_visualization_config.assay_name = assay_name
            assay_visualization_config.assay_display_name = (
                assay.assay_definition.display_name
            )
            self._visualization_config.append_to_list(assay_visualization_config)

        # Create configuration for each assay column not already available
        for col in selected_columns:
            if col not in assay_visualization_config.column_config:
                config = AssayColumnVisualizationConfigFactory().create_config(
                    assay.assay_definition, col
                )
                config.assay_display_name = assay.assay_definition.display_name
                config.assay = assay
                config.assay_name = assay_name

                # For categorical assay, default configuration doesn't have category
                if isinstance(
                    config, ExtendedCategoricalAssayColumnVisualizationConfig
                ):
                    # special case for Geotic (symbology stored in database)
                    if isinstance(self._openlog_connection, GeoticConnection):
                        config.bar_symbology.init_from_database(
                            session=self._openlog_connection.session, assay=assay
                        )
                    else:
                        config.bar_symbology.init_from_assay(assay)
                elif isinstance(
                    config, DiscreteCategoricalAssayColumnVisualizationConfig
                ):
                    config.point_symbology.init_from_assay(assay)

                assay_visualization_config.add_column_config(config)
        self._visualization_config.update_min_max()

    def _display_splitted(self):
        indexes = self.plot_tree_widget.selectionModel().selectedIndexes()
        is_column = False
        for proxy_index in indexes:
            index = self._proxy_mdl.mapToSource(proxy_index)

            if self._visualization_config_model.is_assay(index):
                config = self._visualization_config_model.data(
                    index, self._visualization_config_model.CONFIG_ROLE
                )
                name = config.assay_name
                column_names = list(config.column_config.keys())
                is_discrete = isinstance(
                    config.column_config[column_names[0]],
                    (
                        DiscreteAssayColumnVisualizationConfig,
                        DiscreteCategoricalAssayColumnVisualizationConfig,
                    ),
                )

            elif self._visualization_config_model.is_assay_column(index):
                config = self._visualization_config_model.data(
                    index, self._visualization_config_model.CONFIG_ROLE
                )
                name = config.assay_name
                column_names = [config.column.name]
                is_discrete = isinstance(
                    config,
                    (
                        DiscreteAssayColumnVisualizationConfig,
                        DiscreteCategoricalAssayColumnVisualizationConfig,
                    ),
                )
                is_column = True
            else:
                continue

            # display layers to QGIS canvas

            for column in column_names:

                column_config = config if is_column else config.column_config[column]

                layer = self._openlog_connection.get_layers_iface().get_splitted_trace_layer(
                    f"{name}_{column}"
                )
                # if doesnt exist because database was created before Openlog update, create it
                if not layer.isValid():
                    assay_column = column_config.column
                    self._openlog_connection.get_assay_iface().split_traces(
                        table_name=name, assay_column=assay_column, discrete=is_discrete
                    )
                    layer = self._openlog_connection.get_layers_iface().get_splitted_trace_layer(
                        f"{name}_{column}"
                    )

                if layer.isValid():
                    assay_def = self._openlog_connection.get_assay_iface().get_assay_table_definition(
                        name
                    )
                    # translate symbology used in logviewer to QGIS layer
                    layer = TraceSymbology(column_config, assay_def).set_symbology(
                        layer
                    )

                    # remove layer if exists
                    for lyr in QgsProject.instance().mapLayers().values():
                        if lyr.name() == layer.name():
                            QgsProject.instance().removeMapLayers([lyr.id()])

                    QgsProject.instance().addMapLayer(layer)

    def _remove_selected_rows(self) -> None:
        """
        Remove selected assay or/and collars of tree model.
        Removed collars are unselected in QGIS canvas

        """

        removed_config = []
        removed_stacked_config = []
        removed_collar_ids = []
        indexes = self.plot_tree_widget.selectionModel().selectedIndexes()
        for proxy_index in indexes:
            index = self._proxy_mdl.mapToSource(proxy_index)
            if self._visualization_config_model.is_assay(index):
                config = self._visualization_config_model.data(
                    index, self._visualization_config_model.CONFIG_ROLE
                )
                removed_config.append(config)

            if self._visualization_config_model.is_stacked(index):
                config = self._visualization_config_model.data(
                    index, self._visualization_config_model.CONFIG_ROLE
                )
                removed_stacked_config.append(config)

            if self._visualization_config_model.is_collar(index):
                collar_name = self._visualization_config_model.data(
                    index, self._visualization_config_model.HOLE_ID_ROLE
                )
                removed_collar_ids.append(collar_name)

        removed_collar_ids = list(set(removed_collar_ids))

        for config in removed_config:
            # Remove from configuration
            self._visualization_config.remove_assay(config.assay_name)
            # Remove from tree model
            self._visualization_config_model.remove_assay(config.assay_name)

        for config in removed_stacked_config:
            # Remove from configuration
            self._stacked_config.remove_config(config)
            # Remove from tree model
            self._visualization_config_model.remove_stacked(config)

        # Refresh displayed visualization
        self._visualization_config_model.refresh()

        self._check_state_changed()

        # unselect collars from layer
        self._openlog_connection.get_layers_iface().unselect_collar_by_id(
            removed_collar_ids
        )

    def _collar_feature_selected(self) -> None:
        """
        Update configuration with selected collar in map

        """
        selected_hole = self._get_selected_hole_id()
        self._remove_unavailable_collar(selected_hole)
        self._add_or_update_selected_collar(selected_hole)

    def _remove_unavailable_collar(self, selected_hole: List[str]) -> None:
        """
        Remove unavailable collar in tree model and visualization configuration

        Args:
            selected_hole: [str] selected collar ids
        """
        unavailable_collar = [
            hole_id
            for hole_id in self._visualization_config_model.get_available_collar()
            if hole_id not in selected_hole
        ]
        for hole_id in unavailable_collar:
            self._visualization_config.remove_collar(hole_id)
            self._visualization_config_model.remove_collar(hole_id)

    def _add_or_update_selected_collar(self, selected_hole: List[str]) -> None:
        """
        Add or update selected collar for visualization config add if needed

        Args:
            selected_hole: [str]  selected collar ids
        """
        hole_len = len(selected_hole)

        # Disconnect check state signal to avoid multiple automatic limits calculation
        self._disconnect_check_state_signal()

        if hole_len:

            progress_dialog = QProgressDialog(
                self.tr("Importing collar assays"), self.tr("Cancel"), 0, hole_len, self
            )
            title = (
                self.tr("Time assays")
                if self.domain == AssayDomainType.TIME
                else self.tr("Depth assays")
            )
            progress_dialog.setWindowTitle(title)
            progress_dialog.setWindowModality(Qt.WindowModal)
            progress_dialog.show()
            i = 0
            for hole_id in selected_hole:
                progress_dialog.setValue(i)
                progress_dialog.setLabelText(
                    self.tr("Importing {0} assays").format(
                        self._openlog_connection.get_read_iface().get_collar_display_name(
                            hole_id
                        )
                    )
                )
                if self._visualization_config.has_collar(hole_id):
                    self._update_collar(hole_id)
                else:
                    self._add_collar(hole_id)
                if progress_dialog.wasCanceled():
                    break
                i = i + 1
                QCoreApplication.processEvents()

            progress_dialog.hide()
            progress_dialog.deleteLater()

        self._visualization_config_model.refresh()

        # Restore check state signal connection for automatic limits calculation
        self._connect_check_state_signal()
        self._check_state_changed()

    def _update_collar(self, hole_id: str) -> None:
        """
        Update collar visualization configuration with added assays

        Args:
            hole_id: (str) collar hole id
        """
        # Add only assay not already available
        added_assay = [
            assay
            for assay in self._visualization_config.assay_list()
            if not self._visualization_config.has_collar_assay(hole_id, assay)
        ]
        for assay in added_assay:
            # Get selected column from config
            columns = self._visualization_config.get_assay_added_columns(assay)
            self._add_visualization_config(hole_id, assay, columns)

    def _add_collar(self, hole_id: str) -> None:
        """
        Add collar visualization configuration

        Args:
            hole_id: (str) collar hole id
        """
        # Add all available assays
        for assay_name in self._visualization_config.assay_list():
            # Get selected column from config
            columns = self._visualization_config.get_assay_added_columns(assay_name)
            self._add_visualization_config(hole_id, assay_name, columns)

    def _add_visualization_config(
        self, hole_id: str, assay_name: str, columns: List[str]
    ) -> None:
        """
        Add a new visualization configuration for collar and assay

        Args:
            hole_id: (str) collar hole id
            assay_name: (str) assay name
            columns: List[str] assay collumns to add
        """

        # Create or update visualization config
        assay_visualization_config = self._create_or_update_visualization_config(
            hole_id, assay_name, columns
        )

        # Create plot from created configuration
        self._create_plot_from_config(assay_visualization_config, columns)

    def _create_or_update_visualization_config(
        self, hole_id: str, assay_name: str, columns: List[str]
    ) -> AssayVisualizationConfig:
        """
        Create or update a configuration for a collar and assay for defined columns

        Args:
            hole_id: (str) collar hole id
            assay_name: (str) assay name
            columns: List[str] assay collumns to add

        Returns: AssayVisualizationConfig created configuration

        """
        # Get assay from connection
        assay = self._openlog_connection.get_assay_iface().get_assay(
            hole_id, assay_name
        )

        default_config = None

        # Get default configuration if available
        try:
            default_config = self._visualization_config.get_default_configuration(
                assay_name
            )
        except AssayVisualizationConfigList.UnavailableDefaultConfig:
            self.log(
                message=self.tr(
                    "No default configuration available for assay {}"
                ).format(assay_name)
            )

        # Create or get assay visualization configuration
        try:
            assay_visualization_config = (
                self._visualization_config.get_assay_hole_configuration(
                    assay_name, hole_id
                )
            )
        except AssayVisualizationConfigList.UnavailableConfig:
            assay_visualization_config = AssayVisualizationConfig(
                hole_id=hole_id, assay=assay_name
            )

            # Init from default config
            if default_config:
                assay_visualization_config.copy_from_config(default_config)

            self._visualization_config.append_to_list(assay_visualization_config)

        assay_visualization_config.hole_display_name = (
            self._openlog_connection.get_read_iface().get_collar_display_name(hole_id)
        )
        assay_visualization_config.assay_display_name = (
            assay.assay_definition.display_name
        )
        # Default visibility
        visible_assays = self._visualization_config.visible_assay_list()
        is_visible = assay_name in visible_assays
        assay_visualization_config.is_visible = is_visible
        # Create configuration for each assay column
        for col in columns:
            config = AssayColumnVisualizationConfigFactory().create_config(
                assay.assay_definition, col
            )
            config.hole_id = assay.hole_id
            config.hole_display_name = (
                self._openlog_connection.get_read_iface().get_collar_display_name(
                    assay.hole_id
                )
            )
            config.assay_name = assay.assay_definition.variable
            config.assay_display_name = assay.assay_definition.display_name

            # Copy column config from default config
            if default_config and col in default_config.column_config:
                config.copy_from_config(default_config.column_config[col])

            assay_visualization_config.add_column_config(config)

        return assay_visualization_config

    def _create_plot_from_config(
        self,
        assay_visualization_config: AssayVisualizationConfig,
        columns: Union[List[str], None] = None,
    ) -> None:
        """
        Create plot from configuration

        Args:
            assay_visualization_config: (AssayVisualizationConfig) configuration
            columns: List[str], Optional assay collumns to add
        """
        hole_id = assay_visualization_config.hole_id
        assay_name = assay_visualization_config.assay_name

        # Get assay from connection
        assay = self._openlog_connection.get_assay_iface().get_assay(
            hole_id, assay_name
        )
        # Add assay to tree model
        self._visualization_config_model.add_assay(assay.assay_definition)

        if not hole_id or not self._openlog_connection.get_assay_iface().is_assay_available_for_collar(
            hole_id, assay_name
        ):
            return None

        # Add collar to tree model
        self._visualization_config_model.add_collar(hole_id)

        assay_visualization_config.hole_display_name = (
            self._openlog_connection.get_read_iface().get_collar_display_name(hole_id)
        )
        assay_visualization_config.assay_display_name = (
            assay.assay_definition.display_name
        )

        # Default visibility
        is_visible = assay_visualization_config.is_visible

        # Create plot widget for assay and collar
        plot_stacked = self._plot_grid.add_plot_widget(assay)
        assay_visualization_config.set_plot_stacked(plot_stacked)
        plot_stacked.setVisible(
            is_visible and assay_visualization_config.stacked_param.value()
        )

        # Create configuration for each assay column
        for col, config in assay_visualization_config.column_config.items():

            # If columns defined, add only columns in list
            if columns and col not in columns:
                continue

            # Create plot widget for assay and collar
            plot = self._plot_grid.add_plot_widget(assay)
            plot.setVisible(
                is_visible and not assay_visualization_config.stacked_param.value()
            )

            config.hole_id = assay.hole_id
            config.hole_display_name = (
                self._openlog_connection.get_read_iface().get_collar_display_name(
                    assay.hole_id
                )
            )
            config.assay_name = assay.assay_definition.variable
            config.assay_display_name = assay.assay_definition.display_name
            config.plot_widget = plot
            config.plot_stacked = plot_stacked
            used_plot = (
                plot_stacked
                if assay_visualization_config.stacked_param.value()
                else plot
            )
            AssayPlotItemFactory().create_plot_items(used_plot, assay, col, config)

            used_plot.getViewBox().autoRange()

            plot.set_assay_column_config(config)

        plot_stacked.set_assay_config(assay_visualization_config)
        self._visualization_config.update_min_max()

    def _check_state_changed(self) -> None:
        """
        Update sync plot and define automatic limits when selection is changed

        """
        self._plot_grid.update_sync_plot()
        self._plot_grid.automatic_limits()
        self._plot_grid.update_empty_cells_visibility()

    def _assay_sort_updated(self) -> None:
        """
        Update plot grid sort

        """
        if self.sort_by_assay_button.isChecked():
            sort = AssaySort.ASSAY_NAME
        elif self.sort_by_collar_button.isChecked():
            sort = AssaySort.HOLE_ID
        self._plot_grid.set_assay_sort(sort)

    def _sort_order_updated(self) -> None:
        """
        Update plot grid row/col sort

        """
        sort_order = self.plot_tree_widget.header().sortIndicatorOrder()
        self._plot_grid.set_sort_order(sort_order)

    def _get_selected_hole_id(self) -> List[str]:
        """
        Get selected collar from OpenLogConnection collar layer

        Returns: [str] list of selected collar hole id

        """
        return (
            self._openlog_connection.get_layers_iface().get_selected_collar_from_layer()
        )

    def _tree_selection_changed(self) -> None:
        """
        Defined remove action enable from current tree model selection

        """

        # enable split button require all selected index are splittable
        splittable = True
        # can remove assay if parent is root, collar if parent is root, stacked
        removable = True
        indexes = self.plot_tree_widget.selectionModel().selectedIndexes()
        for proxy_index in indexes:
            index = self._proxy_mdl.mapToSource(proxy_index)
            if self._visualization_config_model.is_assay(index):
                removable = removable and True
            elif self._visualization_config_model.is_stacked(index):
                removable = removable and True
            elif self._visualization_config_model.is_collar(
                index
            ) and not self._visualization_config_model.is_assay(index.parent()):
                removable = removable and True

            else:
                removable = removable and False

            if self._visualization_config_model.is_splittable(index):
                splittable = splittable and True
            else:
                splittable = splittable and False

        self.remove_assay_action.setEnabled(removable and len(indexes) > 0)

        self.split_traces.setEnabled(
            splittable and self._can_split_geometries() and len(indexes) > 0
        )

        self._update_assay_config_widget_from_selection()
        self._update_stacked_config_widget_from_selection()
        self._update_stacked_config_creation_widget_from_selection()

    def _update_assay_config_widget_from_selection(self) -> None:
        """
        Update current assay config widget from assay selection

        """
        self.assay_config_widget.setVisible(False)
        indexes = self.plot_tree_widget.selectionModel().selectedIndexes()
        if len(indexes) == 1:
            config = indexes[0].data(self._visualization_config_model.CONFIG_ROLE)
            if isinstance(config, AssayVisualizationConfig):
                self.assay_config_widget.set_assay_config(config)
                self.assay_config_widget.setVisible(True)
            elif isinstance(config, AssayColumnVisualizationConfig):
                self.assay_config_widget.set_assay_column_config(config)
                self.assay_config_widget.setVisible(True)

    def _apply_to_all(self) -> None:
        """
        Apply selected assay visualization configuration to all collar and define configuration as default

        """
        selected_config = self.assay_config_widget.get_assay_config()
        if not selected_config:
            return
        for config in self._visualization_config.list:
            # Update configuration if same assay
            if config.assay_name == selected_config.assay_name:
                config.copy_from_config(selected_config)

        for config in self._stacked_config.list:
            config.copy_from_config(selected_config)

    def _apply_to_selected(self) -> None:
        """
        Apply selected assay visualization configuration to selected collar

        """
        selected_config = self.assay_config_widget.get_assay_config()
        if not selected_config:
            return
        for config in self._visualization_config.list:
            # Update configuration if same assay but not if default configuration
            if config.assay_name == selected_config.assay_name and config.hole_id:
                config.copy_from_config(selected_config)

        for config in self._stacked_config.list:
            config.copy_from_config(selected_config)

    def _apply_to_visible(self) -> None:
        """
        Apply selected assay visualization configuration to visible collar

        """
        selected_config = self.assay_config_widget.get_assay_config()
        if not selected_config:
            return
        for config in self._visualization_config.list:
            # Update configuration if same assay but not if default configuration or collar not visible
            if (
                config.assay_name == selected_config.assay_name
                and config.hole_id
                and config.is_visible
            ):
                config.copy_from_config(selected_config)

        for config in self._stacked_config.list:
            config.copy_from_config(selected_config, only_visible=True)

    def _update_stacked_config_widget_from_selection(self) -> None:
        """
        Update current stacked config widget from assay selection

        """
        self.stacked_config_widget.setVisible(False)
        indexes = self.plot_tree_widget.selectionModel().selectedIndexes()
        if len(indexes) == 1:
            config = indexes[0].data(self._visualization_config_model.CONFIG_ROLE)
            if isinstance(config, StackedConfiguration):
                self.stacked_config_widget.set_stacked_config(config)
                self.stacked_config_widget.setVisible(True)

    def _update_stacked_config_creation_widget_from_selection(self):
        """
        Update current stacked config creation widget from assay column selection

        """
        self.stacked_config_creation_widget.setVisible(False)
        stacked_config = self._stacked_config_from_selected_indexes()

        if len(stacked_config.config_list) > 1:
            self.stacked_config_creation_widget.set_stacked_configuration(
                stacked_config
            )
            self.stacked_config_creation_widget.setVisible(True)

    def _stacked_config_from_selected_indexes(self) -> StackedConfiguration:
        """
        Create StackedConfiguration from treeview selected indexes

        Returns: StackedConfiguration

        """
        indexes = self.plot_tree_widget.selectionModel().selectedIndexes()
        stacked_config = StackedConfiguration()
        for proxy_index in indexes:
            index = self._proxy_mdl.mapToSource(proxy_index)
            self._extract_column_config_from_index(index, stacked_config)
        return stacked_config

    def _extract_column_config_from_index(
        self, index, stacked_config: StackedConfiguration
    ) -> None:
        """

        Args:
            index:
            stacked_config: StackedConfiguration
        """
        config = index.data(self._visualization_config_model.CONFIG_ROLE)
        if isinstance(config, AssayVisualizationConfig):
            # Add only configuration with defined hole_id
            if config.hole_id:
                for colum in config.column_config.values():
                    stacked_config.add_column_config(colum)
        elif isinstance(config, AssayColumnVisualizationConfig):
            stacked_config.add_column_config(config)

        # Extract configuration from children
        for i in range(0, self._visualization_config_model.rowCount(index)):
            child = self._visualization_config_model.index(i, 0, index)
            self._extract_column_config_from_index(child, stacked_config)

    def _add_current_stacked_config(self) -> None:
        """
        Add stacked configuration from current config creation widget

        """
        stacked_config = self.stacked_config_creation_widget.get_stacked_configuration()
        ok = True
        while ok:
            name, ok = QInputDialog.getText(
                self, self.tr("New stacked plot"), self.tr("Stacked name")
            )
            if name:
                used_stacked_config_names = [c.name for c in self._stacked_config.list]
                if name not in used_stacked_config_names:
                    stacked_config.name = name
                    self._add_stacked_config_plot(stacked_config)
                    break
                else:
                    QMessageBox.warning(
                        self,
                        self.tr("Invalid stacked name"),
                        self.tr("Stacked name already used."),
                    )

    def _add_stacked_config_plot(self, stacked_config: StackedConfiguration) -> None:
        """
        Add stacked configuration plot and save StackedConfiguration to model

        Args:
            stacked_config: StackedConfiguration
        """
        # Create plot widget
        plot = self._plot_grid.add_stacked_plot_widget()
        stacked_config.set_plot(plot)

        copy_config = []
        # Create configuration for each assay column
        for config in stacked_config.config_list:
            # Get assay from connection
            assay = self._openlog_connection.get_assay_iface().get_assay(
                config.hole_id, config.assay_name
            )
            # Create new configuration
            new_config = AssayColumnVisualizationConfigFactory().create_config(
                assay.assay_definition, config.column.name
            )
            # Copy style
            new_config.copy_from_config(config)
            # Define hole id and assay
            new_config.hole_id = config.hole_id
            new_config.hole_display_name = config.hole_display_name
            new_config.assay_display_name = config.assay_display_name
            new_config.assay_name = config.assay_name
            new_config.plot_stacked = plot
            # Update visibility param to display hole id
            new_config.update_name_for_stacked_config()
            # Add plot items
            AssayPlotItemFactory().create_plot_items(
                plot, assay, new_config.column.name, new_config
            )
            copy_config.append(new_config)

        plot.getViewBox().autoRange()

        # Use created config list and update model
        stacked_config.config_list = copy_config
        plot.set_stacked_config(stacked_config)

        self._stacked_config.list.append(stacked_config)
        # Refresh view
        self._visualization_config_model.refresh()

    def _save_current_configuration(self) -> None:
        filename, _ = QFileDialog.getSaveFileName(
            self,
            self.tr("Select file"),
            "visualization_configuration.json",
            "JSON file(*.json)",
        )
        if filename:
            with open(filename, "w") as f:
                f.write(self._visualization_config.to_json())

    def _add_default_config_file_if_needed(self) -> None:
        default_visualization_config_file = (
            PlgOptionsManager.get_plg_settings().default_visualization_config_file
        )
        if (
            not default_visualization_config_file
            or not Path(default_visualization_config_file).is_file()
        ):

            ret = QMessageBox.question(
                self,
                self.tr("Default configuration file unavailable"),
                self.tr(
                    "Default configuration file unavailable, do you want to select or create one ?"
                ),
            )
            if ret == QMessageBox.Yes:
                filename, _ = QFileDialog.getSaveFileName(
                    self,
                    self.tr("Define default configuration file path"),
                    "visualization_configuration.json",
                    "JSON file(*.json)",
                )
                if filename:
                    PlgOptionsManager.set_value_from_key(
                        "default_visualization_config_file", str(filename)
                    )

    def _add_config_to_default(self) -> None:

        # Ask user for default config file if not available
        self._add_default_config_file_if_needed()

        default_visualization_config_file = (
            PlgOptionsManager.get_plg_settings().default_visualization_config_file
        )
        if default_visualization_config_file:
            # Merge content if file already available
            if Path(default_visualization_config_file).is_file():
                with open(default_visualization_config_file, "r") as f:
                    configs, _ = config_loader.from_json(
                        json.load(f), self._openlog_connection, self.domain
                    )
                    configs.merge(self._visualization_config)
            # Otherwise, use current config
            else:
                configs = self._visualization_config

            with open(default_visualization_config_file, "w") as f:
                f.write(configs.to_json())

            self._read_settings_default_assay_visualization_config(force_update=True)

    def _load_configuration(self) -> None:
        filename, _ = QFileDialog.getOpenFileName(
            self, self.tr("Select file"), "", "JSON file(*.json)"
        )
        if filename:
            with open(filename, "r") as f:
                configs, errors = config_loader.from_json(
                    json.load(f), self._openlog_connection, self.domain
                )
                if len(errors):
                    error_msg = self.tr(
                        "There was some errors during import, check details for more information."
                    )
                    buttons = QMessageBox.StandardButton.NoButton
                    if len(configs.list):
                        error_msg += self.tr(
                            "\n Do you want to import {} valid configuration(s) ?".format(
                                len(configs.list)
                            )
                        )
                        buttons = (
                            QMessageBox.StandardButton.Yes
                            | QMessageBox.StandardButton.No
                        )
                    msg_box = QMessageBox(
                        QMessageBox.Warning,
                        self.tr("Errors in imported file."),
                        error_msg,
                        buttons=buttons,
                    )
                    msg_box.setDetailedText("\n".join(errors))
                    result = msg_box.exec()

                    if len(configs.list) and result != QMessageBox.StandardButton.Yes:
                        return

                self._visualization_config.clear()

                self._disconnect_collar_selection_signal()
                self._openlog_connection.get_layers_iface().select_collar_by_id(
                    configs.collar_list()
                )
                self._connect_collar_selection_signal()

                for config in configs.list:
                    self._visualization_config.append_to_list(config)
                    self._create_plot_from_config(config)

                # Refresh displayed visualization
                self._visualization_config_model.refresh()
                self._plot_grid.refresh_grid()
