from statsmodels.regression.linear_model import RegressionResults

from src.analysis_functions.visualization import DataVisualizer
from src.data_classes.parameters_plot_barchart import BarChartPlotParameters
from src.data_classes.parameters_plot_boxplot import BoxPlotParameters
from src.data_classes.parameters_plot_contourplot import ContourPlotParameters
from src.data_classes.parameters_plot_count_distribution import (
    CountDistributionPlotParameters,
)
from src.data_classes.parameters_plot_densityplot import DensityPlotParameters
from src.data_classes.parameters_plot_forestplot import ForestPlotParameters
from src.data_classes.parameters_plot_forestplot_paired_ttest import (
    ForestPlotPairedTTestParameters,
)
from src.data_classes.parameters_plot_grouped_histogram import (
    GroupedHistogramParameters,
)
from src.data_classes.parameters_plot_heatmap import HeatmapParameters
from src.data_classes.parameters_plot_hexbinplot import HexbinPlotParameters
from src.data_classes.parameters_plot_histogram import HistogramPlotParameters
from src.data_classes.parameters_plot_percentage_stacked_barchart import (
    PercentageStackedBarChartPlotParameters,
)
from src.data_classes.parameters_plot_ridgelineplot import RidgelineParameters
from src.data_classes.parameters_plot_simple_scatterplot import (
    SimpleScatterPlotParameters,
)
from src.data_classes.parameters_plot_stacked_barchart import (
    StackedBarChartPlotParameters,
)
from src.data_classes.parameters_plot_surfaceplot import SurfacePlotParameters
from src.data_classes.parameters_plot_violinplot import ViolinPlotParameters
from src.data_classes.parameters_visualization import PlotParameters
from src.utils.helper_functions import FunctionData
from src.utils.helper_logging import log_visualization_details


