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

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import seaborn as sns
from statsmodels.regression.linear_model import RegressionResults

from src.data_loading_and_saving.save_plot import save_plot
from src.utils.helper_functions import FunctionData


class DataVisualizer:
    """
    A class used to perform various visualizations.

    This class provides methods to create different types of plots including bar charts, stacked bar charts,
    percentage stacked bar charts, count distribution plots, grouped histograms, simple scatter plots, heatmaps,
    surface plots, contour plots, density plots, ridgeline plots, hexbin plots with optional trendline, box plots,
    violin plots, and forest plots for regression model coefficients or paired t-test results. It supports saving
    plots to files.

    Attributes
    ----------
    save_plots : bool
        A flag used to indicate if the function should save the plots to a file.
    filepath : str
        The directory path where the plot files will be saved if save_plots is True.
    """

    def __init__(self, save_plots: bool = True, filepath="results/"):
        self.save_plots: bool = save_plots
        self.filepath: str = filepath
        if not os.path.isdir(filepath):
            os.makedirs(filepath)

    def create_bar_chart(
        self,
        data: pd.DataFrame,
        variable_x_axis: str,
        variable_y_axis: Optional[str] = None,
        x_axis_label: Optional[str] = None,
        y_axis_label: Optional[str] = None,
        title: Optional[str] = None,
        chart_orientation: str = "h",
        name_save_file: Optional[Path] = None,
        sort_order="ascending",
        custom_order: Optional[list[str]] = None,
    ) -> None:
        """
        Creates a bar chart based on the provided DataFrame and plotting parameters.

        Parameters
        ----------
        data : pd.DataFrame
            The DataFrame containing the data to plot.
        variable_x_axis : str
            The column name in `data` to be used as the x-axis variable.
        variable_y_axis : Optional[str], optional
            The column name in `data` to be used as the y-axis variable. If None, a count of occurrences is used.
        x_axis_label : Optional[str], optional
            The label for the x-axis. If None, `variable_x_axis` is used.
        y_axis_label : Optional[str], optional
            The label for the y-axis. If None and `variable_y_axis` is None, "Count" is used; otherwise,
            `variable_y_axis` is used.
        title : Optional[str], optional
            The title of the chart. If None, a default title is generated.
        chart_orientation : str, default "h"
            The orientation of the chart. "h" for horizontal, any other value for vertical.
        name_save_file : Optional[Path], optional
            The path (including filename) where the chart should be saved. If None, the chart is not saved.
        sort_order : str, default "ascending"
            The order in which the bars are sorted. Can be "ascending", "descending", or "custom".
        custom_order : Optional[list[str]], optional
            Specifies a custom order for the x-axis categories. Only effective if `sort_order` is "custom".

        Returns
        -------
        None
        """
        plt.figure(figsize=(10, 8))

        x_axis_label: str = x_axis_label if x_axis_label else variable_x_axis
        y_axis_label: str = (
            "Count"
            if variable_y_axis is None
            else y_axis_label
            if y_axis_label
            else variable_y_axis
        )

        if variable_y_axis is None:
            data: pd.DataFrame = data.copy()
            data["count"] = 1
            variable_y_axis: str = "count"

        title: str = (
            title if title else f"Bar Chart {variable_y_axis} by {variable_x_axis}"
        )

        aggregated_data: pd.DataFrame = (
            data.groupby(variable_x_axis)[variable_y_axis].sum().reset_index()
        )

        if sort_order == "ascending":
            aggregated_data: pd.DataFrame = aggregated_data.sort_values(
                by=variable_y_axis
            )
        elif sort_order == "descending":
            aggregated_data: pd.DataFrame = aggregated_data.sort_values(
                by=variable_y_axis, ascending=False
            )
        elif sort_order == "custom" and custom_order is not None:
            data[variable_x_axis] = pd.Categorical(
                data[variable_x_axis], categories=custom_order, ordered=True
            )
            aggregated_data: pd.DataFrame = (
                data.groupby([variable_x_axis]).sum().unstack().fillna(0)
            )

        if chart_orientation == "h":
            aggregated_data.plot(
                x=variable_x_axis, y=variable_y_axis, kind="barh", legend=False
            )
            plt.xlabel(y_axis_label)
            plt.ylabel(x_axis_label)
        else:
            aggregated_data.plot(
                x=variable_x_axis, y=variable_y_axis, kind="bar", legend=False
            )
            plt.xlabel(x_axis_label)
            plt.ylabel(y_axis_label)

        plt.title(title)

        save_plot(self.save_plots, self.filepath, plt, name_save_file)

    def create_stacked_bar_chart(
        self,
        data: pd.DataFrame,
        variable_x_axis: str,
        variable_y_axis: Optional[str],
        hue: str,
        x_axis_label: Optional[str] = None,
        y_axis_label: Optional[str] = None,
        title: Optional[str] = None,
        chart_orientation: str = "h",
        name_save_file: Optional[Path] = None,
        sort_order: str = "descending",
        custom_order: Optional[list[str]] = None,
    ) -> None:
        """
        Creates a stacked bar chart from the given DataFrame.

        This method generates a stacked bar chart visualizing the relationship between two variables,
        with the option to categorize data based on a third variable (`hue`). The chart can be sorted
        and customized in various ways, including orientation, sorting order, and custom sorting.

        Warning: This can take a very long time for large datasets due to sorting.

        Parameters
        ----------
        data : pd.DataFrame
            The DataFrame containing the data to be plotted.
        variable_x_axis : str
            The column name in `data` to be used as the x-axis variable.
        variable_y_axis : Optional[str]
            The column name in `data` to be used as the y-axis variable. If None, a count of occurrences is used.
        hue : str
            The column name to be used for categorizing the data into different sections of the stacked bar.
        x_axis_label : Optional[str], optional
            The label for the x-axis. If None, `variable_x_axis` is used.
        y_axis_label : Optional[str], optional
            The label for the y-axis. If None and `variable_y_axis` is None, "Count" is used; otherwise,
            `variable_y_axis` is used.
        title : Optional[str], optional
            The title of the chart. If None, a default title is generated based on the provided parameters.
        chart_orientation : str, default "h"
            The orientation of the chart. "h" for horizontal, any other value for vertical.
        name_save_file : Optional[Path], optional
            The path (including filename) where the chart should be saved. If None, the chart is not saved.
        sort_order : str, default "descending"
            The order in which the bars are sorted. Can be "ascending", "descending", or "custom".
        custom_order : Optional[list[str]], optional
            Specifies a custom order for the x-axis categories. Only effective if `sort_order` is "custom".

        Returns
        -------
        None
        """
        plt.figure(figsize=(10, 8))

        x_axis_label: str = x_axis_label if x_axis_label else variable_x_axis
        y_axis_label: str = (
            "Count"
            if variable_y_axis is None
            else y_axis_label
            if y_axis_label
            else variable_y_axis
        )

        if variable_y_axis is None:
            data: pd.DataFrame = data.copy()
            data["count"] = 1
            variable_y_axis: str = "count"

        title: str = (
            title
            if title
            else f"Stacked Bar Chart {variable_y_axis} by {variable_x_axis} stacked by {hue}"
        )

        aggregated_data: pd.DataFrame = (
            data.groupby([variable_x_axis, hue]).sum().unstack()
        )

        sum_of_each_group: pd.Series = aggregated_data.sum(axis=1)

        if sort_order == "ascending":
            aggregated_data: pd.DataFrame = aggregated_data.loc[
                sum_of_each_group.sort_values().index
            ]
        elif sort_order == "descending":
            aggregated_data: pd.DataFrame = aggregated_data.loc[
                sum_of_each_group.sort_values(ascending=False).index
            ]
        elif sort_order == "custom" and custom_order is not None:
            data[variable_x_axis] = pd.Categorical(
                data[variable_x_axis], categories=custom_order, ordered=True
            )
            aggregated_data: pd.DataFrame = (
                data.groupby([variable_x_axis, hue]).sum().unstack().fillna(0)
            )

        if chart_orientation == "h":
            aggregated_data.loc[:, variable_y_axis].plot(
                kind="barh", stacked=True, legend=True
            )
            plt.xlabel(y_axis_label)
            plt.ylabel(x_axis_label)
        else:
            aggregated_data.loc[:, variable_y_axis].plot(
                kind="bar", stacked=True, legend=True
            )
            plt.xlabel(x_axis_label)
            plt.ylabel(y_axis_label)

        plt.title(title)
        plt.legend(title=hue, title_fontsize="13", loc="best")
        save_plot(self.save_plots, self.filepath, plt, name_save_file)

    def create_percentage_stacked_bar_chart(
        self,
        data: pd.DataFrame,
        variable_x_axis: str,
        variables_to_compare: list[str],
        x_axis_label: Optional[str] = None,
        y_axis_label: str = "Percentage",
        title: Optional[str] = None,
        chart_orientation: str = "h",
        name_save_file: Optional[Path] = None,
        sort_order: str = "ascending",
    ) -> None:
        """
        Creates a percentage stacked bar chart for the specified variables in a DataFrame.

        This method visualizes the distribution of the specified variables as a percentage of their total
        across the categories defined by the `variable_x_axis`. It supports both horizontal and vertical
        orientations and allows for sorting the data in ascending or descending order.

        Parameters
        ----------
        data : pd.DataFrame
            The DataFrame containing the data to be plotted.
        variable_x_axis : str
            The column name in `data` to be used as the x-axis variable.
        variables_to_compare : list[str]
            A list of column names whose values are to be compared and visualized as stacked percentages.
        x_axis_label : Optional[str], optional
            The label for the x-axis. If None, `variable_x_axis` is used as the label. Default is None.
        y_axis_label : str, default "Percentage"
            The label for the y-axis. Default is "Percentage".
        title : Optional[str], optional
            The title of the chart. If None, a default title is generated based on the provided parameters.
            Default is None.
        chart_orientation : str, default "h"
            The orientation of the chart. "h" for horizontal, any other value for vertical. Default is "h".
        name_save_file : Optional[Path], optional
            The path (including filename) where the chart should be saved. If None, the chart is not saved.
            Default is None.
        sort_order : str, default "ascending"
            The order in which the bars are sorted. Can be "ascending" or "descending". Default is "ascending".

        Returns
        -------
        None
        """
        plt.figure(figsize=(10, 8))

        x_axis_label: str = x_axis_label if x_axis_label else variable_x_axis
        title: str = (
            title
            if title
            else f"Percentage Stacked Bar Chart {variable_x_axis} by {variables_to_compare}"
        )

        sum_data: pd.DataFrame = data.groupby(variable_x_axis)[
            variables_to_compare
        ].sum()

        data_percentages: pd.DataFrame = (
            sum_data.div(sum_data.sum(axis=1), axis=0) * 100
        )

        total_row: pd.Series = (
            data[variables_to_compare].sum() / data[variables_to_compare].sum().sum()
        ) * 100
        total_row.name = "Total"

        if sort_order == "ascending":
            data_percentages: pd.DataFrame = data_percentages.loc[
                data_percentages.sum(axis=1).sort_values().index
            ]
        elif sort_order == "descending":
            data_percentages: pd.DataFrame = data_percentages.loc[
                data_percentages.sum(axis=1).sort_values(ascending=False).index
            ]

        data_percentages: pd.DataFrame = pd.concat(
            [data_percentages, total_row.to_frame().T]
        )

        if chart_orientation == "h":
            data_percentages.plot(kind="barh", stacked=True)
            plt.xlabel(y_axis_label)
            plt.ylabel(x_axis_label)
        else:
            data_percentages.plot(kind="bar", stacked=True)
            plt.xlabel(x_axis_label)
            plt.ylabel(y_axis_label)

        plt.title(title)
        plt.legend(title="Variable", title_fontsize="13", loc="best")

        save_plot(self.save_plots, self.filepath, plt, name_save_file)


    def create_histogram(
        self,
        data: pd.DataFrame,
        variable: str,
        x_axis_limits: Optional[list[Union[int, float]]] = None,
        x_axis_logarithmic_scaling: bool = False,
        y_axis_logarithmic_scaling: bool = False,
        x_axis_label: Optional[str] = None,
        y_axis_label: Optional[str] = "Count",
        title: Optional[str] = None,
        name_save_file: Optional[Path] = None,
    ) -> None:

        """
        Creates a histogram from the given DataFrame.

        This method generates a histogram visualizing the distribution of values in a given column of the DataFrame.
        It supports optional logarithmic scaling of the x-axis, y-axis, or both to better visualize data with a wide
        distribution range.

        Parameters
        ----------
        data : pd.DataFrame
            The DataFrame containing the data to be plotted.
        variable : str
            The column name in `data` for which the histogram will be plotted.
        x_axis_limits : Optional[list[Union[int, float]]], optional
            A list containing two elements [min, max] that define the limits of the x-axis.
            If None, the limits are determined automatically.
        x_axis_logarithmic_scaling : bool, optional
            If True, applies logarithmic scaling to the x-axis. Useful for data with a wide range of values.
        y_axis_logarithmic_scaling : bool, optional
            If True, applies logarithmic scaling to the y-axis. Useful for data with a wide range of values.
        x_axis_label : str, optional
            The label for the x-axis.
        y_axis_label : str, optional
            The label for the y-axis. Defaults to "Count".
        title : Optional[str], optional
            The title of the plot. If None, defaults to "Histogram of <variable>".
        name_save_file : Optional[Path], optional
            The path (including filename) where the plot should be saved. If None, the plot is not saved.

        Returns
        -------
        None
        """
        plt.figure(figsize=(10, 8))

        x_axis_label: str = x_axis_label if x_axis_label else f"Number of {variable}"
        title: str = title if title else f"Histogram of {variable}"

        plt.hist(data[variable], bins=50)
        plt.xlabel(x_axis_label)
        plt.ylabel(y_axis_label)
        plt.title(title)

        if x_axis_limits:
            plt.xlim(x_axis_limits)

        if x_axis_logarithmic_scaling:
            plt.xscale("log")

        if y_axis_logarithmic_scaling:
            plt.yscale("log")

        save_plot(self.save_plots, self.filepath, plt, name_save_file)

    def create_count_distribution(
        self,
        data: pd.DataFrame,
        variable: str,
        x_axis_limits: Optional[list[Union[int, float]]] = None,
        x_axis_logarithmic_scaling: bool = False,
        y_axis_logarithmic_scaling: bool = False,
        x_axis_label: str = "Number of Occurances",
        y_axis_label: Optional[str] = None,
        title: Optional[str] = None,
        name_save_file: Optional[Path] = None,
    ) -> None:
        """
        Creates a distribution plot showing the frequency of unique values for a specified column in a DataFrame.

        This method visualizes how often unique values occur within a given column of the DataFrame. It can also apply
        logarithmic scaling to the x-axis, y-axis, or both to better visualize data with a wide distribution range.

        Parameters
        ----------
        data : pd.DataFrame
            The DataFrame containing the data to plot.
        variable : str
            The name of the column in `data` for which the distribution of unique values will be plotted.
        x_axis_limits : Optional[list[Union[int, float]]], optional
            A list containing two elements [min, max] that define the limits of the x-axis.
            If None, the limits are determined automatically.
        x_axis_logarithmic_scaling : bool, optional
            If True, applies logarithmic scaling to the x-axis. Useful for data with a wide range of values.
        y_axis_logarithmic_scaling : bool, optional
            If True, applies logarithmic scaling to the y-axis. Useful for data with a wide range of values.
        x_axis_label : str, optional
            The label for the x-axis. Defaults to "Number of Occurances".
        y_axis_label : Optional[str], optional
            The label for the y-axis. Defaults to "Number of <variable>" where
            <variable> is the column name specified in `variable`.
        title : Optional[str], optional
            The title of the plot. If None, defaults to "Distribution of Count over <variable>".
        name_save_file : Optional[Path], optional
            The file path (including the name) where the plot should be saved. If None, the plot is not saved.

        Returns
        -------
        None
        """
        y_axis_label: str = y_axis_label if y_axis_label else f"Number of {variable}"
        title: str = title if title else f"Distribution of Count over {variable}"

        plt.figure(figsize=(10, 6))

        uniques_per_column: pd.DataFrame = data[variable].value_counts()

        uniques_count: pd.Series = uniques_per_column.value_counts()

        uniques_count: pd.Series = uniques_count.sort_index()

        plt.bar(uniques_count.index, uniques_count.values)
        plt.xlabel(x_axis_label)
        plt.ylabel(y_axis_label)
        plt.title(title)

        if x_axis_limits:
            plt.xlim(x_axis_limits)

        if x_axis_logarithmic_scaling:
            plt.xscale("log")

        if y_axis_logarithmic_scaling:
            plt.yscale("log")

        save_plot(self.save_plots, self.filepath, plt, name_save_file)

    def create_grouped_histogram(
        self,
        data: pd.DataFrame,
        group_by: str,
        aggregation_column: str,
        aggregation_function: str,
        x_axis_label: Optional[str] = None,
        y_axis_label: Optional[str] = None,
        title: Optional[str] = None,
        name_save_file: Optional[Path] = None,
        bins: int = 10,
        edgecolor: str = "black",
    ):
        """
        This method creates a grouped histogram based on and aggregated function.

        Parameters
        ----------
        data : DataFrame
            The pandas DataFrame to create a histogram from.
        group_by : str
            The column in DataFrame to group by.
        aggregation_column : str
            The column to aggregate.
        aggregation_function : str
            The string passed to pandas agg function.
        x_axis_label : str, default data[group_by]
            The label for X-axis.
        y_axis_label : str, default 'Count'
            The label for Y-axis.
        title : str, default 'Histogram of {agg_func} {group_by} per {agg_col}'
            The title for the plot.
        name_save_file : str, default None
            The name to save the generated plot.
        bins : int, default 10
            The number of bins to divide the data into.
        edgecolor : str, default 'black'
            The edge color for the bins in the histogram.
        """

        agg_data: pd.DataFrame = data.groupby(group_by)[aggregation_column].agg(aggregation_function)

        plt.figure(figsize=(10, 8))

        x_axis_label: str = (
            x_axis_label if x_axis_label else f"{aggregation_function} {group_by}"
        )
        y_axis_label: str = (
            y_axis_label if y_axis_label else f"Number of {aggregation_column}"
        )
        title: str = (
            title
            if title
            else f"Histogram of {aggregation_function.title()} {aggregation_column.title()} per {group_by.title()}"
        )

        plt.hist(agg_data, bins=bins, edgecolor=edgecolor)
        plt.xlabel(x_axis_label)
        plt.ylabel(y_axis_label)
        plt.title(title)

        save_plot(self.save_plots, self.filepath, plt, name_save_file)

    def create_scatter_plot_simple(
        self,
        data: pd.DataFrame,
        variable_x_axis: str,
        variable_y_axis: str,
        title=None,
        hue=None,
        name_save_file: Optional[Path] = None,
    ) -> None:
        """
         Creates a simple scatter plot from the given DataFrame.

         This method generates a scatter plot visualizing the relationship between two variables. It supports
         optional grouping by color through the 'hue' parameter and saving the plot to a file.

         Parameters
         ----------
         data : pd.DataFrame
             The DataFrame containing the data to be plotted.
         variable_x_axis : str
             The column name in `data` to be used as the x-axis variable.
         variable_y_axis : str
             The column name in `data` to be used as the y-axis variable.
         title : str, optional
             The title of the plot. If None, a default title is generated based on the x and y variables.
         hue : str, optional
             The column name to be used for color encoding. This allows for the visualization of a third variable.
         name_save_file : Optional[Path], optional
             The path (including filename) where the plot should be saved. If None, the plot is not saved.

         Returns
         -------
         None
         """
        plt.figure(figsize=(10, 8))
        title: str = (
            title if title else f"Scatter Plot {variable_y_axis} vs {variable_x_axis}"
        )
        sns.scatterplot(data=data, x=variable_x_axis, y=variable_y_axis, hue=hue)
        plt.title(title)

        save_plot(self.save_plots, self.filepath, plt, name_save_file)

    def create_heatmap(
        self,
        data: pd.DataFrame,
        axis_variables: list[str],
        heat_variable: str,
        max_values_axes: list[Union[int, float]],
        min_values_axes: list[Union[int, float]],
        log_option: str,
        name_save_file: Path,
    ):
        """
        This method creates a heatmap.

        Parameters
        ----------
        data: pd.DataFrame
            The input dataframe which contains the data.
        axis_variables: list[str]
            The list of column names to be used as independent variables in the regression model.
        heat_variable: str
            The column name to be used as dependent variable in the regression model.
        max_values_axes: list
            The list of maximum values for each column defined in X.
            Any row with a value higher than this for the respective column is excluded.
        min_values_axes: list
            The list of minimum values for each column defined in X.
            Any row with a value less than this for the respective column is excluded.
        log_option: str
            Scale specific axis logarithmically. It could take four options: "false", "x_log", "y_log", "double_log".
            If an invalid value is supplied, the axes will not be log scaled.
        name_save_file: Path
            The name of the file to which the result will be saved (if self.save_result is True).

        Returns
        -------
        None
        """
        data_filtered: pd.DataFrame = data.copy()

        for col, max_val, min_val in zip(
            axis_variables, max_values_axes, min_values_axes
        ):
            data_filtered = data_filtered[
                (data_filtered[col] <= max_val) & (data_filtered[col] >= min_val)
            ]

        if log_option == "x_log":
            data_filtered[axis_variables] = np.log(data_filtered[axis_variables])
        elif log_option == "y_log":
            data_filtered[heat_variable] = np.log(data_filtered[heat_variable])
        elif log_option == "double_log":
            data_filtered[axis_variables] = np.log(data_filtered[axis_variables])
            data_filtered[heat_variable] = np.log(data_filtered[heat_variable])

        heatmap_data = data_filtered.pivot_table(
            index=axis_variables[0],
            columns=axis_variables[1],
            values=heat_variable,
            aggfunc="mean",
        )

        heatmap_data.sort_index(axis=0, ascending=False, inplace=True)

        heatmap_data = heatmap_data.interpolate(
            method="linear", limit_direction="both", axis=0
        )

        plt.figure(figsize=(10, 8))
        sns.heatmap(heatmap_data, cmap="YlGnBu")
        save_plot(self.save_plots, self.filepath, plt, name_save_file)

    def generate_surface_plot(
        self,
        function: FunctionData,
        x_axis_maximum: Union[int, float],
        y_axis_maximum: Union[int, float],
        x_axis_label: str,
        y_axis_label: str,
        z_axis_label: str,
        elevation_angle: Optional[Union[int, float]] = None,
        azimuth_angle: Optional[Union[int, float]] = None,
        title: Optional[str] = None,
        name_save_file: Optional[Path] = None,
        x_steps: int = 1000,
        y_steps: int = 1000,
    ) -> None:
        """
        Generates a 3D surface plot for a function.

        Parameters
        ----------
        function: FunctionData
            The function to plot as well as the parameters.
        x_axis_maximum: Union[int, float]
            The range of x values to plot.
        y_axis_maximum: Union[int, float]
            The range of y values to plot.
        x_axis_label: str
            The label for the x-axis.
        y_axis_label: str
            The label for the y-axis.
        z_axis_label: str
            The label for the z-axis.
        elevation_angle: int
            set elevation angle in the z plane
        azimuth_angle: int
            set azimuth angle in the x,y plane
        title: str, default None
            The title of the plot.
        name_save_file: str, default None
            The name of the file to which the result will be saved.
        x_steps: int
            The step size of the grid x direction
        y_steps: int
            The step size of the grid y direction
        """

        x_range: np.ndarray = np.linspace(1, x_axis_maximum, x_steps)
        y_range: np.ndarray = np.linspace(1, y_axis_maximum, y_steps)

        x, y = np.meshgrid(x_range, y_range)
        z = function(x, y)
        origin_z = function(1, 1)

        fig = plt.figure(figsize=(14, 8))
        ax = fig.add_subplot(111, projection="3d")

        if elevation_angle is not None:
            ax.view_init(elev=elevation_angle)

        if azimuth_angle is not None:
            ax.view_init(azim=azimuth_angle)

        ax.plot_surface(x, y, z)
        ax.set_xlabel(x_axis_label)
        ax.set_ylabel(y_axis_label)
        ax.set_zlabel(z_axis_label)

        title: str = title if title else f"3D surface plot of {function.func.__name__}"
        plt.title(title)

        plt.subplots_adjust(left=0.0, right=1.0, bottom=0.0, top=1.0)

        save_plot(self.save_plots, self.filepath, plt, name_save_file)

    def create_contour_plot(
        self,
        function: FunctionData,
        x_axis_maximum: Union[int, float],
        y_axis_maximum: Union[int, float],
        x_axis_label: str,
        y_axis_label: str,
        title: Optional[str] = None,
        name_save_file: Optional[Path] = None,
        x_steps: int = 1000,
        y_steps: int = 1000,
    ) -> None:
        """
        Generates a contour plot for a function.

        Parameters
        ----------
        function: FunctionData
            The function to plot as well as the parameters.
        x_axis_maximum: int
            The range of x values to plot.
        y_axis_maximum: int
            The range of y values to plot.
        x_axis_label: str
            The label for the x-axis.
        y_axis_label: str
            The label for the y-axis.
        title: str, default None
            The title of the plot.
        name_save_file: str, default None
            The name of the file to which the result will be saved.
        x_steps: int
            The step size of the grid x direction
        y_steps: int
            The step size of the grid y direction
        """
        x_range: np.ndarray = np.linspace(1, x_axis_maximum, x_steps)
        y_range: np.ndarray = np.linspace(1, y_axis_maximum, y_steps)

        x, y = np.meshgrid(x_range, y_range)
        z = function(x, y)

        fig = plt.figure()
        ax = fig.add_subplot(111)
        ax.contour(x, y, z)
        ax.set_xlabel(x_axis_label)
        ax.set_ylabel(y_axis_label)

        title: str = title if title else f"Contour plot of {function.func.__name__}"
        plt.title(title)

        save_plot(self.save_plots, self.filepath, plt, name_save_file)

    def create_density_plot(
        self,
        data: pd.DataFrame,
        variable_x_axis: str,
        variable_y_axis: str,
        data_breakpoints: list[Union[int, float]],
        name_save_file: Optional[Path] = None,
        title: Optional[str] = None,
        group_names: Optional[list[str]] = None,
        x_axis_limits: Optional[Tuple[float, float]] = None,
        y_axis_limits: Optional[Tuple[float, float]] = None,
    ) -> None:
        """
        This method creates a density plot.

        Parameters
        ----------
        data : pd.DataFrame
            DataFrame that has the data.
        variable_x_axis : str
            The name of the column of the provided DataFrame that this method applies breakpoints to.
        variable_y_axis : str
            The name of the column of the DataFrame that is plotted in the density plot.
        data_breakpoints : list[Union[int, float]]
            List of breakpoints or bins where the dataset is split.
        name_save_file : str, optional
            The name of the file to which the plot will be saved (if self.save_plots is True).
        title : str, optional
            The title of the plot.
        group_names : list[str], optional
            A list of group names used to rename the boolean categories in the legend.
        x_axis_limits : tuple[float, float], optional
            The lower and upper limits of the x-axis. Format as (lower_limit, upper_limit).
        y_axis_limits : tuple[float, float], optional
            The lower and upper limits of the y-axis. Format as (lower_limit, upper_limit).

        Returns
        -------
        None
        """
        data_temp: pd.DataFrame = data.copy()
        data_breakpoints: list[Union[int, float]] = sorted(data_breakpoints)

        data_temp["group"] = pd.cut(
            data_temp[variable_x_axis], bins=[-np.inf] + data_breakpoints + [np.inf]
        )

        fig, ax = plt.subplots(figsize=(10, 5))

        for group, values in data_temp.groupby("group", observed=False):
            sns.kdeplot(
                values[variable_y_axis],
                bw_adjust=0.5,
                clip_on=False,
                fill=True,
                alpha=0.5,
                linewidth=1.5,
                ax=ax,
                label=str(group),
            )

        plt.tight_layout()

        ax.legend()

        if title:
            ax.set_title(title)

        if group_names:
            legend = ax.legend_
            for t, l in zip(legend.texts, group_names):
                t.set_text(l)

        if x_axis_limits:
            ax.set_xlim(x_axis_limits)
        if y_axis_limits:
            ax.set_ylim(y_axis_limits)

        save_plot(self.save_plots, self.filepath, plt, name_save_file)

    def create_ridgeline_plot(
        self,
        data: pd.DataFrame,
        x_axis_variable: str,
        y_axis_variable: str,
        data_breakpoints: list[Union[int, float]],
        name_save_file: Optional[Path] = None,
        title: Optional[str] = None,
        group_names: Optional[list[str]] = None,
        x_axis_limits: Optional[Tuple[float, float]] = None,
        y_axis_limits: Optional[Tuple[float, float]] = None,
        y_label: Optional[str] = None,
    ) -> None:
        """
        This method creates a ridgeline plot.

        Parameters
        ----------
        data : pd.DataFrame
            DataFrame that has the data.
        x_axis_variable : str
            The name of the column of the provided DataFrame that this method applies breakpoints to.
        y_axis_variable : str
            The name of the column of the DataFrame that is plotted in the ridgeline plot.
        data_breakpoints : list[Union[int, float]]
            List of breakpoints or bins where the dataset is split.
        name_save_file : str, optional
            The name of the file to which the plot will be saved (if self.save_plots is True).
        title : str, optional
            The title of the plot.
        group_names : list[str], optional
            A list of group names used to rename the boolean categories in the legend.
        x_axis_limits : tuple[float, float], optional
            The lower and upper limits of the x-axis. Format as (lower_limit, upper_limit).
        y_axis_limits : tuple[float, float], optional
            The lower and upper limits of the y-axis. Format as (lower_limit, upper_limit).
        y_label : str, optional
            The label for the y-axis

        Returns
        -------
        None
        """
        data_temp: pd.DataFrame = data.copy()
        data_breakpoints: list[Union[int, float]] = sorted(data_breakpoints)

        data_temp["group"] = pd.cut(
            data_temp[x_axis_variable], bins=[-np.inf] + data_breakpoints + [np.inf]
        )

        g = sns.FacetGrid(
            data_temp, row="group", hue="group", aspect=15, height=2, palette="tab10"
        )
        g.map(
            sns.kdeplot,
            y_axis_variable,
            bw_adjust=0.5,
            clip_on=False,
            fill=True,
            alpha=1,
            linewidth=1.5,
        )
        g.map(
            sns.kdeplot, y_axis_variable, clip_on=False, color="w", lw=2, bw_adjust=0.5
        )
        g.map(plt.axhline, y=0, lw=2, clip_on=False)

        def label(x, color, label):
            ax = plt.gca()
            ax.text(
                0,
                0.2,
                label,
                fontweight="bold",
                fontsize=15,
                color=color,
                ha="left",
                va="center",
                transform=ax.transAxes,
            )

        g.map(label, y_axis_variable)

        g.fig.subplots_adjust(hspace=-0.25)

        g.set_titles("")
        g.set_ylabels("")
        g.set(yticks=[])
        g.despine(bottom=True, left=True)

        plt.tight_layout()

        if title:
            plt.title(title)

        if group_names:
            g.set(
                yticklabels=group_names[::-1]
            )  # Reversed because of the order how seaborn plots

        if y_axis_limits:
            g.set(ylim=y_axis_limits)
            g.map(plt.axhline, y=y_axis_limits[0], lw=2, clip_on=False)
        if x_axis_limits:
            g.set(xlim=x_axis_limits)

        plt.xlabel(y_axis_variable, fontsize=15)
        axes = g.axes.flatten()
        middle_plot = len(axes) // 2
        if y_label is None:
            axes[middle_plot].set_ylabel(x_axis_variable, fontsize=15)
        else:
            axes[middle_plot].set_ylabel(y_label, fontsize=15)
        plt.tick_params(axis="both", which="major", labelsize=15)
        plt.tick_params(axis="both", which="minor", labelsize=12)

        save_plot(self.save_plots, self.filepath, plt, name_save_file)

    def create_hexbin_plot(
        self,
        data: pd.DataFrame,
        variable_x_axis: str,
        variable_y_axis: str,
        x_axis_maximum: Union[int, float],
        y_axis_maximum: Union[int, float],
        trendline: bool = False,
        log_scale: bool = False,
        x_label: Optional[str] = None,
        y_label: Optional[str] = None,
        title: Optional[str] = None,
        name_save_file: Optional[Path] = None,
    ):
        """
        Generates a hexbin plot with an optional trendline.

        Parameters
        ----------
        data: DataFrame
            The DataFrame object that contains the data.

        variable_x_axis: str
            The column name in data for x values to plot.

        variable_y_axis: str
            The column name in data for y values to plot.
        x_axis_maximum: int
            The range maximum of x values to plot.
        y_axis_maximum: int
            The range maximum of y values to plot.
        trendline: bool, default False
            Whether to plot a trend line.
        log_scale: bool, default False
            Changes the scaling of the hue of hexbins to logarithmic.
        x_label: str, default None
            The label for the x-axis.
        y_label: str, default None
            The label for the y-axis.
        title: str, default None
            The title of the plot.
        name_save_file: str, default None
            The name of the file to which the result will be saved.
        """

        x_data: pd.DataFrame = data[variable_x_axis]
        y_data: pd.DataFrame = data[variable_y_axis]

        bins = None
        if log_scale:
            bins = "log"

        plt.hexbin(x_data, y_data, gridsize=50, cmap="Blues", bins=bins)

        data: pd.DataFrame = data.replace([np.inf, -np.inf], np.nan).dropna()

        x_data_clean, y_data_clean = data[variable_x_axis], data[variable_y_axis]

        if trendline:
            z = np.polyfit(x_data_clean, y_data_clean, 1)
            p = np.poly1d(z)
            plt.plot(x_data_clean, p(x_data_clean), "r--")

        if x_axis_maximum is not None:
            plt.xlim(0, x_axis_maximum)
        if y_axis_maximum is not None:
            plt.ylim(0, y_axis_maximum)

        if x_label:
            plt.xlabel(x_label)
        else:
            plt.xlabel(variable_x_axis)
        if y_label:
            plt.ylabel(y_label)
        else:
            plt.ylabel(variable_y_axis)
        if title:
            plt.title(title)

        save_plot(self.save_plots, self.filepath, plt, name_save_file)

    def create_box_plot(
        self,
        data: pd.DataFrame,
        variable_1: str,
        variable_2: str,
        x_axis_label: str,
        y_axis_label: str,
        title: Optional[str] = None,
        name_save_file: Optional[Path] = None,
    ):
        """
        Generates a box plot for two columns from a data frame.

        Parameters
        ----------
        data: dataframe
            The dataframe that contains the data.
        variable_1: str
            The name of the first column to plot.
        variable_2: str
            The name of the second column to plot.
        x_axis_label: str
            The label for the x-axis.
        y_axis_label: str
            The label for the y-axis.
        title: str, default None
            The title of the plot.
        name_save_file: str, default None
            The name of the file to which the result will be saved.
        """

        column_one_data: list = data[variable_1].dropna().values.tolist()
        column_two_data: list = data[variable_2].dropna().values.tolist()

        fig = plt.figure(figsize=(10, 6))
        ax = fig.add_subplot(111)
        ax.boxplot([column_one_data, column_two_data], labels=[variable_1, variable_2])
        ax.set_xlabel(x_axis_label)
        ax.set_ylabel(y_axis_label)

        title: str = title if title else f"Box plot of {variable_1} and {variable_2}"
        plt.title(title)

        save_plot(self.save_plots, self.filepath, plt, name_save_file)

    def create_violin_plot(
        self,
        data: pd.DataFrame,
        variable_x_axis: str,
        variable_y_axis: str,
        x_axis_label: str,
        y_axis_label: str,
        title: Optional[str] = None,
        name_save_file: Optional[Path] = None,
    ):
        """
        Generates a violin plot for two columns from a data frame.
        Violin plots show the distribution of data as a frequency distribution.
        Varying sample size between the columns will influence the violin plot.

        Parameters
        ----------
        data: dataframe
            The dataframe that contains the data.
        variable_x_axis: str
            The name of the first column to plot.
        variable_y_axis: str
            The name of the second column to plot.
        x_axis_label: str
            The label for the x-axis.
        y_axis_label: str
            The label for the y-axis.
        title: str, default None
            The title of the plot.
        name_save_file: str, default None
            The name of the file to which the result will be saved.
        """
        column_one_data: list = data[variable_x_axis].dropna().values.tolist()
        column_two_data: list = data[variable_y_axis].dropna().values.tolist()

        fig = plt.figure(figsize=(10, 6))
        ax = fig.add_subplot(111)
        ax.violinplot([column_one_data, column_two_data])
        ax.set_xticks([1, 2])
        ax.set_xticklabels([variable_x_axis, variable_y_axis])
        ax.set_xlabel(x_axis_label)
        ax.set_ylabel(y_axis_label)

        title: str = (
            title
            if title
            else f"Violin plot of {variable_x_axis} and {variable_y_axis}"
        )
        plt.title(title)

        save_plot(self.save_plots, self.filepath, plt, name_save_file)

    def create_forest_plot(
        self,
        regression_models: Dict[str, RegressionResults],
        coefficient_names: list[str],
        sort_by_size: bool = False,
        x_axis_minimum: Optional[Union[int, float]] = None,
        x_axis_maximum: Optional[Union[int, float]] = None,
        x_axis_label: Optional[str] = None,
        dotsize: float = 5,
        colors: list[str] = ["orange", "royalblue", "forestgreen", "firebrick"],
        name_save_file: Optional[Path] = None,
    ) -> None:
        """
        The function creates a forest plot of coefficients from multiple regression models.

        Parameters
        ----------
        regression_models: dict
            A dictionary containing category names as keys and RegressionResults as values.
        coefficient_names: list[str]
            The coefficients of the regression model to be plotted in the forest plot
        sort_by_size: bool
            If one want to sort the datapoints by order
            # Set x-axis limits if provided
        x_axis_minimum: float, default None
            The minimum value of the x-axis. Ignored if not given.
        x_axis_maximum: float, default None
            The maximum value of the x-axis.
        x_axis_label: str, default None
            The label of the x-axis
        dotsize: float, default 5
            The size of the dots indicating values in the forestplot
        colors: list[str]
            The colors used to plot the coefficients
        name_save_file: str, default None
            The name of the file to which the result will be saved.


        Returns
        -------
        None
        """
        data: dict = {"categories": []}

        for coefficient_name in coefficient_names:
            data[coefficient_name] = []
            data[f"conf_interval_{coefficient_name}"] = []
            data[f"color_{coefficient_name}"] = []

        for category, model in regression_models.items():
            data["categories"].append(category)
            for idx, coefficient_name in enumerate(coefficient_names):
                color = colors[idx % len(colors)]
                coef = model.params.get(coefficient_name, None)
                if coef is not None:
                    data[coefficient_name].append(coef)
                    conf_int = model.conf_int().loc[coefficient_name]
                    data[f"conf_interval_{coefficient_name}"].append(
                        [conf_int[0], conf_int[1]]
                    )
                    data[f"color_{coefficient_name}"].append(color)
                else:
                    data[coefficient_name].append(None)
                    data[f"conf_interval_{coefficient_name}"].append([None, None])
                    data[f"color_{coefficient_name}"].append(None)

        data_dataframe: pd.DataFrame = pd.DataFrame(data)

        if sort_by_size:
            data_dataframe.sort_values(by=coefficient_names, key=abs, inplace=True)
        data_dataframe.reset_index(drop=True, inplace=True)

        f, ax = plt.subplots(figsize=(7, 3.5 + 0.5 * len(coefficient_names)))
        ax.plot([0, 0], [0.5, len(data_dataframe) + 0.5], "--", color="black")

        for index, row in data_dataframe.iterrows():
            for idx, coefficient_name in enumerate(coefficient_names):
                color = row[f"color_{coefficient_name}"]
                coef = row[coefficient_name]
                conf_int = row[f"conf_interval_{coefficient_name}"]
                if color and coef is not None:
                    ax.plot(
                        conf_int,
                        [index + 1, index + 1],
                        "-",
                        color=color,
                        solid_capstyle="round",
                    )
                    ax.plot(coef, index + 1, "o", color=color, markersize=dotsize)

        ax.set_yticks(range(1, len(data_dataframe) + 1))
        ax.set_yticklabels(data_dataframe["categories"])
        ax.invert_yaxis()

        if x_axis_minimum is not None or x_axis_maximum is not None:
            ax.set_xlim(left=x_axis_minimum, right=x_axis_maximum)

        if x_axis_label is not None:
            ax.set_xlabel(x_axis_label)

        if len(coefficient_names) > 1:
            handles = [
                plt.Line2D(
                    [0], [0], color=colors[i % len(colors)], linewidth=2, linestyle="-"
                )
                for i in range(len(coefficient_names))
            ]
            ax.legend(
                handles,
                coefficient_names,
                title="Coefficients",
                loc="best",
                fontsize="small",
                title_fontsize="small",
            )

        save_plot(self.save_plots, self.filepath, plt, name_save_file)

    def create_forest_plot_paired_ttest(
        self,
        paired_ttests: dict[str, tuple[float, float, float]],
        x_axis_minimum: Optional[Union[int, float]] = None,
        x_axis_maximum: Optional[Union[int, float]] = None,
        x_axis_label: Optional[str] = None,
        dotsize: float = 5,
        name_save_file: Optional[Path] = None,
    ) -> None:
        """
        The function creates a forest plot of mean differences from multiple paired t-tests.

        Parameters
        ----------
        paired_ttests: dict
            A dictionary containing category names as keys and paired t-test results as values.
            Each test result is a tuple of (mean difference, lower CI, upper CI).
        x_axis_minimum: float, default None
            The minimum value of the x-axis. Ignored if not given.
        x_axis_maximum: float, default None
            The maximum value of the x-axis. Ignored if not given.
        x_axis_label: str, default None
            The label of the x-axis
        dotsize: float, default 5
            Size of the dots in the plot.
        name_save_file: str, default None
            The name of the file to which the result will be saved.

        Returns
        -------
        None
        """
        data: dict = {"categories": [], "mean_diff": [], "conf_interval": []}

        for category, (mean_diff, lower_ci, upper_ci) in paired_ttests.items():
            data["categories"].append(category)
            data["mean_diff"].append(mean_diff)
            data["conf_interval"].append([lower_ci, upper_ci])

        data_dataframe: pd.DataFrame = pd.DataFrame(data)

        f, ax = plt.subplots(figsize=(8, len(data_dataframe) * 0.5))

        ax.plot([0, 0], [0.5, len(data_dataframe) + 0.5], "--", color="black")

        for index, row in data_dataframe.iterrows():
            ax.plot(
                row["conf_interval"],
                [index + 1, index + 1],
                "-",
                color="orange",
                solid_capstyle="round",
            )
            ax.plot(row["mean_diff"], index + 1, "o", color="orange", markersize=dotsize)

        ax.set_yticks(range(1, len(data_dataframe) + 1))
        ax.set_yticklabels(data_dataframe["categories"])
        ax.invert_yaxis()

        if x_axis_minimum is not None or x_axis_maximum is not None:
            ax.set_xlim(left=x_axis_minimum, right=x_axis_maximum)

        if x_axis_label is not None:
            ax.set_xlabel(x_axis_label)

        save_plot(self.save_plots, self.filepath, plt, name_save_file)
