import math
import os
from typing import Union
from pathlib import Path

import numpy as np
import pandas as pd
from rpy2 import robjects as ro
from rpy2.robjects import pandas2ri, Formula
from sklearn.preprocessing import StandardScaler
from statsmodels import api as sm
from statsmodels.iolib.summary import Summary
from statsmodels.regression.linear_model import RegressionResults

from src.data_loading_and_saving.print_and_save_results import print_and_save_result
from src.utils.handle_r_dependencies import BAS


class Regression:
    """
    A class used to perform Linear and Bayesian regression analyses.

    ...

    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
    -------
    linear_regression(data, x_vector, y, standardize, name_save_file) -> RegressionResults:
        Performs an OLS (Ordinary Least Squares) linear regression on the provided dataset.
    linear_regression_grouped(data, x_vector, y, dictionary_aggregation_methods_for_data_columns,
        column_to_group_by, standardize, name_save_file) -> RegressionResults:
        Performs Linear Regression on the data grouped by the given column as per the aggregation dictionary.
    predict_percentage_increase_between_liner_model_points(model, data_point_1, data_point_2) -> float:
        Calculates the percentage increase in fitted values between two different data points using a given model.
    report_effect_size_of_model(model, name_save_file) -> None:
        Computes and reports the effect size of the model
        (r-squared, Cohen's f, and equivalent Cohen's d for the regression model).
    bayesian_regression(data, x_vector, y, name_save_file) -> None:
        Performs a Bayesian linear regression on the provided dataset.
    _create_dataframe_of_bayesian_regression(summary, x_vector) -> pd.DataFrame:
        Transforms the provided Bayesian regression summary into a DataFrame.
    _clean_names(name) -> str:
        Cleans the input string by replacing certain characters with underscores or removing them.
    _clean_column_names(dataframe) -> pd.DataFrame:
        Cleans the DataFrame column names using the "_clean_names" method.
    _get_clean_x_vector(x_vector) -> list[str]:
        Cleans the elements of the input list using the "_clean_names" method.
    """

    def __init__(
        self,
        print_result: bool = True,
        save_result: bool = True,
        filepath: str = "results/",
    ):
        """
        Initializes the Regression object with the provided parameters.

        Parameters
        ----------
        print_result: bool, optional
            A flag to indicate if the function should print the result to the standard output. Default is True.
        save_result: bool, optional
            A flag to indicate if the function should save the result to a file. Default is True.
        filepath: str, optional
            The directory path where the result files will be saved if save_result is True. Default is "results/".
        """
        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 linear_regression(
        self,
        data: pd.DataFrame,
        x_vector: list[str],
        y: str,
        standardize: bool = False,
        report_effect_size: bool = False,
        name_save_file: Path = "",
    ) -> RegressionResults:
        """
        This method performs Ordinary Least Squares (OLS) Linear Regression using the given parameters
        and either prints or saves the result based on the Regression object properties.

        Parameters
        ----------
        data: pandas.DataFrame
            The input dataframe which contains the data.
        x_vector: list
            The list of column names to be used as independent variables in the regression model.
        y: str
            The column name to be used as dependent variable in the regression model.
        standardize: bool
            If the independent variables should be standardized before fitting for better comparison.
        report_effect_size: bool
            If the effect size of the model should be reported (r-squared, Cohen's f, and equivalent Cohen's d).
        name_save_file: str
            The name of the file to which the result will be saved (if self.save_result is True).

        Returns
        -------
        model: RegressionResults
            The fitted regression model.
        """
        data: pd.DataFrame = data.dropna(subset=x_vector)
        data: pd.DataFrame = data.dropna(subset=y)

        x_vector_data: pd.DataFrame = data[x_vector]
        y_data: pd.Series = data[y]

        if standardize:
            scaler: StandardScaler = StandardScaler()
            x_vector_data_standardized: np.ndarray = scaler.fit_transform(x_vector_data)
            x_vector_data: pd.DataFrame = pd.DataFrame(
                x_vector_data_standardized, columns=x_vector
            )
            x_vector_data.set_index(y_data.index, inplace=True)

        x_vector_data: pd.DataFrame = sm.add_constant(x_vector_data)

        model: RegressionResults = sm.OLS(y_data, x_vector_data).fit()

        linear_regression_summary: Summary = model.summary()

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

        if report_effect_size:
            name_save_file_effect_size: Path = name_save_file.with_stem(
                name_save_file.stem + "_effect_size"
            )
            self._report_effect_size_of_model(model, name_save_file_effect_size)

        return model

    def linear_regression_grouped(
        self,
        data: pd.DataFrame,
        x_vector: list[str],
        y: str,
        dictionary_aggregation_methods_for_data_columns: dict[str, str],
        column_to_group_by: str,
        standardize: bool = False,
        report_effect_size: bool = False,
        print_detailed_coefficients: bool = False,
        name_save_file: Path = "",
    ) -> RegressionResults:
        """
        This method performs OLS Linear Regression on the data grouped by the group_by
        as per the aggregation dictionary and then the performs regression analysis.

        Parameters
        ----------
        data: pandas.DataFrame
            The input dataframe which contains the data.
        x_vector: list
            The list of column names to be used as independent variables in the regression model.
        y: str
            The column name to be used as dependent variable in the regression model.
        dictionary_aggregation_methods_for_data_columns: dict
            The dictionary specifying how to aggregate each column before regression.
            The keys of the dictionary should include all columns in X and y.
        column_to_group_by: str
            The column name to group the data by before aggregating.
        name_save_file: str
            The name of the file to which the result will be saved (if self.save_result is True).
        standardize: bool
            If the regression should be performed standardized instead to return beta factors.
        report_effect_size: bool
            If the effect size of the model should be reported (r-squared, Cohen's f, and equivalent Cohen's d).
        print_detailed_coefficients: bool
            If the coefficients should be printed separately with 10 point float accuracy and 95% CI.

        Returns
        -------
        model: RegressionResults
            The fitted regression model.
        """

        all_cols = set(x_vector + [y])
        if not all_cols.issubset(
            set(dictionary_aggregation_methods_for_data_columns.keys())
        ):
            raise ValueError(
                "dictionary_aggregation_methods_for_data_columns should contain all columns from X and y."
            )

        if column_to_group_by not in data.columns:
            raise ValueError(
                f"'{column_to_group_by}' column to group by not found in data."
            )

        grouped_data = (
            data.groupby(column_to_group_by)
            .agg(dictionary_aggregation_methods_for_data_columns)
            .reset_index()
        )

        model = self.linear_regression(
            data=grouped_data,
            x_vector=x_vector,
            y=y,
            standardize=standardize,
            report_effect_size=report_effect_size,
            name_save_file=name_save_file,
        )

        if print_detailed_coefficients:
            coefficients = model.params
            confidence_intervals = model.conf_int()
            for name, coefficient in coefficients.items():
                confidence_interval = confidence_intervals.loc[name]
                detailed_coefficient_information: str = f"{name}: {coefficient: .10f} (CI: [{confidence_interval[0]: .10f}, {confidence_interval[1]: .10f}])"
                print_and_save_result(
                    self.print_result,
                    self.save_result,
                    self.filepath,
                    detailed_coefficient_information,
                    name_save_file,
                )

        return model

    @staticmethod
    def predict_percentage_increase_between_liner_model_points(
        model: RegressionResults,
        data_point_1: Union[list, np.array],
        data_point_2: Union[list, np.array],
    ) -> float:
        """
        Calculate the percentage increase in fitted values between two different data points using a given model.

        Parameters
        ----------
        model: statsmodels.regression.linear_model.OLS
            The fitted regression model.
        data_point_1: list or numpy.array
            The first data point (a vector of X values).
        data_point_2: list or numpy.array
            The second data point (a vector of X values).

        Returns
        -------
        percentage_increase: float
            The percentage increase in fitted value from data_point1 to data_point2.
        """
        if not isinstance(data_point_1, pd.Series):
            data_point_1 = pd.Series(data_point_1, index=model.model.exog_names[1:])
        if not isinstance(data_point_2, pd.Series):
            data_point_2 = pd.Series(data_point_2, index=model.model.exog_names[1:])

        if len(data_point_1) != len(data_point_2):
            raise ValueError("Data points must have the same dimension.")
        if len(data_point_1) != len(model.params) - 1:
            raise ValueError("Dimensions of data points and model do not match.")

        x_vector_data: pd.DataFrame = pd.DataFrame(
            [data_point_1, data_point_2], columns=model.model.data.orig_exog.columns[1:]
        )
        x_vector_data: pd.DataFrame = sm.add_constant(x_vector_data)

        predictions_for_both_datapoints: np.ndarray = model.predict(x_vector_data)

        percentage_increase: float = (
            (predictions_for_both_datapoints[1] - predictions_for_both_datapoints[0])
            / predictions_for_both_datapoints[0]
        ) * 100

        return percentage_increase

    def _report_effect_size_of_model(
        self, model: RegressionResults, name_save_file: Path
    ) -> None:
        """
        Computes and reports the effect size of the model (r-squared, Cohen's f and equivalent Cohen's d
        for the regression model) and either print or save it based on the Regression object properties.

        Parameters
        ----------
        model: RegressionResults
            The regression model results obtained from statsmodels regression.
        name_save_file: Path
            The name of the file to which the result will be saved if self.save_result is True.
        """
        r_squared: int = model.rsquared
        f_squared: float = r_squared / (1 - r_squared)
        cohens_f: float = math.sqrt(f_squared)
        equivalent_cohens_d: float = 2 * cohens_f

        effect_size_model: str = f"""
        R-Square R^2: {r_squared}
        Cohen's f: {cohens_f}
        Equivalent Cohen's d: {equivalent_cohens_d}
        """

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

    def bayesian_regression(
        self, data: pd.DataFrame, x_vector: list[str], y: str, name_save_file: Path
    ) -> None:
        """
        This method performs Bayesian Regression using the given parameters and either
        prints or saves the result based on the Regression object properties.

        the results can be interpreted in the following way:
        - `P(B != 0 | Y)`: This column gives the probability that the coefficient,
            associated with the factor (the B-value), is not zero given the data Y.
            This is a way of measuring the relevance and significance of the predictive factor.
            The higher the score, the more likely it is that the factor has an impact on your outcome variable.
        - `Model 1, Model 2, Model 3, etc`: These columns represent different models that have been built.
            The values in these columns are the specific coefficient values for each factor within each model.
        - `BF`: BF is short for Bayes Factor which indicates the strength of evidence for a model as
            compared to an alternative model. Higher BF values indicate stronger evidence for a model.
        - `PostProbs`: Posterior probabilities for each model. These add up to 1 across all models and
            give the probability of each model being the "best" model given the data.
        - `R2`:  Shows the proportion of the variance for a dependent variable that's explained by
            independent variables in the model. Close to 1 indicates the model explains a large amount of the variance.
        - `dim`: The dimensionality of the model, which is basically the number of parameters in each model.
        - `logmarg`: It is the logarithm of the marginal likelihood for each model.
            This is a factor that is used in determining the posterior probabilities of each model.

        Parameters
        ----------
        data: pandas.DataFrame
            The input dataframe which contains the data.
        x_vector: list
            The list of column names to be used as independent variables in the regression model.
        y: str
            The column name to be used as dependent variable in the regression model.
        name_save_file: str
            The name of the file to which the result will be saved (if self.save_result is True).

        Returns
        -------
        None
        """
        data: pd.DataFrame = data.dropna(subset=x_vector)

        pandas2ri.activate()

        data: pd.DataFrame = self._clean_column_names(data)
        x_vector: list[str] = self._get_clean_x_vector(x_vector)
        y: str = self._clean_names(y)

        formula: Formula = Formula(y + " ~ " + " + ".join(x_vector))

        data_r = pandas2ri.py2rpy(data)

        bayesian_model = BAS.bas_lm(
            formula=formula,
            data=data_r,
            method="MCMC",
            prior="ZS-null",
            modelprior=BAS.uniform(),
        )

        bayesian_regression_summary: np.ndarray = ro.r["summary"](bayesian_model)

        try:
            dataframe_bayesian_regression_summary = (
                self._create_dataframe_of_bayesian_regression(
                    bayesian_regression_summary, x_vector
                )
            )
        except ValueError:
            raise ValueError(
                "The number of bayesian models does not match the expected values."
            )

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

    @staticmethod
    def _create_dataframe_of_bayesian_regression(
        summary: np.ndarray, x_vector: list[str]
    ) -> pd.DataFrame:
        """
        Create a DataFrame from the summary of Bayesian regression results.

        Parameters
        ----------
        summary: np.ndarray
            The Bayesian regression summary obtained from performing a Bayesian linear regression.
        x_vector: list[str]
            The list of column names used as independent variables in the Bayesian regression model.

        Returns
        -------
        pd.DataFrame
            A DataFrame that presents the Bayesian regression's summary, including the probability
            of the regression coefficients not being zero, model number, intercept, independent variables,
            Bayes Factor (BF), posterior probabilities (PostProbs), R squared (R2), dimension (dim),
            and log marginal likelihood (logmarg).
        """
        base_columns: list[str] = ["P(B != 0 | Y)"]
        model_columns: list[str] = [
            "model {}".format(i + 1) for i in range(summary.shape[1] - 1)
        ]
        columns: list[str] = base_columns + model_columns

        index: list[str] = [
            "Intercept",
            *x_vector,
            "BF",
            "PostProbs",
            "R2",
            "dim",
            "logmarg",
        ]

        dataframe_summary: pd.DataFrame = pd.DataFrame(
            summary, columns=columns, index=index
        )

        return dataframe_summary

    @staticmethod
    def _clean_names(name: str) -> str:
        """
        This method performs cleaning of individual name (could be column name or value in X).
        Cleaning includes replacing certain characters with underscores or removing them.

        Parameters
        ----------
        name: str
            Original name (un-cleaned).

        Returns
        -------
        str
            Cleaned up name.
        """
        return (
            name.replace(" ", "_")
            .replace("(", "")
            .replace(")", "")
            .replace("-", "_")
            .replace("+", "_")
            .replace("=", "")
            .replace(".", "")
        )

    def _clean_column_names(self, dataframe: pd.DataFrame) -> pd.DataFrame:
        """
        This method performs cleaning of DataFrame column names using the method clean_names.
        It cleans a deep copy to not alter the original dataframe.

        Parameters
        ----------
        dataframe: pd.DataFrame
            DataFrame with original (un-cleaned) column names.

        Returns
        -------
        pd.DataFrame
            DataFrame with cleaned column names.
        """
        dataframe: pd.DataFrame = dataframe.copy()
        dataframe.columns = [self._clean_names(col) for col in dataframe.columns]
        return dataframe

    def _get_clean_x_vector(self, x_vector: list[str]) -> list[str]:
        """
        This method uses clean_names to clean the individual elements/column names in x_vector.

        Parameters
        ----------
        x_vector: list[str]
            List with original (un-cleaned) elements.

        Returns
        -------
        list[str]
            List where all elements have been cleaned.
        """
        return [self._clean_names(x) for x in x_vector]