def run_visualization(
    visualizer: DataVisualizer,
    analyses_list: list[PlotParameters],
    regression_results: dict[str, RegressionResults],
    functions: dict[str, FunctionData],
    ttest_results: dict[str, tuple[float, float, float]],
) -> None:
    """
    Run visualization analyses based on a list of parameter objects.

    This function iterates over a list of PlotParameters objects,
    logs the details of each analysis, and then performs the appropriate
    visualization based on the type of the PlotParameters object.

    Parameters
    ----------
    visualizer: DataVisualizer
        An instance of the DataVisualizer class.
    analyses_list: list[PlotParameters]
        A list of PlotParameters objects.
        Each object contains the parameters for a single visualization.
    regression_results: dict[str, RegressionResults]
        A dictionary mapping regression model names to their results.
    functions: dict[str, FunctionData]
        A dictionary mapping function names to their data.
    ttest_results: dict[str, tuple[float, float, float]]
        A dictionary mapping paired t-test names to their results.

    Raises
    ------
    ValueError: If the type of the PlotParameters object is not recognized.

    Returns
    -------
    None
    """
    for analysis_config in analyses_list:
        log_visualization_details(analysis_config)
        if isinstance(analysis_config, BarChartPlotParameters):
            visualizer.create_bar_chart(
                data=analysis_config.data,
                variable_x_axis=analysis_config.variable_x_axis,
                variable_y_axis=analysis_config.variable_y_axis,
                x_axis_label=analysis_config.x_axis_label,
                y_axis_label=analysis_config.y_axis_label,
                chart_orientation=analysis_config.chart_orientation,
                sort_order=analysis_config.sort_order,
                title=analysis_config.title,
                name_save_file=analysis_config.name_save_file,
                custom_order=analysis_config.custom_order,
            )
        elif isinstance(analysis_config, BoxPlotParameters):
            visualizer.create_box_plot(
                data=analysis_config.data,
                variable_1=analysis_config.variable_1,
                variable_2=analysis_config.variable_2,
                x_axis_label=analysis_config.x_axis_label,
                y_axis_label=analysis_config.y_axis_label,
                title=analysis_config.title,
                name_save_file=analysis_config.name_save_file,
            )
        elif isinstance(analysis_config, ContourPlotParameters):
            visualizer.create_contour_plot(
                function=functions[analysis_config.function_name],
                x_axis_maximum=analysis_config.x_axis_maximum,
                y_axis_maximum=analysis_config.y_axis_maximum,
                x_axis_label=analysis_config.x_axis_label,
                y_axis_label=analysis_config.y_axis_label,
                title=analysis_config.title,
                name_save_file=analysis_config.name_save_file,
            )
        elif isinstance(analysis_config, HistogramPlotParameters):
            visualizer.create_histogram(
                data=analysis_config.data,
                variable=analysis_config.variable,
                x_axis_limits=analysis_config.x_axis_limits,
                x_axis_logarithmic_scaling=analysis_config.x_axis_logarithmic_scaling,
                y_axis_logarithmic_scaling=analysis_config.y_axis_logarithmic_scaling,
                x_axis_label=analysis_config.x_axis_label,
                y_axis_label=analysis_config.y_axis_label,
                title=analysis_config.title,
                name_save_file=analysis_config.name_save_file,
            )
        elif isinstance(analysis_config, CountDistributionPlotParameters):
            visualizer.create_count_distribution(
                data=analysis_config.data,
                variable=analysis_config.variable,
                x_axis_limits=analysis_config.x_axis_limits,
                x_axis_logarithmic_scaling=analysis_config.x_axis_logarithmic_scaling,
                y_axis_logarithmic_scaling=analysis_config.y_axis_logarithmic_scaling,
                x_axis_label=analysis_config.x_axis_label,
                y_axis_label=analysis_config.y_axis_label,
                title=analysis_config.title,
                name_save_file=analysis_config.name_save_file,
            )
        elif isinstance(analysis_config, DensityPlotParameters):
            visualizer.create_density_plot(
                data=analysis_config.data,
                variable_x_axis=analysis_config.variable_x_axis,
                variable_y_axis=analysis_config.variable_y_axis,
                data_breakpoints=analysis_config.data_breakpoints,
                title=analysis_config.title,
                name_save_file=analysis_config.name_save_file,
            )
        elif isinstance(analysis_config, ForestPlotParameters):
            visualizer.create_forest_plot(
                regression_models=_get_labeled_models_for_forest_plot(
                    analysis_config, regression_results
                ),
                coefficient_names=analysis_config.coefficient_names,
                sort_by_size=analysis_config.sort_by_size,
                x_axis_minimum=analysis_config.x_axis_minimum,
                x_axis_maximum=analysis_config.x_axis_maximum,
                dotsize=analysis_config.dotsize,
                colors=analysis_config.colors,
                name_save_file=analysis_config.name_save_file,
            )
        elif isinstance(analysis_config, ForestPlotPairedTTestParameters):
            visualizer.create_forest_plot_paired_ttest(
                paired_ttests=_get_labeled_ttests_for_forest_plot(
                    analysis_config, ttest_results
                ),
                x_axis_minimum=analysis_config.x_axis_minimum,
                x_axis_maximum=analysis_config.x_axis_maximum,
                dotsize=analysis_config.dotsize,
                name_save_file=analysis_config.name_save_file,
            )
        elif isinstance(analysis_config, GroupedHistogramParameters):
            visualizer.create_grouped_histogram(
                data=analysis_config.data,
                group_by=analysis_config.group_by,
                aggregation_column=analysis_config.aggregation_variable,
                aggregation_function=analysis_config.aggregation_function,
                x_axis_label=analysis_config.x_axis_label,
                y_axis_label=analysis_config.y_axis_label,
                title=analysis_config.title,
                name_save_file=analysis_config.name_save_file,
            )
        elif isinstance(analysis_config, HeatmapParameters):
            visualizer.create_heatmap(
                data=analysis_config.data,
                axis_variables=analysis_config.axis_variables,
                heat_variable=analysis_config.heat_variable,
                max_values_axes=analysis_config.axis_maxima,
                min_values_axes=analysis_config.axis_minima,
                log_option=analysis_config.logarithmic_heat_scaling,
                name_save_file=analysis_config.name_save_file,
            )
        elif isinstance(analysis_config, HexbinPlotParameters):
            visualizer.create_hexbin_plot(
                data=analysis_config.data,
                variable_x_axis=analysis_config.variable_x_axis,
                variable_y_axis=analysis_config.variable_y_axis,
                x_axis_maximum=analysis_config.x_axis_maximum,
                y_axis_maximum=analysis_config.y_axis_maximum,
                trendline=analysis_config.trendline,
                log_scale=analysis_config.logarithmic_hex_scaling,
                title=analysis_config.title,
                name_save_file=analysis_config.name_save_file,
            )
        elif isinstance(analysis_config, PercentageStackedBarChartPlotParameters):
            visualizer.create_percentage_stacked_bar_chart(
                data=analysis_config.data,
                variable_x_axis=analysis_config.variable_x_axis,
                variables_to_compare=analysis_config.variables_to_compare,
                chart_orientation=analysis_config.chart_orientation,
                sort_order=analysis_config.sort_order,
                title=analysis_config.title,
                name_save_file=analysis_config.name_save_file,
            )
        elif isinstance(analysis_config, RidgelineParameters):
            visualizer.create_ridgeline_plot(
                data=analysis_config.data,
                x_axis_variable=analysis_config.variable_x_axis,
                y_axis_variable=analysis_config.variable_y_axis,
                data_breakpoints=analysis_config.data_breakpoints,
                title=analysis_config.title,
                name_save_file=analysis_config.name_save_file,
            )
        elif isinstance(analysis_config, StackedBarChartPlotParameters):
            visualizer.create_stacked_bar_chart(
                data=analysis_config.data,
                variable_x_axis=analysis_config.variable_x_axis,
                variable_y_axis=analysis_config.variable_y_axis,
                x_axis_label=analysis_config.x_axis_label,
                y_axis_label=analysis_config.y_axis_label,
                hue=analysis_config.hue,
                chart_orientation=analysis_config.chart_orientation,
                sort_order=analysis_config.sort_order,
                title=analysis_config.title,
                name_save_file=analysis_config.name_save_file,
                custom_order=analysis_config.custom_order,
            )
        elif isinstance(analysis_config, SimpleScatterPlotParameters):
            visualizer.create_scatter_plot_simple(
                data=analysis_config.data,
                variable_x_axis=analysis_config.variable_x_axis,
                variable_y_axis=analysis_config.variable_y_axis,
                title=analysis_config.title,
                name_save_file=analysis_config.name_save_file,
            )
        elif isinstance(analysis_config, SurfacePlotParameters):
            visualizer.generate_surface_plot(
                function=functions[analysis_config.function_name],
                x_axis_maximum=analysis_config.x_axis_maximum,
                y_axis_maximum=analysis_config.y_axis_maximum,
                x_axis_label=analysis_config.x_axis_label,
                y_axis_label=analysis_config.y_axis_label,
                z_axis_label=analysis_config.z_axis_label,
                elevation_angle=analysis_config.elevation_angle,
                azimuth_angle=analysis_config.azimuth_angle,
                title=analysis_config.title,
                name_save_file=analysis_config.name_save_file,
            )
        elif isinstance(analysis_config, ViolinPlotParameters):
            visualizer.create_violin_plot(
                data=analysis_config.data,
                variable_x_axis=analysis_config.variable_x_axis,
                variable_y_axis=analysis_config.variable_y_axis,
                x_axis_label=analysis_config.x_axis_label,
                y_axis_label=analysis_config.y_axis_label,
                title=analysis_config.title,
                name_save_file=analysis_config.name_save_file,
            )
        else:
            raise ValueError(
                f"Invalid type of visualization requested {type(analysis_config)}"
            )


