import matplotlib.pyplot as plt
import numpy as np
from typing import Dict, List, Any, Optional, Callable, Tuple, Literal
from dataclasses import dataclass, field

from aboba.experiment.experiment_data import ExperimentData
from aboba.utils.draw import (
    draw_interval,
    draw_pvalue_distribution,
    draw_ecdf,
)
from aboba.utils.alpha_interval import calculate_real_alpha
from aboba.utils.translations import Language, t


@dataclass
class GroupVisualizationConfig:
    color: Optional[str] = None
    kwargs: Dict[str, Any] = field(default_factory=dict)


def _draw_group_aa_layout_2x2(
    group_name: str,
    data: ExperimentData,
    ax_interval: plt.Axes,
    ax_hist: plt.Axes,
    ax_ecdf: plt.Axes,
    alpha: float,
    lang: Language = "en",
):
    """
    Internal helper function to draw the standardized 3-panel layout for a single group.

    This function orchestrates the visualization of three key plots for a given
    experiment group:
    1. Confidence interval for the Type I error rate (alpha).
    2. Histogram of p-values.
    3. Empirical Cumulative Distribution Function (ECDF) of p-values.

    This layout is designed for A/A testing diagnostics but can be used for
    other scenarios where such a combination of plots is insightful.

    Parameters
    ----------
    group_name : str
        Name of the experiment group being visualized.
    data : ExperimentData
        Object containing the history of test results (p-values, etc.) for the group.
    ax_interval : plt.Axes
        Axes object for plotting the confidence interval.
    ax_hist : plt.Axes
        Axes object for plotting the p-value histogram.
    ax_ecdf : plt.Axes
        Axes object for plotting the p-value ECDF.
    alpha : float
        The significance level used for determining rejections (p < alpha).
    lang : Language, optional
        Language code for labels ('en' or 'ru'), by default 'en'.

    Returns
    -------
    None
    """
    n_iter = len(data.history)
    n_errors = sum(int(tr.pvalue < alpha) for tr in data.history)
    real_alpha, left_alpha, right_alpha = calculate_real_alpha(n_iter=n_iter, n_errors=n_errors)
    pvals = [tr.pvalue for tr in data.history]

    # Top: interval (full width)
    draw_interval(real_alpha, left_alpha, right_alpha, ax_interval, alpha, group_name, lang=lang)

    # Bottom-left: histogram
    draw_pvalue_distribution(pvals, ax_hist, group_name, lang=lang)

    # Bottom-right: ECDF
    draw_ecdf(pvals, ax_ecdf, group_name, lang=lang)


