from typing import Any, List, Union

from qgis.PyQt import QtCore
from qgis.PyQt.QtCore import (
    QItemSelectionModel,
    QModelIndex,
    QObject,
    QSortFilterProxyModel,
    Qt,
    QVariant,
)
from qgis.PyQt.QtGui import QStandardItemModel

from openlog.datamodel.assay.generic_assay import AssayDefinition
from openlog.datamodel.connection.openlog_connection import OpenLogConnection
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.multiplot.multiplot_config import (
    MultiplotConfiguration,
    MultiplotConfigurationList,
)
from openlog.plugins.extensions.crossplots.classes import (
    CrossPlotConfig,
    CrossPlotConfigList,
)
from openlog.plugins.extensions.stacked_charts.classes import (
    StackedPlotConfig,
    StackedPlotConfigList,
)


class AssayVisualizationConfigProxyModel(QSortFilterProxyModel):
    def __init__(self, parent: QObject = None):
        super().__init__(parent)

    def lessThan(self, left: QModelIndex, right: QModelIndex) -> bool:
        if left.parent().isValid():
            return super().lessThan(left, right)
        else:
            return False


class AssayVisualizationConfigTreeModel(QStandardItemModel):

    CONFIG_ROLE = QtCore.Qt.UserRole
    DISPLAY_NAME_ROLE = QtCore.Qt.UserRole + 1
    HOLE_ID_ROLE = QtCore.Qt.UserRole + 2
    itemCheckStateChanged = QtCore.pyqtSignal(QModelIndex)

    def __init__(
        self,
        assay_visualization_config_list: AssayVisualizationConfigList,
        multiplot_config_list: MultiplotConfigurationList,
        crossplot_config_list: CrossPlotConfigList,
        stacked_config_list: StackedPlotConfigList,
        openlog_connection: OpenLogConnection,
        parent: QObject = None,
    ):
        """
        Qt TreeModel implementation for AssayVisualizationConfigList and MultiplotConfigurationList display and edition.
        Configuration is displayed with 4 parent item :
        - collar display : _collar_root_index
        - assay display : _assay_root_index
        - multiplot display : _multiplot_root_index
        - crossplots : _crossplot_root_index

        The configuration are synchronized by setData override:
         - all configuration are searched and we apply new check state for each configuration indexes

        To display collar and assay without associated configuration, use :
        - add_assay / remove_assay
        - add_collar / remove_collar


        Args:
            assay_visualization_config_list: AssayVisualizationConfigList to be displayed
            openlog_connection: OpenLogConnection used to get collar display name
            parent: parent object
        """
        super().__init__(0, 1, parent)
        # store parameters with its config
        self.visibility_params_registry = {}
        # store check state : keys are (assay,column,collar), values are (checkstate, expanded, selected)
        self.state = {}

        self.selected_config = None
        self.assay_visualization_config_list = assay_visualization_config_list
        self.multiplot_config_list = multiplot_config_list
        self.crossplot_config_list = crossplot_config_list
        self.stacked_config_list = stacked_config_list
        self.openlog_connection = openlog_connection
        self.setHorizontalHeaderLabels(["Name"])
        self.refresh()

    @property
    def _collar_root_index(self) -> QModelIndex:
        res = self._get_root_child_from_name("Collars")
        if not res.isValid():
            res = self._insert_root_item("Collars")
        return res

    @property
    def _assay_root_index(self) -> QModelIndex:
        res = self._get_root_child_from_name("Downhole data")
        if not res.isValid():
            res = self._insert_root_item("Downhole data")
        return res

    @property
    def _multiplot_root_index(self) -> QModelIndex:
        res = self._get_root_child_from_name("Multiplots")
        if not res.isValid():
            res = self._insert_root_item("Multiplots")
        return res

    @property
    def _crossplot_root_index(self) -> QModelIndex:
        res = self._get_root_child_from_name("Crossplots")
        if not res.isValid():
            res = self._insert_root_item("Crossplots")
        return res

    @property
    def _stacked_root_index(self) -> QModelIndex:
        res = self._get_root_child_from_name("Stacked plots")
        if not res.isValid():
            res = self._insert_root_item("Stacked plots")
        return res

    def set_view(self, view):

        self.view = view

    def refresh(self):
        """
        Update tree model with current AssayVisualizationConfigList.

        """
        self._store_check_state()
        # clear tree
        self._clear(self._assay_root_index)
        self._clear(self._collar_root_index)
        self._clear(self._multiplot_root_index)
        self._clear(self._crossplot_root_index)
        self._clear(self._stacked_root_index)

        # main tree
        self._add_assay_visualization_config()
        self._restore_check_state()
        # other sections
        self._add_other_configs()
        # connect visibility signals
        self._connect_config_visibility_change()

    def _clear(self, parent: QModelIndex) -> None:
        """
        Remove all childrens of parent.
        """
        for i in reversed(range(self.rowCount(parent=parent))):
            self.removeRow(i, parent)

    def _is_index_expanded(self, index: QModelIndex) -> bool:
        """
        Use proxy model to know if index is expanded in TreeView.
        """
        assay_widget = self.parent()
        proxy_index = assay_widget._proxy_mdl.mapFromSource(index)
        return self.view.isExpanded(proxy_index)

    def _set_index_expanded(self, index: QModelIndex, expanded: bool) -> None:
        """
        Use proxy model to expand index.
        """
        assay_widget = self.parent()
        proxy_index = assay_widget._proxy_mdl.mapFromSource(index)
        self.view.setExpanded(proxy_index, expanded)

    def _is_index_selected(self, index: QModelIndex) -> bool:
        """
        Use proxy model to know if index is selected in TreeView.
        """
        assay_widget = self.parent()
        proxy_index = assay_widget._proxy_mdl.mapFromSource(index)
        return self.view.selectionModel().isSelected(proxy_index)

    def _set_index_selected(self, index: QModelIndex, selected: bool) -> None:
        """
        Use proxy model to select index.
        """
        assay_widget = self.parent()
        proxy_index = assay_widget._proxy_mdl.mapFromSource(index)
        self.view.selectionModel().select(
            proxy_index,
            QItemSelectionModel.SelectionFlag.Select
            if selected
            else QItemSelectionModel.SelectionFlag.Deselect,
        )

    def _get_check_state(self) -> dict:
        """
        Get check state of main tree.
        """
        state = {}
        for i in range(self.rowCount(self._assay_root_index)):
            assay_index = self.index(i, 0, self._assay_root_index)
            config = self.data(assay_index, self.CONFIG_ROLE)
            check_state = self.data(assay_index, Qt.CheckStateRole)
            is_expanded = self._is_index_expanded(assay_index)
            is_selected = self._is_index_selected(assay_index)
            state[(config.assay_display_name, config.hole_id)] = (
                check_state,
                is_expanded,
                is_selected,
            )

            for j in range(self.rowCount(assay_index)):
                column_index = self.index(j, 0, assay_index)
                config = self.data(column_index, self.CONFIG_ROLE)
                check_state = self.data(column_index, Qt.CheckStateRole)
                is_expanded = self._is_index_expanded(column_index)
                is_selected = self._is_index_selected(column_index)
                state[
                    (
                        config.assay_display_name,
                        config.column.display_name,
                        config.hole_id,
                    )
                ] = (check_state, is_expanded, is_selected)
                for k in range(self.rowCount(column_index)):
                    collar_column_index = self.index(k, 0, column_index)
                    col_config = self.data(collar_column_index, self.CONFIG_ROLE)
                    check_state = self.data(collar_column_index, Qt.CheckStateRole)
                    is_expanded = self._is_index_expanded(collar_column_index)
                    is_selected = self._is_index_selected(collar_column_index)
                    state[
                        (
                            col_config.assay_display_name,
                            col_config.column.display_name,
                            col_config.hole_id,
                        )
                    ] = (check_state, is_expanded, is_selected)
        return state

    def _store_check_state(self) -> None:
        """
        Store current check state.
        """
        self.state.clear()
        self.state = self._get_check_state()
        # since tree is cleared after, we should remove partial checks.
        self._update_partial_checks()

    def _restore_check_state(self) -> None:
        """
        Restore check state.
        """

        for i in range(self.rowCount(self._assay_root_index)):
            assay_index = self.index(i, 0, self._assay_root_index)
            config = self.data(assay_index, self.CONFIG_ROLE)
            index_data = self.state.get((config.assay_display_name, config.hole_id))
            if index_data is not None:
                old_state, is_expanded, is_selected = index_data
                super().setData(assay_index, old_state, Qt.CheckStateRole)
                self._propagate_check_state(assay_index, old_state != Qt.Unchecked)
                self._set_index_expanded(assay_index, is_expanded)
                self._set_index_selected(assay_index, is_selected)

            for j in range(self.rowCount(assay_index)):
                column_index = self.index(j, 0, assay_index)
                config = self.data(column_index, self.CONFIG_ROLE)
                index_data = self.state.get(
                    (
                        config.assay_display_name,
                        config.column.display_name,
                        config.hole_id,
                    )
                )
                if index_data is not None:
                    old_state, is_expanded, is_selected = index_data
                    super().setData(column_index, old_state, Qt.CheckStateRole)
                    self._propagate_check_state(column_index, old_state != Qt.Unchecked)
                    self._set_index_expanded(column_index, is_expanded)
                    self._set_index_selected(column_index, is_selected)

                for k in range(self.rowCount(column_index)):
                    collar_column_index = self.index(k, 0, column_index)
                    col_config = self.data(collar_column_index, self.CONFIG_ROLE)
                    index_data = self.state.get(
                        (
                            col_config.assay_display_name,
                            col_config.column.display_name,
                            col_config.hole_id,
                        )
                    )
                    if index_data is not None:
                        old_state, is_expanded, is_selected = index_data
                        super().setData(
                            collar_column_index, old_state, Qt.CheckStateRole
                        )
                        self._propagate_check_state(
                            collar_column_index, old_state != Qt.Unchecked
                        )
                        self._set_index_expanded(collar_column_index, is_expanded)
                        self._set_index_selected(collar_column_index, is_selected)

        self._apply_checkstate_to_mirror(self._assay_root_index)
        self._apply_checkstate_to_configs()

    def _update_partial_checks(self) -> None:
        """
        Scan state and replace partial checks by Checked.
        """
        for k, v in self.state.items():
            check, expanded, selected = v
            if check == Qt.CheckState.PartiallyChecked:
                self.state[k] = (Qt.CheckState.Checked, expanded, selected)

    def setData(
        self, index: QtCore.QModelIndex, value: Any, role: int = Qt.DisplayRole
    ) -> bool:
        """
        Override QStandardItemModel setData for child and parent CheckStateRole synchronization.

        Args:
            index: QModelIndex
            value: new value
            role: Qt role

        Returns: True if data set, False otherwise

        """

        if role == Qt.CheckStateRole:
            # called only when user check inside tree
            newState = value
            oldState = self.data(index, role)
            if newState == oldState:
                return True
            # main tree
            if self._get_root_index(index) in [
                self._assay_root_index,
                self._collar_root_index,
            ]:
                self._update_main_tree_checkstate(index, value)
                res = True
            else:
                self._set_checkstate_to_other_configs(index, value)
                res = True

        else:
            res = super().setData(index, value, role)

        return res

    def _set_checkstate_to_other_configs(
        self, index: QtCore.QModelIndex, value: Any, role: int = Qt.DisplayRole
    ) -> bool:
        """
        Synchronize check state to composition plots.

        Args:
            index: QModelIndex
            value: new value
            role: Qt role

        Returns: True if data set, False otherwise

        """
        done_ = False
        visible = value == Qt.Checked
        config = self.data(index, self.CONFIG_ROLE)
        super().setData(index, value, Qt.CheckStateRole)

        # Update configuration and plot visibility
        if isinstance(config, MultiplotConfiguration):
            config.is_visible = visible
            if config.plot:
                config.plot.setVisible(visible)

            self.itemCheckStateChanged.emit(index)
            done_ = True
        elif isinstance(config, CrossPlotConfig):
            config.is_visible = visible
            if config.plot:
                config.plot.setVisible(visible)
            self.itemCheckStateChanged.emit(index)
            done_ = True
        elif isinstance(config, StackedPlotConfig):
            config.is_visible = visible
            if config.plot:
                config.plot.setVisible(visible)
            self.itemCheckStateChanged.emit(index)
            done_ = True
        return done_

    def data(
        self, index: QtCore.QModelIndex, role: int = QtCore.Qt.DisplayRole
    ) -> QVariant:
        """
        Override QStandardItemModel data() for collar name display instead of collar id

        Args:
            index: QModelIndex
            role: Qt role

        Returns: QVariant

        """
        result = super().data(index, role)
        if (
            self.openlog_connection
            and role == QtCore.Qt.DisplayRole
            and index.column() == 0
            and result
            and super().data(index, self.DISPLAY_NAME_ROLE)
        ):
            result = super().data(index, self.DISPLAY_NAME_ROLE)

        return result

    def get_available_assay(self, only_visible: bool = False) -> List[str]:
        """
        Return assay available in tree model.

        Args:
            only_visible: True to returns only visible assay (default False)

        Returns: list of assay name

        """
        return self._get_available_items(self._assay_root_index, only_visible)

    def get_available_collar(self, only_visible: bool = False) -> List[str]:
        """
        Return collar available in tree model.

        Args:
            only_visible: True to returns only visible collar (default False)

        Returns: list of collar name

        """
        return self._get_available_items(self._collar_root_index, only_visible)

    def is_assay(self, index: QModelIndex) -> bool:
        """
        Check if an index is used to display an assay

        Args:
            index: QModelIndex

        Returns: True if index is used to display an assay, False otherwise

        """
        return index.parent() == self._assay_root_index

    def is_collar(self, index: QModelIndex) -> bool:
        """
        Check if an index is used to display a collar

        Args:
            index: QModelIndex

        Returns: True if index is used to display an assay, False otherwise

        """
        return index.parent() == self._collar_root_index

    def is_assay_column(self, index: QModelIndex) -> bool:
        """
        Check if an index is used to display an assay column (when multiple column within an assay)

        Args:
            index: QModelIndex

        Returns: True if index is used to display an assay column, False otherwise

        """
        return index.parent().parent() == self._assay_root_index

    def is_multiplot(self, index: QModelIndex) -> bool:
        """
        Check if an index is used to display a multiplot

        Args:
            index: QModelIndex

        Returns: True if index is used to display an multiplot, False otherwise

        """
        return index.parent() == self._multiplot_root_index

    def is_crossplot(self, index: QModelIndex) -> bool:
        """
        Check if an index is used to display a crossplot

        Args:
            index: QModelIndex

        Returns: True if index is used to display an crossplot, False otherwise

        """
        return index.parent() == self._crossplot_root_index

    def is_stacked(self, index: QModelIndex) -> bool:
        """
        Check if an index is used to display a stacked plot

        Args:
            index: QModelIndex

        Returns: True if index is used to display an stacked plot, False otherwise

        """
        return index.parent() == self._stacked_root_index

    def is_orthogonally_splittable(self, index: QModelIndex) -> bool:
        """
        Check if an index can be orthogonally projected onto traces.
        Check also if cross symbology is off.

        Args:
            index: QModelIndex

        Returns: True if item is orthogonally splittable, False otherwise
        """
        if self.is_assay(index):
            config = index.data(self.CONFIG_ROLE)
            splittable = True
            for column in config.column_config.values():
                splittable = (
                    splittable
                    and column.is_splittable
                    and not column.is_discrete
                    and column.is_minmax_registered
                )
                if hasattr(column, "color_ramp_handler"):
                    splittable = (
                        splittable
                        and column.color_ramp_handler.symbology_widget.is_original_variable()
                    )

            return splittable
        elif self.is_assay_column(index):
            config = index.data(self.CONFIG_ROLE)
            splittable = (
                config.is_splittable
                and not config.is_discrete
                and config.is_minmax_registered
            )
            if hasattr(config, "color_ramp_handler"):
                splittable = (
                    splittable
                    and config.color_ramp_handler.symbology_widget.is_original_variable()
                )
            return splittable
        else:
            return False

    def is_splittable(self, index: QModelIndex) -> bool:
        """
        Check if an index is splittable for display in canvas.
        Check also if cross symbology is off.

        Args:
            index: QModelIndex

        Returns: True if item is splittable, False otherwise

        """

        if self.is_assay(index):
            config = index.data(self.CONFIG_ROLE)
            splittable = True
            for column in config.column_config.values():
                splittable = splittable and column.is_splittable
                if hasattr(column, "color_ramp_handler"):
                    splittable = (
                        splittable
                        and column.color_ramp_handler.symbology_widget.is_original_variable()
                    )

            return splittable
        elif self.is_assay_column(index):

            config = index.data(self.CONFIG_ROLE)
            splittable = config.is_splittable
            if hasattr(config, "color_ramp_handler"):
                splittable = (
                    splittable
                    and config.color_ramp_handler.symbology_widget.is_original_variable()
                )
            return splittable
        else:
            return False

    def flags(self, index: QtCore.QModelIndex) -> QtCore.Qt.ItemFlags:
        """
        Override QStandardItemModel flags.

        Args:
            index: QModelIndex

        Returns: index flags

        """
        # All item are enabled and selectable
        flags = Qt.ItemFlag.ItemIsEnabled | Qt.ItemFlag.ItemIsSelectable
        # All item should be checkable
        flags = flags | Qt.ItemIsUserCheckable | Qt.ItemIsAutoTristate
        return flags

    def _add_assay_visualization_config(self) -> None:
        """
        Add AssayVisualizationConfig items from AssayVisualizationConfigList.
        Added to collar and DH root indexes.
        """

        # get collars and assays
        collars = self.assay_visualization_config_list.collar_list()
        assays = {}
        for avc in self.assay_visualization_config_list.list:
            assay_name = avc.assay_name
            if assays.get(assay_name) is None:
                assays[assay_name] = avc.get_columns_name()

        # insert indexes
        for assay in assays.keys():
            columns = assays.get(assay)
            is_expanded = len(columns) > 1
            # DH
            ## global assay
            config = self.assay_visualization_config_list.get_assay_hole_configuration(
                assay=assay, hole_id=""
            )
            assay_index = self._insert_index(self._assay_root_index)
            self.setData(assay_index, config.assay_display_name)
            self.setData(assay_index, config, self.CONFIG_ROLE)
            super().setData(assay_index, Qt.Checked, Qt.CheckStateRole)

            for column in columns:
                # global column
                if is_expanded:
                    col_config = config.column_config[column]
                    column_index = self._insert_index(assay_index)
                    self.setData(column_index, col_config.column.display_name)
                    self.setData(column_index, col_config, self.CONFIG_ROLE)
                    super().setData(column_index, Qt.Checked, Qt.CheckStateRole)
                else:
                    column_index = assay_index
                # columns associated with collars
                for collar in collars:
                    collar_config = self.assay_visualization_config_list.get_assay_hole_configuration(
                        assay=assay, hole_id=collar
                    )
                    collar_column_config = collar_config.column_config[column]
                    collar_column_index = self._insert_index(column_index)
                    self.setData(collar_column_index, collar)
                    self.setData(
                        collar_column_index, collar_column_config, self.CONFIG_ROLE
                    )
                    super().setData(collar_column_index, Qt.Checked, Qt.CheckStateRole)

        # Collars
        ## collars
        for collar in collars:
            collar_index = self._insert_index(self._collar_root_index)
            self.setData(collar_index, collar)
            super().setData(collar_index, Qt.Checked, Qt.CheckStateRole)
            for assay in assays.keys():
                columns = assays.get(assay)
                is_expanded = len(columns) > 1
                config = (
                    self.assay_visualization_config_list.get_assay_hole_configuration(
                        assay=assay, hole_id=""
                    )
                )
                if is_expanded:
                    collar_assay_index = self._insert_index(collar_index)
                    self.setData(collar_assay_index, config.assay_display_name)
                    self.setData(collar_assay_index, config, self.CONFIG_ROLE)
                    super().setData(collar_assay_index, Qt.Checked, Qt.CheckStateRole)
                else:
                    collar_assay_index = collar_index

                for column in columns:
                    config = self.assay_visualization_config_list.get_assay_hole_configuration(
                        assay=assay, hole_id=collar
                    )
                    col_config = config.column_config.get(column)
                    collar_column_index = self._insert_index(collar_assay_index)
                    self.setData(
                        collar_column_index,
                        col_config.column.display_name
                        if is_expanded
                        else col_config.assay_display_name,
                    )
                    self.setData(collar_column_index, col_config, self.CONFIG_ROLE)
                    super().setData(collar_column_index, Qt.Checked, Qt.CheckStateRole)

    def _add_other_configs(self) -> None:
        """
        Add crossplots and multiplot plots.
        """
        for config in self.multiplot_config_list.list:
            multiplot_index = self._insert_index(self._multiplot_root_index)
            self.setData(multiplot_index, config.name)
            self.setData(multiplot_index, config, self.CONFIG_ROLE)
            super().setData(
                multiplot_index,
                Qt.Checked if config.is_visible else Qt.Unchecked,
                Qt.CheckStateRole,
            )

        for config in self.crossplot_config_list.list:
            crossplot_index = self._insert_index(self._crossplot_root_index)
            self.setData(crossplot_index, config.name)
            self.setData(crossplot_index, config, self.CONFIG_ROLE)
            super().setData(
                crossplot_index,
                Qt.Checked if config.is_visible else Qt.Unchecked,
                Qt.CheckStateRole,
            )

        for config in self.stacked_config_list.list:
            stacked_index = self._insert_index(self._stacked_root_index)
            self.setData(stacked_index, config.name)
            self.setData(stacked_index, config, self.CONFIG_ROLE)
            super().setData(
                stacked_index,
                Qt.Checked if config.is_visible else Qt.Unchecked,
                Qt.CheckStateRole,
            )

    def _connect_config_visibility_change(self) -> None:
        """
        Connect slot to config's visibility_param change.
        """

        self.visibility_params_registry.clear()

        for config in self.assay_visualization_config_list.list:
            for col_name, col_config in config.column_config.items():
                self.visibility_params_registry[
                    col_config.visibility_param
                ] = col_config

                col_config.disconnect_slot_to_visibility_param(
                    self._on_config_visibility_change
                )
                col_config.connect_slot_to_visibility_param(
                    self._on_config_visibility_change
                )

    def _update_main_tree_checkstate(
        self, index: QModelIndex, check_state: Qt.CheckState
    ) -> None:
        """
        Main method executed after visibility change in DH and Collar sections.
        """
        # set checkstate to corresponding index
        super().setData(index, check_state, Qt.CheckStateRole)
        # update children and parent
        self._propagate_check_state(index, check_state == Qt.Checked)
        # apply changes to mirror section
        self._apply_checkstate_to_mirror(source=self._get_root_index(index))
        # apply checkstate to configs
        self._apply_checkstate_to_configs()

    def _on_config_visibility_change(self) -> None:
        """
        Update checkstate of main tree according to configs visibility.
        Triggered by change in config.
        """
        # get config
        param = self.sender()
        config = self.visibility_params_registry.get(param)
        indexes = self._get_config_indexes(config)
        # special case in dh section : only one column is displayed
        if len(indexes) == 0:
            index = self._get_first_child_from_name(
                parent=self._assay_root_index, name=config.assay_display_name
            )
        else:
            # keep first index
            index = indexes[0]

        # update
        self._update_main_tree_checkstate(
            index, Qt.Checked if param.value() else Qt.Unchecked
        )

    def _apply_checkstate_to_mirror(self, source: QModelIndex) -> None:
        """ "
        Apply checkstate from source section to target section.
        """

        if source == self._assay_root_index:
            self._apply_dh_checkstate_to_collar_section()
        elif source == self._collar_root_index:
            self._apply_collar_checkstate_to_dh_section()

    def _apply_collar_checkstate_to_dh_section(self) -> None:
        """
        Apply Collar checkstate to Downhole section
        """
        for config in self.assay_visualization_config_list.list:
            for col_config in config.column_config.values():
                if col_config.hole_id == "":
                    continue
                dh_index = self._get_config_indexes_from_parent(
                    parent=self._assay_root_index, config=col_config
                )[0]
                collar_index = self._get_config_indexes_from_parent(
                    parent=self._collar_root_index, config=col_config
                )[0]
                check_state = self.data(collar_index, Qt.CheckStateRole)
                super().setData(dh_index, check_state, Qt.CheckStateRole)
                self._propagate_check_state(dh_index, check_state)

    def _apply_dh_checkstate_to_collar_section(self) -> None:
        """
        Apply Downhole section checkstate to Collar
        """
        for config in self.assay_visualization_config_list.list:
            for col_config in config.column_config.values():
                if col_config.hole_id == "":
                    continue
                dh_index = self._get_config_indexes_from_parent(
                    parent=self._assay_root_index, config=col_config
                )[0]
                collar_index = self._get_config_indexes_from_parent(
                    parent=self._collar_root_index, config=col_config
                )[0]
                check_state = self.data(dh_index, Qt.CheckStateRole)
                super().setData(collar_index, check_state, Qt.CheckStateRole)
                self._propagate_check_state(collar_index, check_state)

    def _get_first_child_from_name(self, parent: QModelIndex, name: str) -> QModelIndex:
        """
        Get first parent child from display name

        Args:
            parent: parent QModelIndex
            name: child name

        Returns: invalid QModelIndex if child not found, child index otherwise

        """
        res = QModelIndex()
        for i in range(0, self.rowCount(parent)):
            index = parent.child(i, 0)
            if name == super().data(index, Qt.DisplayRole):
                res = index
        return res

    def _get_root_index(self, index: QModelIndex) -> QModelIndex:
        """
        Given an index, return root.
        """
        candidate = index
        while True:
            if candidate == self._assay_root_index:
                return self._assay_root_index
            elif candidate == self._collar_root_index:
                return self._collar_root_index
            elif candidate == self._multiplot_root_index:
                return self._multiplot_root_index
            elif candidate == self._crossplot_root_index:
                return self._crossplot_root_index
            elif candidate == self._stacked_root_index:
                return self._stacked_root_index
            else:
                candidate = candidate.parent()

    def _propagate_check_state(self, index: QModelIndex, check_state: bool) -> None:
        """
        Propagate check state to children or parent.
        We use super().setData
        Args:
            - index : changing index
            - check state : new state
        """
        # first, update children
        self.setChildrenChecked(index, check_state)

        # then, parent
        parent = index.parent()
        if super().data(parent, Qt.CheckStateRole) is not None:
            super().setData(parent, self.childrenCheckState(parent), Qt.CheckStateRole)
            # and parent's parent
            if super().data(parent.parent(), Qt.CheckStateRole) is not None:
                super().setData(
                    parent.parent(),
                    self.childrenCheckState(parent.parent()),
                    Qt.CheckStateRole,
                )

    def _apply_checkstate_to_configs(self) -> None:
        """
        Set checkstate value to each config in tree.
        """
        index = None
        # disconnect visibility change slot
        for config in self.assay_visualization_config_list.list:
            indexes = self._get_config_indexes(config)
            for index in indexes:
                check_state = self.data(index, Qt.CheckStateRole)
                visibility = False if check_state == Qt.Unchecked else True
                config.is_visible = visibility

            for col_config in config.column_config.values():
                indexes = self._get_config_indexes(col_config)
                for index in indexes:
                    check_state = self.data(index, Qt.CheckStateRole)
                    visibility = False if check_state == Qt.Unchecked else True
                    col_config.disconnect_slot_to_visibility_param(
                        self._on_config_visibility_change
                    )
                    col_config.visibility_param.setValue(visibility)
                    col_config.connect_slot_to_visibility_param(
                        self._on_config_visibility_change
                    )

        if index is not None:
            # trigger grid refresh
            self.itemCheckStateChanged.emit(index)

    def setChildrenChecked(self, parent: QModelIndex, checked: bool) -> None:
        """
        Update check state of parent child

        Args:
            parent: parent QModelIndex
            checked: (bool) parent is checked
        """
        for i in range(0, self.rowCount(parent)):
            index = self.index(i, 0, parent)
            if index.isValid():
                check_state = Qt.Checked if checked else Qt.Unchecked
                super().setData(index, check_state, Qt.CheckStateRole)
                # propagate to children's children
                self.setChildrenChecked(index, checked)

    def childrenCheckState(self, parent: QModelIndex) -> Qt.CheckState:
        """
        Define parent CheckState from children CheckStateRole

        Args:
            parent: parent QModelIndex

        Returns: Qt.CheckState

        """
        total = self.rowCount(parent)
        nb_checked = 0
        nb_unchecked = 0
        nb_partial = 0

        for i in range(0, self.rowCount(parent)):
            check_state = self.data(parent.child(i, 0), Qt.CheckStateRole)
            if check_state == Qt.Checked:
                nb_checked = nb_checked + 1
            elif check_state == Qt.Unchecked:
                nb_unchecked = nb_unchecked + 1
            elif check_state == Qt.PartiallyChecked:
                nb_partial = nb_partial + 1

        if total == nb_checked:
            res = Qt.Checked
        elif total == nb_unchecked:
            res = Qt.Unchecked
        else:
            res = Qt.PartiallyChecked
        return res

    def _insert_index(self, parent: QModelIndex) -> QModelIndex:
        """
        Insert a row as parent's child and return associated index.
        """
        row = self.rowCount(parent=parent)
        self.insertRow(row, parent)
        new_index = self.index(row, 0, parent)
        self.insertColumn(0, new_index)
        return new_index

    def _get_config_indexes(
        self, config: Union[AssayVisualizationConfig, AssayColumnVisualizationConfig]
    ) -> List[QModelIndex]:
        """
        Get AssayVisualizationConfig indexes

        Args:
            config: AssayVisualizationConfig

        Returns: list of indexes associated with config

        """
        res = self._get_config_indexes_from_parent(self._collar_root_index, config)
        res += self._get_config_indexes_from_parent(self._assay_root_index, config)
        return res

    def replace_config(
        self,
        old_config: AssayColumnVisualizationConfig,
        new_config: AssayColumnVisualizationConfig,
    ) -> None:

        indexes = self._get_config_indexes(config=old_config)
        for index in indexes:
            self.setData(index, new_config, self.CONFIG_ROLE)

    def _get_config_indexes_from_parent(
        self,
        parent: QModelIndex,
        config: Union[AssayVisualizationConfig, AssayColumnVisualizationConfig],
    ) -> List[QModelIndex]:
        """
        Get AssayVisualizationConfig indexes from a parent index

        Args:
            parent: QModelIndex
            config: AssayVisualizationConfig

        Returns: list of indexes associated with config inside this parent

        """
        res = []
        for i in range(0, self.rowCount(parent)):
            index = parent.child(i, 0)
            if config == self.data(index, self.CONFIG_ROLE):
                res.append(index)
            res += self._get_config_indexes_from_parent(index, config)
        return res

    def _insert_root_item(self, name: str) -> QModelIndex:
        """
        Insert an item in root and returns created QModelIndex

        Args:
            name: item display name

        Returns: created item QModelIndex

        """
        # Add row in root
        row = self.rowCount()
        self.insertRow(row)
        index = self.index(row, 0)
        self.setData(index, name)

        # Insert column in index for child display
        self.insertColumn(0, index)
        return index

    def _get_root_child_from_name(self, name: str) -> QModelIndex:
        """
        Get root child from display name

        Args:
            name: child name

        Returns: invalid QModelIndex if child not found, child index otherwise

        """
        res = QModelIndex()
        for i in range(0, self.rowCount()):
            index = self.index(i, 0)
            if name == super().data(index, Qt.DisplayRole):
                res = index
        return res

    def _get_first_child_from_name(self, parent: QModelIndex, name: str) -> QModelIndex:
        """
        Get first parent child from display name

        Args:
            parent: parent QModelIndex
            name: child name

        Returns: invalid QModelIndex if child not found, child index otherwise

        """
        res = QModelIndex()
        for i in range(0, self.rowCount(parent)):
            index = parent.child(i, 0)
            if name == super().data(index, Qt.DisplayRole):
                res = index
        return res

    def _remove_item_by_name(self, parent: QModelIndex, name: str) -> None:
        """
        Remove child from parent by name

        Args:
            parent: parent QModelIndex
            name: child name
        """
        index = self._get_first_child_from_name(parent, name)
        if index.isValid():
            self.removeRow(index.row(), parent)

    def _remove_item_not_available(
        self, parent: QModelIndex, available_names: List[str]
    ) -> None:
        """
        Remove item in parent with display name not in available list

        Args:
            parent: parent QModelIndex
            available_names: available name list
        """
        for i in reversed(range(0, self.rowCount(parent))):
            if super().data(parent.child(i, 0), Qt.DisplayRole) not in available_names:
                self.removeRow(i, parent)

    def _get_available_items(
        self, parent: QModelIndex, only_visible: bool = False
    ) -> List[str]:
        """
        Return list of child names from a parent.

        Args:
            parent: parent QModelIndex
            only_visible: True to returns only checked child (default False)

        Returns: list of child name

        """
        res = []
        for i in range(0, self.rowCount(parent)):
            index = parent.child(i, 0)
            name = super().data(index, Qt.DisplayRole)
            if only_visible:
                check_state = self.data(index, Qt.CheckStateRole)
                if check_state != Qt.Unchecked:
                    res.append(name)
            else:
                res.append(name)
        return res