def _get_labeled_models_for_forest_plot(
    analysis_config: ForestPlotParameters,
    regression_results: dict[str, RegressionResults],
) -> dict:
    """
    Maps regression model names to their labels for forest plot visualization.

    This function takes an analysis configuration for a forest plot and a dictionary
    of regression results. It maps the specified regression model names to their
    corresponding labels as defined in the analysis configuration. This mapping is
    used for labeling the models in the forest plot visualization.

    Parameters
    ----------
    analysis_config : ForestPlotParameters
        The configuration parameters for the forest plot, including the names and
        labels of the regression models to be used.
    regression_results : dict[str, RegressionResults]
        A dictionary mapping model names to their corresponding fitted regression models.

    Returns
    -------
    dict
        A dictionary mapping the labels (as specified in the analysis configuration)
        to the corresponding regression models.
    """
    mapping_regression_model_names_to_labels: dict = {
        name: label
        for name, label in zip(
            analysis_config.regression_model_names,
            analysis_config.regression_model_labels,
        )
        if name in regression_results.keys()
    }

    selected_models: dict = {
        name: regression_results[name]
        for name in mapping_regression_model_names_to_labels.keys()
    }

    labeled_models: dict = {
        mapping_regression_model_names_to_labels[name]: model
        for name, model in selected_models.items()
    }

    return labeled_models


def _get_labeled_ttests_for_forest_plot(
    analysis_config: ForestPlotPairedTTestParameters,
    ttest_results: dict[str, tuple[float, float, float]],
) -> dict:
    """
    Maps paired t-test names to their labels for forest plot visualization.

    This function takes an analysis configuration for a forest plot that includes
    paired t-tests and a dictionary of t-test results. It maps the specified t-test
    names to their corresponding labels as defined in the analysis configuration.
    This mapping is used for labeling the t-tests in the forest plot visualization.

    Parameters
    ----------
    analysis_config : ForestPlotPairedTTestParameters
        The configuration parameters for the forest plot, including the names and
        labels of the paired t-tests to be used.
    ttest_results : dict[str, tuple[float, float, float]]
        A dictionary mapping t-test names to their results (t-statistic, p-value, and
        degrees of freedom).

    Returns
    -------
    dict
        A dictionary mapping the labels (as specified in the analysis configuration)
        to the corresponding t-test results.
    """
    mapping_ttest_names_to_labels: dict = {
        name: label
        for name, label in zip(
            analysis_config.paired_ttest_names, analysis_config.paired_ttest_labels
        )
        if name in ttest_results.keys()
    }

    selected_ttests: dict = {
        name: ttest_results[name] for name in mapping_ttest_names_to_labels.keys()
    }

    labeled_ttests: dict = {
        mapping_ttest_names_to_labels[name]: ttest
        for name, ttest in selected_ttests.items()
    }

    return labeled_ttests
