"""Group analysis module."""

from typing import Any, Callable, Dict, List, Optional, Union

import colorlover as cl
import nbformat.v4 as nbfv4
import numpy as np
import pandas as pd
import plotly.graph_objects as go
from IPython.display import Markdown, display
from matplotlib import cm
from matplotlib.colors import Normalize, to_hex
from plotly.subplots import make_subplots

from edvart import utils
from edvart.data_types import DataType, infer_data_type
from edvart.report_sections.code_string_formatting import code_dedent, get_code, total_dedent
from edvart.report_sections.section_base import Section, Verbosity


class GroupAnalysis(Section):
    """Generate the group analysis section of the report.

    Parameters
    ----------
    df : pd.DataFrame
        Data for which to perform analysis.
    groupby : Union[str, List[str]]
        Name of column or list of columns names to group by.
    verbosity : Verbosity (default = Verbosity.LOW)
        Generated code verbosity global to the Group analysis sections.
        If subsection verbosities are None, then they will be overridden by this parameter.
    columns : List[str], optional
        Columns on which to do group analysis. All columns are used by default.
    show_within_group_statistics : bool (default = True)
        Whether to show per-group statistics.
    show_group_missing_values : bool (default = True)
        Whether to show per-group missing values.
    show_group_distribution_plots : bool (default = True)
        Whether to show per-group distribution plots.

    Raises
    ------
    ValueError
        If groupby columns are not a subset of the columns of the input DataFrame df.
    """

    def __init__(
        self,
        df: pd.DataFrame,
        groupby: Union[str, List[str]],
        verbosity: Verbosity = Verbosity.LOW,
        columns: Optional[List[str]] = None,
        show_within_group_statistics: bool = True,
        show_group_missing_values: bool = True,
        show_group_distribution_plots: bool = True,
    ):
        if isinstance(groupby, str):
            groupby = [groupby]
        if not set(groupby) <= set(df.columns):
            raise ValueError("Grouping by a column which is not in columns of df.")
        self.columns = columns
        self.df = df
        self.groupby = groupby
        self.show_statistics = show_within_group_statistics
        self.show_missing_vals = show_group_missing_values
        self.show_dist = show_group_distribution_plots

        super().__init__(verbosity, columns)

    @property
    def name(self) -> str:
        return "Group Analysis"

    @staticmethod
    def default_group_descriptive_stats() -> Dict[str, Callable[[pd.Series], float]]:
        """Descriptive statistic functions.

        Returns
        -------
        Dict[str, Callable[[pd.Series], float]]
            A dictionary of statistic function names and functions.
        """
        return {
            "# Unique values": utils.num_unique_values,
            "Sum": utils.sum_,
            "Mode": utils.mode,
            "Mean": utils.mean,
            "Std": utils.std,
            "Mean abs dev": utils.mad,
            "Median abs dev": utils.median_absolute_deviation,
            "Relative std": utils.coefficient_of_variation,
            "Kurtosis": utils.kurtosis,
            "Skewness": utils.skewness,
        }

    @staticmethod
    def default_group_quantile_stats() -> Dict[str, Callable[[pd.Series], float]]:
        """Quantile statistic functions.

        Returns
        -------
        Dict[str, Callable[[pd.Series], float]]
            A dictionary of statistic function names and functions.
        """
        return {
            "Min": utils.minimum,
            "Q1": utils.quartile1,
            "Median": utils.median,
            "Q3": utils.quartile3,
            "Max": utils.maximum,
            "Range": utils.value_range,
            "IQR": utils.iqr,
        }

    @staticmethod
    def within_group_descriptive_stats(
        df: pd.DataFrame, groupby: List[str], column: str, round_decimals: int = 2
    ):
        """Display within-group descriptive statistics for a column.

        Parameters
        ----------
        df : pd.DataFrame
            Data to display statistics for.
        groupby : List[str]
            List of column names to group data by.
        column : str
            Which column to display statistics for.
        round_decimals : int (default = 2)
            Number of decimals to round displayed results to.
        """
        GroupAnalysis.within_group_stats(
            df=df,
            groupby=groupby,
            column=column,
            stats=GroupAnalysis.default_group_descriptive_stats(),
            round_decimals=round_decimals,
        )

    @staticmethod
    def within_group_quantile_stats(
        df: pd.DataFrame, groupby: List[str], column: str, round_decimals: int = 2
    ) -> None:
        """Display within-group quantile statistics for a column.

        Parameters
        ----------
        df : pd.DataFrame
            Data to display statistics for.
        groupby : List[str]
            List of column names to group data by.
        column : str
            Which column to display statistics for.
        round_decimals : int (default = 2)
            Number of decimals to round displayed results to.
        """
        GroupAnalysis.within_group_stats(
            df=df,
            groupby=groupby,
            column=column,
            stats=GroupAnalysis.default_group_quantile_stats(),
            round_decimals=round_decimals,
        )

    @staticmethod
    def within_group_stats(
        df: pd.DataFrame,
        groupby: List[str],
        column: str,
        stats: Dict[str, Callable[[pd.Series], float]] = None,
        round_decimals: int = 2,
    ) -> None:
        """Display withing group statistics for a column of df grouped by one or other more columns.

        Parameters
        ----------
        df : pd.DataFrame
            Data to display statistics for.
        groupby : List[str]
            List of column names to group by.
        column : str
            Name of column to display statistics for.
        stats : Dict[str, Callable[[pd.Series], float]], optional
            A dictionary of statistic function names and functions.
            If None, default_group_quantile_stats() and default_group_descriptive_stats() will be
            used.
        round_decimals : int (default = 2)
            Number of decimals to round displayed results to.
        """
        if stats is None:
            GroupAnalysis.within_group_descriptive_stats(df, groupby, column, round_decimals)
            GroupAnalysis.within_group_quantile_stats(df, groupby, column, round_decimals)
            return

        df_grouped = df.groupby(groupby)[column]
        group_stats = []
        for name, func in stats.items():
            group_stats.append(df_grouped.apply(func).rename(name))
        stats_table = pd.concat(group_stats, axis=1)
        stats_table = stats_table.round(decimals=round_decimals)
        display(stats_table)

    # pylint: disable=too-many-locals
    @staticmethod
    def group_missing_values(
        df: pd.DataFrame,
        groupby: Union[str, List[str]],
        columns: Optional[List[str]] = None,
        round_decimals: int = 2,
        heatmap: bool = True,
        foreground_colormap: str = "bone",
        background_colormap: str = "OrRd",
        sort: bool = True,
        sort_by: Optional[List[str]] = None,
        ascending: bool = False,
    ) -> None:
        """Display per-group number and percentage of missing values in each column.

        Parameters
        ----------
        df : pd.DataFrame
            Data to display missing values for.
        groupby : str or List[str]
            Name of column or list of columns names to group by.
        columns : List[str], optional
            Subset of columns to analyze. All columns except those for grouping are used by default.
        round_decimals : int (default = 2)
            Number of decimals to round displayed results to.
        heatmap : bool (default = True)
            Whether to color missing value percentage cells according to the corresponding value.
        foreground_colormap : str (default = "bone")
            Color map of the foreground.
        background_colormap : str (default = "OrRd)
            Color map of the background.
        sort : bool (default = True)
            Whether to sort the results.
        sort_by : List[str], optional
            List of column names to sort the results by. Sort by all column by default.
        ascending : bool (default = False)
            If True, sort in ascending order, otherwise sort in descending order.

        Raises
        ------
        ValueError
            If groupby columns are not a subset of the columns of the input DataFrame df.
        """
        if isinstance(groupby, str):
            groupby = [groupby]
        if not set(groupby) <= set(df.columns):
            raise ValueError("Grouping by a column which is not in columns of df.")
        if columns is None:
            columns = [col for col in df.columns if col not in groupby]
        df_grouped = df.groupby(groupby)[columns]

        # Calculate number of samples in each group
        sizes = df_grouped.size().rename("Group Size")

        # Calculate missing values percentage of each column for each group
        missing = df_grouped.apply(lambda g: g.isna().sum(axis=0))
        missing = missing.divide(sizes, axis=0) * 100  # `/` divides by axis=-1
        missing = missing.fillna(value=0)
        missing = missing.round(decimals=round_decimals)

        if missing.sum().sum() == 0:
            print("There are no missing values")
            return

        # Concatenate group sizes and missing value percentages
        final_table = pd.concat([sizes, missing], axis=1)

        # Sort columns to better identify groups with missing data
        all_columns = [col for col in missing.columns if col not in groupby]
        if sort:
            if sort_by is None:
                sort_by = all_columns
            final_table = final_table.sort_values(by=sort_by, axis=0, ascending=ascending)

        # Drop columns with no missing data
        missing = final_table.sum(axis=0) != 0
        final_table = final_table.loc[:, missing]

        colored_columns = [col for col in final_table if col in all_columns]

        # Apply conditional formatting to each cell except group size column
        if heatmap:
            fg_cmap = cm.get_cmap(foreground_colormap)
            bg_cmap = cm.get_cmap(background_colormap)
            norm = Normalize(vmin=0, vmax=100)

            def color_cell(value):
                fg_hex = to_hex(fg_cmap(norm(value)))
                bg_hex = to_hex(bg_cmap(norm(value)))
                return f"""
                    color: {fg_hex};
                    background-color: {bg_hex};
                """

            render = final_table.style.applymap(
                func=color_cell, subset=pd.IndexSlice[:, colored_columns]
            ).format(formatter="{0:.2f} %", subset=pd.IndexSlice[:, colored_columns])
        else:
            render = final_table.style.format(
                formatter="{0:.2f} %", subset=pd.IndexSlice[:, colored_columns]
            )

        # Render table
        display(render)

    @staticmethod
    def group_barplot(
        df: pd.DataFrame,
        groupby: List[str],
        column: str,
        group_count_threshold: int = 20,
        conditional_probability: bool = True,
        xaxis_tickangle: float = 0,
        alpha: float = 0.5,
    ):
        """Display a per-group barplot for a column.

        Parameters
        ----------
        df : pd.DataFrame
            Data to analyze.
        groupby : List[str]
            List of column names to group by.
        column : str
            Which column to analyze.
        group_count_threshold : int (default = 20)
            Maximum number of unique values in column to plot. If the number of unique values
            is higher, a warning will be issued and plot will not be shown.
        conditional_probability : bool (default = True)
            If True, conditional probability contitioned on group will be displayed,
            otherwise conditional frequency will be displayed.
        xaxis_tickangle : float (default = 0)
            Rotation angle of ticks on the x axis.
        alpha : float (default = 0.5)
            Opacity of bars in the plot.
        """
        num_cat = df[column].nunique()
        if num_cat > group_count_threshold:
            display(
                Markdown(
                    f"Number of unique values ({num_cat}) greater than threshold."
                    " Not plotting distribution"
                )
            )
            return

        pivot = df.pivot_table(index=groupby, columns=column, aggfunc="size", fill_value=0)

        if conditional_probability:
            pivot = pivot.divide(pivot.sum(axis=1), axis=0)  # `/` divides by axis=-1
            pivot = pivot.fillna(value=0)

        # Choose color palette
        colors = cl.scales["9"]["qual"]["Set1"]
        color_idx = 0

        fig = go.Figure()
        for idx, row in pivot.iterrows():
            if hasattr(idx, "__len__") and not isinstance(idx, str):
                group_name = "_".join([str(i) for i in idx])
            else:
                group_name = idx
            color = colors[color_idx % len(colors)]
            color_idx += 1
            fig.add_trace(
                go.Bar(
                    x=pivot.columns,
                    y=row,
                    name=group_name,
                    opacity=alpha,
                    marker_color=color,
                )
            )

        groupby_repr = ", ".join(groupby)
        if conditional_probability:
            yaxis_title = f"P({column} | {groupby_repr})"
        else:
            yaxis_title = f"Freq({column} | {groupby_repr})"

        fig.update_layout(
            barmode="group",
            xaxis_tickangle=xaxis_tickangle,
            xaxis_title=column,
            yaxis_title=yaxis_title,
        )
        fig.show()

    @staticmethod
    def overlaid_histograms(
        df: pd.DataFrame,
        groupby: List[str],
        column: str,
        bins: Optional[int] = None,
        density: bool = True,
        alpha: float = 0.5,
    ):
        """Show per-group distribution histograms in a single plot overlaid over each other.

        Parameters
        ----------
        df : pd.DataFrame
            Data to analyze.
        groupby : List[str]
            List of column names to group by.
        column : str
            Name of column to analyze.
        bins : int, optional
            Number of bins in the histogram. If None, number of bin will be inferred using
            Modified Freedman-Diaconis bin number inference.
        density : bool (default = True)
            If True, histograms will be normalized to display density.
        alpha : float
            Opacity of individual histograms.
        """
        # Modified Freedman-Diaconis bin number inference if bins is None
        if bins is None:
            iqr = df[column].quantile(0.75) - df[column].quantile(0.25)
            bin_width = 1 / np.cbrt(len(df)) * iqr
        else:
            bin_width = (df[column].max() - df[column].min()) / bins
        bin_config = {
            "start": df[column].min(),
            "end": df[column].max(),
            "size": bin_width,
        }

        # Choose color palette
        colors = cl.scales["9"]["qual"]["Set1"]
        color_idx = 0

        # Distribution plot
        fig = make_subplots(
            rows=2,
            cols=1,
            shared_xaxes=True,
            row_heights=[0.3, 0.7],
            vertical_spacing=0.02,
        )

        for name, group in df.groupby(
            by=(groupby[0] if isinstance(groupby, list) and len(groupby) == 1 else groupby)
        ):
            if hasattr(name, "__len__") and not isinstance(name, str):
                group_name = "_".join([str(i) for i in name])
            else:
                group_name = name
            color = colors[color_idx % len(colors)]
            color_idx += 1
            # Add to boxplot
            fig.add_trace(
                go.Box(
                    x=group[column],
                    name=group_name,
                    legendgroup=group_name,
                    showlegend=False,
                    marker_color=color,
                ),
                row=1,
                col=1,
            )
            # Add to histogram
            fig.add_trace(
                go.Histogram(
                    x=group[column],
                    name=group_name,
                    legendgroup=group_name,
                    xbins=bin_config,
                    histnorm="density" if density else "",
                    marker_color=color,
                    opacity=alpha,
                ),
                row=2,
                col=1,
            )
        fig.update_layout(barmode="overlay")
        fig.update_xaxes(title_text=column, row=2, col=1)
        yaxis_title = "Density" if density else "Frequency"
        fig.update_yaxes(title_text=yaxis_title, row=2, col=1)
        fig.show()

    @staticmethod
    def group_analysis(
        df: pd.DataFrame,
        groupby: Union[str, List[str]],
        columns: Optional[List[str]] = None,
        show_within_group_statistics: bool = True,
        show_group_missing_values: bool = True,
        show_distribution_plots: bool = True,
    ) -> None:
        """Generate group analysis for df.

        Parameters
        ----------
        df : pd.DataFrame
            Data to be analyzed.
        groupby : Union[str, List[str]]
            Name of column or list of columns names to group by.
        columns : List[str], optional
            Subset of columns to analyze. All columns except those used for grouping are used
            by default.
        show_within_group_statistics : bool (default = True)
            Whether to show per-group statistics.
        show_group_missing_values : bool (default = True)
            Whether to show per-group missing values.
        show_distribution_plots : bool (default = True)
            Whether to show per-group distribution plots.

        Raises
        ------
        ValueError
            If groupby columns are not a subset of the columns of the input DataFrame df.
        """
        if isinstance(groupby, str):
            groupby = [groupby]
        if not set(groupby) <= set(df.columns):
            raise ValueError("Grouping by a column which is not in columns of df.")
        if columns is None:
            columns = [col for col in df.columns if col not in groupby]

        if show_group_missing_values:
            display(Markdown("## Missing values for each group"))
            GroupAnalysis.group_missing_values(df, groupby)
        if not show_distribution_plots and not show_within_group_statistics:
            return
        for col in columns:
            display(Markdown("---"))
            display(Markdown(f"### *{col}*"))
            datatype = infer_data_type(df[col])
            if datatype == DataType.NUMERIC:
                if show_within_group_statistics:
                    GroupAnalysis.within_group_stats(df, groupby, col)
                if show_distribution_plots:
                    GroupAnalysis.overlaid_histograms(df, groupby, col)
            else:
                GroupAnalysis.group_barplot(df, groupby, col)

    def required_imports(self) -> List[str]:
        """Returns a list of imports to be put at the top of a generated notebook.

        Returns
        -------
        List[str]
            List of import strings to be added at the top of the generated notebook,
            e.g. ["import pandas as pd", "import numpy as np"]
        """
        ga = "GroupAnalysis"  # pylint:disable=invalid-name
        if self.verbosity == Verbosity.LOW:
            return [
                f"from edvart.report_sections.group_analysis import {ga}\n"
                f"group_analysis = {ga}.group_analysis"
            ]
        if self.verbosity == Verbosity.MEDIUM:
            return [
                total_dedent(
                    f"""
                        from edvart.report_sections.group_analysis import {ga}
                        group_missing_values = {ga}.group_missing_values
                        within_group_stats = {ga}.within_group_stats
                        within_group_descriptive_stats = {ga}.within_group_descriptive_stats
                        within_group_quantile_stats = {ga}.within_group_quantile_stats
                        group_barplot = {ga}.group_barplot
                        overlaid_histograms = {ga}.overlaid_histograms
                        default_group_quantile_stats = {ga}.default_group_quantile_stats
                        default_group_descriptive_stats = {ga}.default_group_descriptive_stats
                    """
                )
            ]
        # verbosity HIGH
        return [
            "import colorlover as cl",
            "import matplotlib.pyplot as plt",
            "import matplotlib.cm as cm",
            "from matplotlib.colors import Normalize, to_hex",
            "import numpy as np",
            "import pandas as pd",
            "import plotly.graph_objects as go",
            "from edvart.data_types import infer_data_type, DataType",
            "from edvart import utils",
            "from typing import List, Dict, Optional, Callable",
            "from plotly.subplots import make_subplots",
        ]

    def _add_function_defs(self, cells: List[Dict[str, Any]]):
        """Add function definition cell to the list of cells.

        Cells can be either code cells or markdown cells.

        Parameters
        ----------
        cells : List[Dict[str, Any]]
            List of generated notebook cells which are represented as dictionaries.
        """
        code = (
            get_code(GroupAnalysis.default_group_quantile_stats)
            + "\n\n"
            + get_code(GroupAnalysis.default_group_descriptive_stats)
            + "\n\n"
            + (get_code(GroupAnalysis.within_group_descriptive_stats).replace("GroupAnalysis.", ""))
            + "\n\n"
            + (get_code(GroupAnalysis.within_group_quantile_stats).replace("GroupAnalysis.", ""))
            + "\n\n"
            + (get_code(GroupAnalysis.within_group_stats).replace("GroupAnalysis.", ""))
            + "\n\n"
            + get_code(GroupAnalysis.group_barplot)
            + "\n\n"
            + get_code(GroupAnalysis.overlaid_histograms)
            + "\n\n"
            + get_code(GroupAnalysis.group_missing_values)
        )

        cells.append(nbfv4.new_code_cell(code))

    def _add_cells_numeric_col(self, cells: List[Dict[str, Any]], column_name: str):
        """Add code cells for a numeric column at verbosity MEDIUM or HIGH.

        Parameters
        ----------
        cells : List[Dict[str, Any]]
            List of generated notebook cells which are represented as dictionaries.
        column_name : str
            Name of column for which to generate code.
        """
        code = ""
        if self.show_statistics:
            if self.verbosity == Verbosity.MEDIUM:
                code += (
                    f"within_group_stats(df=df, groupby={self.groupby}, column='{column_name}')\n"
                )
            else:
                code += code_dedent(
                    f"""
                            within_group_stats(
                                df=df,
                                groupby={self.groupby},
                                column='{column_name}',
                                stats=default_group_descriptive_stats()
                            )

                            within_group_stats(
                                df=df,
                                groupby={self.groupby},
                                column='{column_name}',
                                stats=default_group_quantile_stats()
                            )
                        """
                )

        if self.show_dist:
            code += f"overlaid_histograms(df=df, groupby={self.groupby}, column='{column_name}')"
        cells.append(nbfv4.new_code_cell(code))

    def add_cells(self, cells: List[Dict[str, Any]]):
        """Add cells to the list of cells.

        Cells can be either code cells or markdown cells.

        Parameters
        ----------
        cells : List[Dict[str, Any]]
            List of generated notebook cells which are represented as dictionaries.
        """
        section_header = nbfv4.new_markdown_cell(self.get_title(section_level=1))
        cells.append(section_header)

        if self.verbosity == Verbosity.LOW:
            if self.columns is None:
                code = f"group_analysis(df=df, groupby={self.groupby})"
            else:
                code = f"group_analysis(df=df, groupby={self.groupby}, columns={self.columns})"
            code_cell = nbfv4.new_code_cell(code)
            cells.append(code_cell)
            return

        if self.verbosity == Verbosity.HIGH:
            self._add_function_defs(cells)

        if self.show_missing_vals:
            cells.append(nbfv4.new_markdown_cell("## Missing values for each group"))
            if self.columns is None:
                code = f"group_missing_values(df=df, groupby={self.groupby})"
            else:
                code = (
                    f"group_missing_values(df=df, columns={self.columns}, groupby={self.groupby})"
                )
            cells.append(nbfv4.new_code_cell(code))

        columns = self.columns if self.columns is not None else self.df.columns

        if not self.show_statistics and not self.show_dist:
            return
        for col in columns:
            if col in self.groupby:
                continue
            cells.append(nbfv4.new_markdown_cell(f"## *{col}*"))
            datatype = infer_data_type(self.df[col])
            if datatype == DataType.NUMERIC:
                self._add_cells_numeric_col(cells, col)
            else:
                code = f"group_barplot(df=df, groupby={self.groupby}, column='{col}')"
                cells.append(nbfv4.new_code_cell(code))

    def show(self, df: pd.DataFrame) -> None:
        """Generates cell output of this section in the calling notebook.

        Parameters
        ----------
        df : pd.DataFrame
            Data based on which to generate the cell output.
        """
        columns = (
            self.columns
            if self.columns is not None
            else [col for col in df.columns if col not in self.groupby]
        )
        display(Markdown(self.get_title(section_level=1)))

        if self.show_missing_vals:
            display(Markdown("## Missing values for each group"))
            GroupAnalysis.group_missing_values(df, self.groupby)
        if not self.show_dist and not self.show_statistics:
            return
        for col in columns:
            if col in self.groupby:
                continue
            display(Markdown(f"## *{col}*"))
            datatype = infer_data_type(df[col])
            if datatype == DataType.NUMERIC:
                if self.show_statistics:
                    GroupAnalysis.within_group_stats(df, self.groupby, col)
                if self.show_dist:
                    GroupAnalysis.overlaid_histograms(df, self.groupby, col)
            else:
                GroupAnalysis.group_barplot(df, self.groupby, col)
