import os
from typing import Union, Optional
from pathlib import Path

import numpy as np
import pandas as pd

from src.data_classes.parameters_descriptive_overview import Metric
from src.data_loading_and_saving.print_and_save_results import print_and_save_result


class DescriptiveAnalysis:
    """
    A class used to perform descriptive data analysis.

    Attributes
    ----------
    print_result: bool
        A flag used to indicate if the function should print the result to the standard output.
    save_result: bool
        A flag used to indicate if the function should save the result to a file.
    filepath: str
        The directory path where the result files will be saved if save_result is True.
    ...

    Methods
    -------
    create_descriptives_for_metrics(data, metrics, name_save_file, group_by=None):
        Performs multiple descriptive analyses on the provided dataset and either prints or saves the result.
    create_descriptive_aggregated_for_metrics(data, variables, aggregation_function, group_by, name_save_file):
        Performs descriptive analysis on the provided dataset with aggregation and grouping, and either prints
        or saves the result.
    give_percentage_of_dataset_under_condition(data, variable, comparison, condition, name_save_file):
        Calculates and prints/saves the percentage of the dataset that meets a specified condition.
    _compute_metrics(group_name, data, metrics, group_column=None):
        Helper method to compute specified metrics on the data.
    _count_values(data, column):
        Counts the non-null values in a specified column or in the dataframe if no column is specified.
    _count_nonzero(data, column):
        Counts the non-zero/True values in a specified column.
    _count_unique(data, column):
        Counts the unique values in a specified column.
    _sum_values(data, column):
        Sums the values in a specified column.
    _mean_values(data, column):
        Calculates the mean of the values in a specified column.
    _std_dev(data, column):
        Calculates the standard deviation of the values in a specified column.
    """

    def __init__(
        self,
        print_result: bool = True,
        save_result: bool = True,
        filepath: str = "results/",
    ):
        """
        Constructs all the necessary attributes for the DescriptiveAnalysis object.

        Parameters
        ----------
        print_result: bool
            A flag to determine if the result should be printed.
        save_result: bool
            A flag to determine if the result should be saved.
        filepath: str
            The directory where the result should be saved.
        """
        self.print_result: bool = print_result
        self.save_result: bool = save_result
        self.filepath: str = filepath
        if not os.path.isdir(filepath):
            os.makedirs(filepath)

    def create_descriptives_for_metrics(
        self,
        data: pd.DataFrame,
        metrics: list[Metric],
        name_save_file: Path,
        group_by: Optional[str] = None,
    ) -> None:
        """
        This method performs multiple descriptive analysis on the provided dataset and either prints or saves the result

        Parameters
        ----------
        data: pandas.DataFrame
            The input dataframe which contains the data.
        metrics: MetricList
            The list of metrics to be calculated.
        name_save_file: str
            The name of the file to which the result will be saved (if self.save_result is True).
        group_by: str
            The column to group by. This gives the option to perform the analysis on each unique group member.

        Returns
        -------
        None
        """
        dataframes: list = []
        result_dataframes: pd.DataFrame = self._compute_metrics("Total", data, metrics)
        dataframes.append(result_dataframes)

        if group_by:
            grouped_dataframe = data.groupby(group_by)
            for group_name, group in grouped_dataframe:
                result_dataframes: pd.DataFrame = self._compute_metrics(
                    group_name, group, metrics
                )
                dataframes.append(result_dataframes)

        result_dataframes: pd.DataFrame = pd.concat(dataframes)
        result_dataframes.set_index(result_dataframes.columns[0], inplace=True)

        print_and_save_result(
            self.print_result,
            self.save_result,
            self.filepath,
            result_dataframes,
            name_save_file,
        )

    def create_descriptive_aggregated_for_metrics(
        self,
        data: pd.DataFrame,
        variables: list[str],
        aggregation_function: str,
        group_by: str,
        name_save_file: Path,
    ) -> None:
        """
        This method performs multiple descriptive analysis on the provided dataset with grouping.

        Parameters
        ----------
        data : pd.DataFrame
            The dataset to perform descriptive analysis on.
        variables : list[str]
            The list of columns we want information on.
        aggregation_function : str
            The method to aggregate the data. Either 'sum' or 'mean'.
        group_by : str
            The column to group by.
        name_save_file : str
            The name for saving the result.

        Returns
        -------
        None
        """
        if aggregation_function == "sum":
            if "Count" in variables:
                variables.remove("Count")
                grouped_data: pd.DataFrame = (
                    data.groupby(group_by)[variables]
                    .sum()
                    .aggregate(["mean", "std", "max", "min"])
                )

                count_data: pd.DataFrame = (
                    data.groupby(group_by)
                    .size()
                    .reset_index(name="Count")["Count"]
                    .aggregate(["mean", "std", "max", "min"])
                )
                grouped_data: pd.DataFrame = pd.concat(
                    [count_data, grouped_data], axis=1
                )
            else:
                grouped_data: pd.DataFrame = (
                    data.groupby(group_by)[variables]
                    .sum()
                    .aggregate(["mean", "std", "max", "min"])
                )

        elif aggregation_function == "mean":
            if "Count" in variables:
                variables.remove("Count")
                grouped_data: pd.DataFrame = (
                    data.groupby(group_by)[variables]
                    .mean()
                    .aggregate(["mean", "std", "max", "min"])
                )

                count_data: pd.DataFrame = (
                    data.groupby(group_by)
                    .size()
                    .reset_index(name="Count")["Count"]
                    .aggregate(["mean", "std", "max", "min"])
                )
                grouped_data: pd.DataFrame = pd.concat(
                    [count_data, grouped_data], axis=1
                )
            else:
                grouped_data: pd.DataFrame = (
                    data.groupby(group_by)[variables]
                    .mean()
                    .aggregate(["mean", "std", "max", "min"])
                )

        else:
            raise ValueError(f"Invalid aggregation: {aggregation_function}")

        print_and_save_result(
            self.print_result,
            self.save_result,
            self.filepath,
            grouped_data,
            name_save_file,
        )

    def give_percentage_of_dataset_under_condition(
        self,
        data: pd.DataFrame,
        variable: str,
        comparison: str,
        condition: Union[int, float],
        name_save_file: Path,
    ) -> None:
        """
        Calculates the percentage of the dataset that meets a specified condition and either prints or saves the result.

        This method evaluates a condition on a specified column of the dataset and calculates the percentage of rows
        that meet this condition. The result can be printed to the console or saved to a file, depending on the
        object's attributes.

        Parameters
        ----------
        data : pd.DataFrame
            The dataset to evaluate the condition on.
        variable : str
            The column name on which the condition will be applied.
        comparison : str
            The type of comparison to perform. Valid options are "smaller", "larger", or "not".
        condition : Union[int, float]
            The value to compare against the data in the specified column.
        name_save_file : Path
            The path (including filename) where the result should be saved if saving is enabled.

        Returns
        -------
        None
        """
        total_data: int = len(data)
        conditional_data: int = 0
        if comparison == "smaller":
            conditional_data: int = len(data[data[variable] <= condition])
        elif comparison == "larger":
            conditional_data: int = len(data[data[variable] >= condition])
        elif comparison == "not":
            conditional_data: int = len(data[data[variable] != condition])
        else:
            ValueError("Invalid comparison type")

        percentage: float = conditional_data / total_data * 100

        percentage_result: str = (
            f"The percentage of the dataset under that condition is {percentage} %"
        )

        print_and_save_result(
            self.print_result,
            self.save_result,
            self.filepath,
            percentage_result,
            name_save_file,
        )

    def _compute_metrics(
        self, group_name, data: pd.DataFrame, metrics: list[Metric], group_column=None
    ) -> pd.DataFrame:
        """
        Computes specified metrics on the given dataset and returns the results in a DataFrame.

        This method iterates over a list of Metric objects, each defining an operation (e.g., count, sum, mean) and
        a column on which the operation is to be performed. It constructs a result dictionary where each key is a
        metric name with the operation and column, and the value is the result of the operation. This dictionary is
        then converted to a DataFrame and returned.

        Parameters
        ----------
        group_name : str
            The name of the group for which metrics are being computed. This is used as a prefix in the result
            dictionary keys if not None.
        data : pd.DataFrame
            The dataset on which the metrics are to be computed.
        metrics : list[Metric]
            A list of Metric objects, each specifying an operation and a column.
        group_column : str, optional
            The name of the column by which the data was grouped, if any. This is used in the result dictionary keys.

        Returns
        -------
        pd.DataFrame
            A DataFrame containing the computed metrics, with each row representing the results for a group (if
            group_name is not None) or the entire dataset.
        """
        result_dict: dict = {group_column: group_name} if group_name is not None else {}
        for metric in metrics:
            operation: str = metric.operation
            column: str = metric.column
            if operation == "count":
                result_dict[
                    f'{column if column else "total"}_count'
                ] = self._count_values(data, column)
            elif operation == "count_nonzero":
                result_dict[f"{column}_nonzero"] = self._count_nonzero(data, column)
            elif operation == "count_unique":
                result_dict[f"{column}_unique"] = self._count_unique(data, column)
            elif operation == "sum":
                result_dict[f"{column}_sum"] = self._sum_values(data, column)
            elif operation == "mean":
                result_dict[f"{column}_mean"] = self._mean_values(data, column)
            elif operation == "std_dev":
                result_dict[f"{column}_std_dev"] = self._std_dev(data, column)

        return pd.DataFrame(result_dict, index=[0])

    @staticmethod
    def _count_values(data: pd.DataFrame, column: str) -> int:
        return len(data[column].dropna()) if column else len(data)

    @staticmethod
    def _count_nonzero(data: pd.DataFrame, column: str) -> int:
        data_type = data[column].dtype
        if data_type == bool:
            return len(data[data[column] == True])
        return len(data[data[column] != 0])

    @staticmethod
    def _count_unique(data: pd.DataFrame, column: str) -> int:
        first_element = data[column].dropna().iloc[0]
        if isinstance(first_element, (list, np.ndarray)):
            flattend_entries = pd.Series(
                [item for sublist in data[column] for item in sublist]
            )
            return flattend_entries.nunique()
        else:
            return data[column].nunique()

    @staticmethod
    def _sum_values(data: pd.DataFrame, column: str) -> Union[int, float]:
        return data[column].sum()

    @staticmethod
    def _mean_values(data: pd.DataFrame, column: str) -> Union[int, float]:
        return data[column].mean()

    @staticmethod
    def _std_dev(data: pd.DataFrame, column: str) -> Union[int, float]:
        return data[column].std()
