import json
from typing import List

import pyqtgraph as pg
import pyqtgraph.parametertree as ptree
import pyqtgraph.parametertree.parameterTypes as pTypes

from openlog.core import pint_utilities
from openlog.datamodel.assay.generic_assay import AssayDefinition, AssayDomainType
from openlog.gui.assay_visualization.config.assay_column_visualization_config import (
    AssayColumnVisualizationConfig,
)
from openlog.gui.assay_visualization.stacked.stacked_utils import (
    config_list_default_unit,
    is_config_list_valid_for_stack,
)
from openlog.toolbelt import PlgTranslator


class AssayVisualizationConfig:
    def __init__(self, hole_id: str, assay: str):
        """
        Store visualization configuration for an assay.
        Can also access plot and plot item if configuration created from visualization widget

        Configuration supported :
        - is_visible (bool) : assay visibility (default : True)

        Args:
            hole_id: (str) collar hole id
            assay: (str) assay name
        """

        # translation
        self.tr = PlgTranslator().tr

        self.hole_id = hole_id
        self.hole_display_name = ""
        self.assay_name = assay
        self.assay_display_name = ""
        self.is_visible = True
        self.plot_stacked = None
        self.column_config = {}

        self.stacked_param = ptree.Parameter.create(
            name=self.tr("Stacked"),
            type="bool",
            value=False,
            default=False,
        )
        # Connection to parameter changes
        self.stacked_param.sigValueChanged.connect(self._stacked_changed)

        self.unit_parameter = pTypes.SimpleParameter(
            name=self.tr("Conversion"),
            type="str",
            value="",
            default="",
        )
        self.unit_parameter.sigValueChanged.connect(self._update_conversion_unit)
        self._conversion_unit = ""

    def from_json(self, data: dict) -> None:
        """
        Define AssayVisualizationConfig from json data (see to_dict)

        Args:
            data: json data

        """
        self.hole_id = data["hole_id"]
        self.hole_display_name = data["hole_display_name"]
        self.assay_name = data["assay_name"]
        self.assay_display_name = data["assay_display_name"]
        self.is_visible = data["is_visible"]

        self.stacked_param.setValue(data["stacked"])
        self._conversion_unit = data["unit_conversion"]
        self.unit_parameter.setValue(data["unit_conversion"])

    def to_dict(self) -> dict:
        """
        Convert configuration to dict (use for json export)

        Returns:
            dict: configuration as dict

        """
        result = {}
        result["hole_id"] = self.hole_id
        result["hole_display_name"] = self.hole_display_name
        result["assay_name"] = self.assay_name
        result["assay_display_name"] = self.assay_display_name
        result["is_visible"] = self.is_visible

        result["configs"] = {}
        for col, conf in self.column_config.items():
            result["configs"][col] = conf.to_dict()

        result["stacked"] = self.stacked_param.value()
        result["unit_conversion"] = self.unit_parameter.value()
        return result

    def get_pyqtgraph_params(self) -> List[ptree.Parameter]:
        """
        Get pyqtgraph param to display in pyqtgraph ParameterTree

        Returns: List[ptree.Parameter] containing all configuration params

        """
        params = []
        valid_stack = is_config_list_valid_for_stack(self.column_config.values())
        if len(self.column_config) > 1 and valid_stack:
            params.append(self.stacked_param)

        if valid_stack and pint_utilities.is_pint_unit(self._conversion_unit):
            params.append(self.unit_parameter)

        return params

    def _stacked_changed(self) -> None:
        if self.plot_stacked:
            self.plot_stacked.setVisible(self.stacked_param.value())
            for col, col_config in self.column_config.items():
                if self.stacked_param.value():
                    col_config.stack_widget(self.plot_stacked)
                else:
                    col_config.unstack_widget(self.plot_stacked)

    def _update_conversion_unit(self) -> None:
        new_conversion_unit = self.unit_parameter.value()

        # No check or update if current conversion unit not defined
        if not self._conversion_unit:
            self._conversion_unit = new_conversion_unit
            return

        # Check if conversion possible, if not restore previous value
        if not pint_utilities.can_convert(new_conversion_unit, self._conversion_unit):
            self.unit_parameter.setValue(self._conversion_unit)
            return
        elif new_conversion_unit != self._conversion_unit:
            # Propagate unit change to assay column
            for config in self.column_config.values():
                config.set_conversion_unit(new_conversion_unit)

            self._conversion_unit = new_conversion_unit

            # Update stack plot widget unit
            if self.plot_stacked:
                self.plot_stacked.update_displayed_unit(self._conversion_unit)

    def set_plot_stacked(self, plot: pg.PlotWidget) -> None:
        """
        Define plot containing current assay

        Args:
            plot: pg.PlotWidget
        """
        self.plot_stacked = plot

    def get_columns_name(self) -> List[str]:
        """
        Get list of available columns names

        Returns: List[str] available columns names

        """
        return list(self.column_config.keys())

    def add_column_config(self, config: AssayColumnVisualizationConfig) -> None:
        """
        Add assay column configuration

        Args:
            config: AssayColumnVisualizationConfig
        """
        self.column_config[config.column.name] = config

        display_unit = False

        if is_config_list_valid_for_stack(self.column_config.values()):
            # Define conversion unit
            default_unit = config_list_default_unit(self.column_config.values())
            display_unit = default_unit and pint_utilities.is_pint_unit(default_unit)
            if display_unit:
                self.unit_parameter.setDefault(default_unit)
                if not self.unit_parameter.value():
                    self._conversion_unit = default_unit
                    self.unit_parameter.setValue(default_unit)
        else:
            # Unstack all configuration
            self.stacked_param.setValue(False)

        for col in self.column_config.values():
            col.enable_unit_param_display(not display_unit)

            # Use current display unit for column
            if display_unit:
                col.unit_parameter.setValue(self.unit_parameter.value())

    def copy_from_config(self, other) -> None:
        """
        Copy configuration from another configuration.
        If a plot item is associated it will be updated

        Args:
            other: configuration to be copy
        """
        if isinstance(other, AssayColumnVisualizationConfig):
            self.copy_column_config(other)
            return

        self.is_visible = other.is_visible
        self.stacked_param.setValue(other.stacked_param.value())
        self.unit_parameter.setValue(other.unit_parameter.value())

        # Don't create column configuration but copy common
        for key, config in self.column_config.items():
            if key in other.column_config:
                config.copy_from_config(other.column_config[key])

    def copy_column_config(self, column_config: AssayColumnVisualizationConfig) -> None:
        """
        Copy column configuration from another configuration.
        If a plot item is associated it will be updated

        Args:
            column_config: column configuration to be copy
        """
        for key, config in self.column_config.items():
            if key == column_config.column.name:
                config.copy_from_config(column_config)


