import copy
import json

import numpy as np
import plotly.figure_factory as ff
from dash import callback_context, dcc, html
from dash.dependencies import ALL, Input, Output

from rubicon_ml.viz.base import VizBase
from rubicon_ml.viz.common import dropdown_header
from rubicon_ml.viz.common.colors import light_blue, plot_background_blue


class MetricListsComparison(VizBase):
    """Visualize lists of metrics logged to the experiments `experiments` as
    an annotated heatmap.

    More info on annotated heatmaps can be found here:
    https://plotly.com/python/annotated-heatmap/

    Parameters
    ----------
    column_names : list of str
        Titles to use for each column in the heatmap. Defaults to None.
    experiments : list of rubicon_ml.client.experiment.Experiment, optional
        The experiments to visualize. Defaults to None. Can be set as
        attribute after instantiation.
    selected_metric : str
        The name of the metric to display at launch. Defaults to None, which
        selects the metric loaded first.
    """

    def __init__(
        self,
        column_names=None,
        experiments=None,
        selected_metric=None,
    ):
        super().__init__(dash_title="compare metric lists")

        if column_names is None:
            self.column_names = []
        else:
            self.column_names = column_names
        self.experiments = experiments
        self.selected_metric = selected_metric

    @property
    def layout(self):
        """Defines the layout for the metric lists comparison."""
        return html.Div(
            [
                dropdown_header(
                    list(self.metric_names),
                    self.selected_metric,
                    "comparing metric ",
                    f" over {len(self.experiments)} experiments",
                    "metric",
                ),
                dcc.Loading(
                    html.Div(
                        dcc.Graph(
                            id="metric-heatmap",
                        ),
                        id="metric-heatmap-container",
                    ),
                    color=light_blue,
                ),
            ],
            id="metric-heatmap-layout-container",
        )

    def load_experiment_data(self):
        """Load the experiment data required for the metric lists comparison.

        Extracts all metric metadata from each experiment in `self.experiment`
        if the metric's value is a list.
        """
        self.experiment_records = {}
        self.metric_names = set()

        for experiment in self.experiments:
            for metric in experiment.metrics():
                if isinstance(metric.value, list):
                    self.metric_names.add(metric.name)

                    experiment_record = self.experiment_records.get(experiment.id, {})
                    experiment_record[metric.name] = metric.value
                    self.experiment_records[experiment.id] = experiment_record

                    if self.selected_metric is None:
                        self.selected_metric = metric.name

        if self.selected_metric not in self.metric_names:
            raise ValueError(
                f"no metric named `selected_metric` '{self.selected_metric}'"
                " logged to any experiment in `experiments`."
            )

    def register_callbacks(self, link_experiment_table=False):
        outputs = [
            Output("metric-heatmap", "figure"),
            Output("metric-heatmap", "style"),
            Output("metric-header-right-text", "children"),
            Output("metric-dropdown", "label"),
        ]
        inputs = [Input({"type": "metric-dropdown-button", "index": ALL}, "n_clicks")]
        states = []

        if link_experiment_table:
            inputs.append(
                Input("experiment-table", "derived_virtual_selected_row_ids"),
            )

        @self.app.callback(outputs, inputs, states)
        def update_metric_heatmap(*args):
            """Render the heatmap based on the currently selected metric.

            Returns the Plotly annotated heatmap generated by the values of the
            experiments' selected metric, the heatmap's style, the header text
            with the metric's name, and the selected metric's name.
            """
            if link_experiment_table:
                selected_row_ids = args[-1]
                selected_row_ids = selected_row_ids if selected_row_ids else []
            else:
                selected_row_ids = self.experiment_records.keys()

            property_id = callback_context.triggered[0].get("prop_id")
            property_value = property_id[: property_id.index(".")]

            if not property_value or property_value == "experiment-table":
                selected_metric = self.selected_metric
            else:
                selected_metric = json.loads(property_value).get("index")

                self.selected_metric = selected_metric

            heatmap_data = []
            experiment_ids = []

            for experiment_id, experiment_record in self.experiment_records.items():
                if experiment_id in selected_row_ids:
                    metric_value = experiment_record.get(selected_metric)

                    if metric_value is not None:
                        heatmap_data.append(metric_value)
                        experiment_ids.append(experiment_id[:7])

            header_right_text = (
                f"over {len(experiment_ids)} experiment{'s' if len(experiment_ids) != 1 else ''}"
            )

            if len(heatmap_data) == 0:
                return [], {"display": "none"}, header_right_text, selected_metric

            data_array = np.array(heatmap_data)
            numerator = data_array - data_array.min(axis=0)
            denominator = data_array.max(axis=0) - data_array.min(axis=0)
            denominator[denominator == 0] = 1
            scaled_heatmap_data = numerator / denominator

            annotations = copy.deepcopy(heatmap_data)
            for i, row in enumerate(annotations):
                for j, label in enumerate(row):
                    if isinstance(label, float):
                        annotations[i][j] = round(label, 6)

            heatmap_margin = 30
            heatmap_min_width = 72
            heatmap_cell_width = 6
            heatmap_cell_height = heatmap_cell_width / 2
            heatmap_cell_buffer = 12

            heatmap = ff.create_annotated_heatmap(
                scaled_heatmap_data,
                annotation_text=annotations,
                colorscale="blues",
                hoverinfo="text",
                text=heatmap_data,
                x=self.column_names if len(self.column_names) == len(heatmap_data[0]) else None,
                y=experiment_ids,
            )
            heatmap.update_layout(
                margin_b=heatmap_margin,
                margin_t=heatmap_margin,
                modebar_orientation="v",
                plot_bgcolor=plot_background_blue,
            )
            heatmap.update_xaxes(gridcolor="white")
            heatmap.update_yaxes(gridcolor="white")

            heatmap_height = heatmap_cell_buffer + (len(heatmap_data) * heatmap_cell_height)
            heatmap_width = (
                heatmap_cell_buffer + (len(heatmap_data[0]) * heatmap_cell_width)
                if len(heatmap_data[0]) > 8
                else heatmap_min_width
            )
            heatmap_style = {"height": f"{heatmap_height}rem", "width": f"{heatmap_width}rem"}

            return heatmap, heatmap_style, header_right_text, selected_metric
