import os
from pathlib import Path

import pandas as pd
from scipy.stats import pearsonr

from src.data_loading_and_saving.print_and_save_results import print_and_save_result


class PearsonCorrelation:
    """
    A class used to calculate the Pearson's correlation coefficient between two variables.

    Attributes:
    ----------
        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.

    Methods
    -------
        calculate_correlation(data: pd.DataFrame, first_group_name: str, second_group_name: str, name_save_file: Path):
            This method calculates Pearson's correlation coefficient for the given columns
    """

    def __init__(
        self,
        print_result: bool = True,
        save_result: bool = True,
        filepath: str = "results/",
    ):
        """
        Constructs all the necessary attributes for the PearsonCorrelation 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 calculate_correlation(
        self,
        data: pd.DataFrame,
        first_group_name: str,
        second_group_name: str,
        name_save_file: Path,
    ) -> None:
        """
        This method calculates Pearson's correlation coefficient for the given columns
        and either prints or saves the result based on the object properties.

        Parameters
        ----------
            data: pd.DataFrame
                The data frame containing the data.
            first_group_name: str
                The name of the first group.
            second_group_name: str
                The name of the second group.
            name_save_file: Path
                The name of the file where the result should be saved.
        """
        data: pd.DataFrame = data[[first_group_name, second_group_name]].dropna()
        group_1: pd.DataFrame = data[first_group_name]
        group_2: pd.DataFrame = data[second_group_name]

        pearson_correlation: float
        p_value: float
        pearson_correlation, p_value = pearsonr(group_1, group_2)

        pearson_correlation_summary: str = (
            f"Pearson correlation between {first_group_name} and {second_group_name}: {pearson_correlation}\n"
            f"P-value: {p_value}"
        )

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