class AssayVisualizationConfigList:
    class UnavailableDefaultConfig(Exception):
        pass

    class UnavailableConfig(Exception):
        pass

    def __init__(self):
        """
        Container of AssayVisualizationConfig.
        List can be accessed with list attributes.

        Define several method for list modification and query

        """
        self.list = []
        # Registry for updated min/max for all collars displayed for each discrete assay columns. {assay : {column : {collar : [min,max]}}}
        self.values_registry = ValuesRegistry(self)

        # Registry for updated min/max for all collars displayed for each discrete assay columns. {assay : {column : {collar : [min,max]}}}
        self.depth_registry = DepthRegistry(self)

    def append_to_list(self, config: AssayVisualizationConfig):
        self.list.append(config)

    def update_min_max(self) -> None:
        """
        Reinitialize min/max registry and scan all config to update min/max range.
        """
        self.values_registry.update_min_max()
        self.depth_registry.update_min_max()
        self._update_parent_config()

    def _update_parent_config(self):
        """
        For global configs (without hole_id), set the list of child configs.
        """

        keys = set()
        # scan of all configs
        for config in self.list:
            assay_name = config.assay_name
            for col_name, _ in config.column_config.items():
                keys.add(f"{assay_name}_{col_name}")

        # key : assay-column, value : list of child configs
        d = {key: [] for key in keys}

        # assign child configs list to global configs
        for config in self.list:
            if config.hole_id != "":
                assay_name = config.assay_name
                for col_name, col_config in config.column_config.items():
                    d[f"{assay_name}_{col_name}"].append(col_config)

        for config in self.list:
            if config.hole_id == "":
                assay_name = config.assay_name
                for col_name, col_config in config.column_config.items():
                    col_config.set_children_configs(d[f"{assay_name}_{col_name}"])

    def _remove_index_from_registries(self, index: int):

        self.values_registry._remove_minmax_from_registry(index)
        self.depth_registry._remove_minmax_from_registry(index)

    def _set_minmax_to_assay_config(self):

        self.values_registry._set_minmax_to_assay_config()
        self.depth_registry._set_minmax_to_assay_config()

    def to_dict(self) -> dict:
        """
        Convert configuration to dict (use for json export)

        Returns:
            dict: configuration as dict

        """
        result = {}
        result["available_assays"] = list(self.assay_list())
        result["available_collars"] = list(self.collar_list())

        def __sort(config: AssayVisualizationConfig):
            return f"{config.assay_name}_{config.hole_id}"

        list_by_assay = sorted(self.list, key=__sort)
        result["configs"] = [conf.to_dict() for conf in list_by_assay]
        return result

    def to_json(self) -> str:
        """
        Write AssayColumnVisualizationConfig as JSON data (see to_dict)

        Returns: JSON data

        """
        return json.dumps(
            self.to_dict(),
            default=lambda o: o.__dict__,
            sort_keys=True,
            indent=4,
        )

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

        """
        to_remove = [i for i, val in enumerate(self.list)]
        self._remove_indexes(to_remove)

    def merge(self, other) -> None:
        """
        Merge other config into current.

        Args:
            other: configuration to be copy
        """
        for config in other.list:
            current_confs = [
                x
                for x in self.list
                if x.assay_name == config.assay_name and config.hole_id == x.hole_id
            ]
            if current_confs:
                current_confs[0].copy_from_config(config)
            else:
                self.append_to_list(config)

    def collar_list(self) -> List[str]:
        """
        Return list of unique available hole id

        Returns: [str] unique available hole id list

        """
        return list(set([x.hole_id for x in self.list if x.hole_id]))

    def assay_list(self) -> List[str]:
        """
        Return list of unique available assay

        Returns: [str] unique available assay list

        """
        return list(set([x.assay_name for x in self.list]))

    def visible_assay_list(self) -> List[str]:
        """
        Return list of unique visible available assay

        Returns: [str] unique visible available assay list

        """
        return list(set([x.assay_name for x in self.list if x.is_visible]))

    def assay_display_name_list(self) -> List[str]:
        """
        Return list of unique available assay display name

        Returns: [str] unique available assay display name list

        """
        return list(set([x.assay_display_name for x in self.list]))

    def get_default_configuration(self, assay: str) -> AssayVisualizationConfig:
        """
        Get default configuration for assay

        Raises UnavailableDefaultConfig if no config available for assay

        Args:
            assay: (str) assay name

        Returns: AssayVisualizationConfig

        """
        default_config = [
            x for x in self.list if x.assay_name == assay and not x.hole_id
        ]
        if default_config:
            return default_config[0]
        else:
            raise AssayVisualizationConfigList.UnavailableDefaultConfig()

    def get_assay_hole_configuration(
        self, assay: str, hole_id: str
    ) -> AssayVisualizationConfig:
        """
        Get configuration for assay

        Raises UnavailableConfig if no config available for assay

        Args:
            assay: (str) assay name
            hole_id: (str) collar hole id

        Returns: AssayVisualizationConfig

        """
        configs = [
            x for x in self.list if x.assay_name == assay and x.hole_id == hole_id
        ]
        if configs:
            return configs[0]
        else:
            raise AssayVisualizationConfigList.UnavailableConfig()

    def has_collar_assay(self, hole_id: str, assay: str) -> bool:
        """
        Check if list contains a configuration for a specific collar and assay

        Args:
            hole_id: (str) collar hole id
            assay: (str) assay name

        Returns: True if at least one assay configuration is available, False otherwise

        """
        return any(x.hole_id == hole_id and x.assay_name == assay for x in self.list)

    def has_collar(self, hole_id: str) -> bool:
        """
        Check if list contains collar

        Args:
            hole_id: (str) collar hole id

        Returns: True if at least one assay configuration is available for collar, False otherwise

        """
        return any(x.hole_id == hole_id for x in self.list)

    def has_assay(self, assay: str) -> bool:
        """
        Check if list contains assay

        Args:
            assay: (str) assay name

        Returns: True if at least one assay configuration is available for assay, False otherwise

        """
        return any(x.assay_name == assay for x in self.list)

    def check_if_all_assay_column_are_available(
        self, assay_definition: AssayDefinition
    ) -> bool:
        """
        Check if list contains all assay plotable columns

        Args:
            assay_definition: (AssayDefinition) assay definition

        Returns: True if all plotable columns of assay are available for assay, False otherwise

        """
        expected_columns = list(assay_definition.get_plottable_columns().keys())
        added_columns = self.get_assay_added_columns(assay_definition.variable)

        # ensure only strings are in lists
        expected_columns = [col for col in expected_columns if isinstance(col, str)]
        added_columns = [col for col in added_columns if isinstance(col, str)]
        return sorted(expected_columns) == sorted(added_columns)

    def get_assay_visualization_config_list(
        self, assay: str
    ) -> List[AssayVisualizationConfig]:
        """
        Get assay visualization config list from assay name

        Args:
            assay: (str) assay name

        Returns: List[AssayVisualizationConfig] assay visualization config list from assay name

        """
        return [x for x in self.list if x.assay_name == assay]

    def get_assay_added_columns(self, assay: str) -> List[str]:
        """
        Get added columns from assay name

        Args:
            assay: (str) assay name

        Returns: List[str] added columns for assay

        """
        added_columns = []
        for config in self.get_assay_visualization_config_list(assay):
            added_columns += config.get_columns_name()

        return list(set(added_columns))

    def remove_collar(self, hole_id: str) -> None:
        """
        Remove all configuration associated to collar.
        If plot is associated to configuration, plot is also deleted

        Args:
            hole_id: (str) collar hole id
        """
        to_remove = [i for i, val in enumerate(self.list) if val.hole_id == hole_id]
        self._remove_indexes(to_remove)

    def remove_assay(self, assay: str) -> None:
        """
        Remove all configuration associated to assay.
        If plot is associated to configuration, plot is also deleted

        Args:
            assay: (str) assay name
        """
        to_remove = [i for i, val in enumerate(self.list) if val.assay_name == assay]
        self._remove_indexes(to_remove)

    def remove_column(self, column: str) -> None:
        """
        Remove all column configuration associated to column.
        If plot is associated to configuration, plot is also deleted

        Args:
            assay: (str) assay name
        """
        for config in self.list:
            if config.column_config.get(column) is None:
                continue
            # remove associated plot widget
            config.column_config[column].clear_item()
            if config.column_config[column].plot_widget:
                config.column_config[column].plot_widget.clear()
                config.column_config[column].plot_widget.setVisible(False)
                config.column_config[column].plot_widget.deleteLater()

            config.column_config.pop(column, None)

    def _remove_indexes(self, to_remove: List[int]) -> None:
        """
        Remove indexes from list

        Args:
            to_remove: indexes to remove
        """
        for index in reversed(to_remove):
            plot = self.list[index].plot_stacked
            if plot is not None:
                plot.clear()
                # Update visibility even if plot is deleted later
                # so we can count number of visible plot if main thread as not process event
                plot.setVisible(False)
                plot.deleteLater()

            for col, col_config in self.list[index].column_config.items():
                col_config.clear_item()
                if col_config.plot_widget:
                    col_config.plot_widget.clear()
                    # Update visibility even if plot is deleted later
                    # so we can count number of visible plot if main thread as not process event
                    col_config.plot_widget.setVisible(False)
                    col_config.plot_widget.deleteLater()

            self._remove_index_from_registries(index)
            self._set_minmax_to_assay_config()
            del self.list[index]


class GenericRegistry:
    def __init__(self, assaylist: AssayVisualizationConfigList):
        """
        Abstract class to store and update range of values displayed in logviewer.
        Object designed to work along with AssayVisualizationConfigList.
        Args :
            - assaylist: AssayVisualizationConfigList
        """
        self.assaylist = assaylist
        self.registry = {}

    def update_min_max(self) -> None:
        """
        Reinitialize min/max registry and scan all config to update min/max range.
        """
        self.minmax_registry = {}
        for config in self.assaylist.list:
            if config.hole_id != "":

                self._add_minmax_to_registry(config)

        self._set_minmax_to_assay_config()

    def _find_column_minmax(
        self, assay_name: str, column_name: str, hole_id: str = ""
    ) -> tuple[float, float]:
        """
        For a given column, return min and max value for all collars displayed in logviewer
        """
        res = (0, 0)
        if self.registry.get(assay_name) is None:
            return res

        if self.registry[assay_name].get(column_name) is None:
            return res

        if hole_id != "":
            if self.registry[assay_name][column_name].get(hole_id) is None:
                return res
            else:
                res = self.registry[assay_name][column_name][hole_id]
                return res
        else:
            mins = maxs = []
            for minmax in self.registry[assay_name][column_name].values():
                mini, maxi = minmax
                if mini is not None:
                    mins.append(mini)
                if maxi is not None:
                    maxs.append(maxi)

            if len(mins) == 0:
                res = (0, 0)
            else:
                res = (min(mins), max(maxs))
            return res

    def _set_minmax_to_assay_config(self):
        return NotImplementedError

    def _find_active_config(self):
        return NotImplementedError

    def _add_minmax_to_registry(self):
        return NotImplementedError

    def _add_collar_to_registry(
        self,
        assay_name: str,
        column_name: str,
        hole_id: str,
        min_value: float,
        max_value: float,
    ) -> None:
        """
        Add a collar min/max range to registry.

        """
        # create assay key
        if assay_name not in self.registry.keys():
            self.registry[assay_name] = {}

        # create column key if don't exist and add range
        if column_name in self.registry[assay_name].keys():
            self.registry[assay_name][column_name][hole_id] = [
                min_value,
                max_value,
            ]
        else:
            self.registry[assay_name][column_name] = {hole_id: [min_value, max_value]}

    def _remove_minmax_from_registry(self, index: int):
        """
        Remove min/max values from registry.
        """
        config = self.assaylist.list[index]
        # detect if name config is associated with a collar
        is_collar = config.hole_id != ""
        if is_collar:
            collar_name = config.hole_id
            for assay_name in self.registry.keys():
                for column_name in self.registry[assay_name].keys():
                    self.registry[assay_name][column_name].pop(collar_name, None)
        else:
            self.registry.pop(config.assay_name, None)


class DepthRegistry(GenericRegistry):
    """
    Registry to store depth values.
    Config's attributes are updated within this class.
    """

    def _add_minmax_to_registry(self, config: AssayVisualizationConfig) -> None:
        """
        Find min/max value and add to registry.
        """

        if config.hole_id != "":
            # update min/max for each discrete assay column
            for col_name, col_config in config.column_config.items():

                if not col_config.assay:
                    continue

                if col_config.assay.use_assay_column_plugin_reader:
                    continue

                if col_config.assay.assay_definition.domain == AssayDomainType.TIME:
                    continue
                # set real min/max range because column config are copied from global config
                x, _ = col_config.assay.get_all_values(col_config.column.name)
                x = x.ravel().astype(float)
                hole_id = col_config.hole_id
                assay_name = col_config.assay_name
                if len(x) == 0:
                    min_value = max_value = None
                else:
                    min_value = x.min()
                    max_value = x.max()
                self._add_collar_to_registry(
                    assay_name, col_name, hole_id, min_value, max_value
                )

    def _set_minmax_to_assay_config(self):
        """
        Refresh min and max depths attributes.
        """
        # we update all assays
        for assay_config in self.assaylist.list:

            for col_name in assay_config.column_config.keys():

                mini, maxi = self._find_column_minmax(
                    assay_config.assay_name, col_name, assay_config.hole_id
                )
                col_config = assay_config.column_config[col_name]
                col_config.min_depth = mini
                col_config.max_depth = maxi


class ValuesRegistry(GenericRegistry):
    """
    Registry to store measure values.
    """

    def _add_minmax_to_registry(self, config: AssayVisualizationConfig) -> None:
        """
        Find min/max value and add to registry.
        """

        if config.hole_id != "":
            # update min/max for each discrete assay column
            for col_name, col_config in config.column_config.items():
                if not col_config.is_minmax_registered:
                    continue
                # scan active configuration
                col_config = self._find_active_config(col_config)
                # set real min/max range because column config are copied from global config
                if col_config.plot_item:
                    col_config.define_color_ramp_min_max_from_plot_item(force=True)

                hole_id = col_config.hole_id
                assay_name = col_config.assay_name
                min_value = col_config.min_color_ramp_param.value()
                max_value = col_config.max_color_ramp_param.value()
                self._add_collar_to_registry(
                    assay_name, col_name, hole_id, min_value, max_value
                )

    def _find_active_config(
        self, col_config: AssayColumnVisualizationConfig
    ) -> AssayColumnVisualizationConfig:
        """
        Active configuration is the one displayed in logviewer.
        If discrete assay is switched to extended, active configuration is col_config.extended_config, else it's col_config.
        If extended assay is switched to discrete, active configuration is col_config.discrete_configuration, else it's col_config.
        """
        if col_config.is_discrete and col_config.is_switchable_config:
            if col_config.as_extended:
                return col_config.extended_config
            else:
                return col_config

        elif col_config.is_switchable_config and not col_config.is_discrete:
            if col_config.as_discrete:
                return col_config.discrete_configuration
            else:
                return col_config
        else:
            return col_config

    def _set_minmax_to_assay_config(self):
        """
        Refresh min and max values of the colar ramp for a whole assay.
        """
        # we update all assays
        for assay_config in self.assaylist.list:

            if assay_config.hole_id != "":
                continue

            for col_name in assay_config.column_config.keys():
                # test discrete
                if not assay_config.column_config[col_name].is_minmax_registered:
                    continue

                mini, maxi = self._find_column_minmax(assay_config.assay_name, col_name)
                col_config = assay_config.column_config[col_name]
                # set to active config and hidden one
                active_config = self._find_active_config(col_config)
                col_config.min_color_ramp_param.setValue(mini)
                col_config.max_color_ramp_param.setValue(maxi)
                active_config.min_color_ramp_param.setValue(mini)
                active_config.max_color_ramp_param.setValue(maxi)