def draw_aa_experiment_layout(
    groups: Dict[str, ExperimentData],
    alpha: float = 0.05,
    experiment_name: str = "AB experiment",
    groups_list: Optional[List[str]] = None,
    group_configs: Optional[Dict[str, GroupVisualizationConfig]] = None,
    lang: Language = "en",
    figsize: Optional[Tuple[float, float]] = None,
    filter_empty: bool = True,
) -> Tuple[Optional[plt.Figure], Any]:
    """
    Generate a standardized 3-panel visualization layout for one or more experiment groups.

    Each group is represented by a 2x2 grid of subplots:
    - Top row (spans 2 columns): Confidence interval for the Type I error rate (alpha).
    - Bottom row: Left subplot - histogram of p-values.
                 Right subplot - ECDF of p-values.

    This layout is optimized for A/A testing analysis, allowing for easy comparison
    of multiple groups against the expected null behavior (uniform p-values, correct alpha).

    Parameters
    ----------
    groups : Dict[str, ExperimentData]
        A dictionary mapping group names to their corresponding ExperimentData objects.
    alpha : float, optional
        The significance level used for determining rejections (p < alpha), by default 0.05.
    experiment_name : str, optional
        Title for the overall figure, by default "AB experiment".
    groups_list : List[str], optional
        Specific list of group names to visualize. If None, all groups in `groups` are used.
    group_configs : Dict[str, GroupVisualizationConfig], optional
        Per-group configuration options (e.g., custom colors).
    lang : Language, optional
        Language code for labels ('en' or 'ru'), by default 'en'.
    figsize : Tuple[float, float], optional
        Size of the figure (width, height). If None, a default size is calculated
        based on the number of groups.
    filter_empty : bool, optional
        Whether to exclude groups that have no data (empty history), by default True.

    Returns
    -------
    Tuple[Optional[plt.Figure], Any]
        A tuple containing:
        - fig (plt.Figure or None): The matplotlib figure object. None if no valid groups.
        - axes (Any): An array-like structure of the individual subplot axes.
    """
    assert 0.0 < alpha < 1.0

    group_configs = group_configs or {}
    groups_to_draw = groups_list or list(groups.keys())
    if filter_empty:
        groups_to_draw = [g for g in groups_to_draw if not groups[g].is_empty()]
    if not groups_to_draw:
        print("No non-empty groups to visualize")
        return None, None

    n_groups = len(groups_to_draw)
    if figsize is None:
        figsize = (12.5, 6.0 * n_groups) 

    fig = plt.figure(figsize=figsize)

    height_ratios = []
    for _ in range(n_groups):
        height_ratios.extend([1, 3])
    gs = fig.add_gridspec(
        nrows=2 * n_groups,
        ncols=2,
        height_ratios=height_ratios,
        width_ratios=[2, 1], 
        hspace=0.5,
        wspace=0.3
    )

    axes_list = []
    for i, group_name in enumerate(groups_to_draw):
        data = groups[group_name]
        config = group_configs.get(group_name, GroupVisualizationConfig())

        row_top = 2 * i
        row_bottom = 2 * i + 1

        ax_interval = fig.add_subplot(gs[row_top, :])      # full width
        ax_hist = fig.add_subplot(gs[row_bottom, 0])       # left
        ax_ecdf = fig.add_subplot(gs[row_bottom, 1])       # right

        _draw_group_aa_layout_2x2(
            group_name, data, ax_interval, ax_hist, ax_ecdf, alpha, lang=lang
        )
        axes_list.append([ax_interval, ax_hist, ax_ecdf])

    fig.suptitle(experiment_name, fontsize=14, fontweight='bold')
    if n_groups == 1:
        fig.subplots_adjust(top=0.90, hspace=0.25)
    else:
        fig.subplots_adjust(top=0.92, hspace=0.35)
    # fig.subplots_adjust(
    #     top=0.92,    
    #     bottom=0.08,  
    #     left=0.08,
    #     right=0.95,
    #     hspace=0.25 if n_groups == 1 else 0.35,   
    #     wspace=0.25   
    # )

    return fig, np.array(axes_list, dtype=object)


def default_visualization_method(
    groups: Dict[str, ExperimentData],
    alpha: float = 0.05,
    experiment_name: str = "AB experiment",
    lang: Language = "en",
    figsize: Optional[Tuple[float, float]] = None,
    filter_empty: bool = True,
    **kwargs
):
    """
    Default visualization method, serving as an entry point to the standardized layout.

    This function provides a simple interface to call the `draw_aa_experiment_layout`
    function, which implements the standard 3-panel per group visualization.
    It is designed to be the default handler for the `AbobaExperiment.draw()` method.

    Parameters
    ----------
    groups : Dict[str, ExperimentData]
        A dictionary mapping group names to their corresponding ExperimentData objects.
    alpha : float, optional
        The significance level used for determining rejections (p < alpha), by default 0.05.
    experiment_name : str, optional
        Title for the overall figure, by default "AB experiment".
    lang : Language, optional
        Language code for labels ('en' or 'ru'), by default 'en'.
    figsize : Tuple[float, float], optional
        Size of the figure (width, height). If None, a default size is calculated
        based on the number of groups.
    filter_empty : bool, optional
        Whether to exclude groups that have no data (empty history), by default True.
    **kwargs : dict
        Additional keyword arguments passed through to `draw_aa_experiment_layout`.

    Returns
    -------
    Tuple[Optional[plt.Figure], Any]
        A tuple containing:
        - fig (plt.Figure or None): The matplotlib figure object. None if no valid groups.
        - axes (Any): An array-like structure of the individual subplot axes.
    """
    return draw_aa_experiment_layout(
        groups=groups,
        alpha=alpha,
        experiment_name=experiment_name,
        lang=lang,
        figsize=figsize,
        filter_empty=filter_empty,
        **kwargs
    )