from typing import Any, Callable, Optional, List, Tuple
import matplotlib.pyplot as plt
import numpy as np
import seaborn as sns

from aboba.utils.translations import t, Language
from tqdm import tqdm

sns.set(font_scale=1.3)


def draw_interval(
    real_alpha: float,
    left_alpha: float,
    right_alpha: float,
    axes: plt.Axes,
    alpha: float = 0.05,
    name: Optional[str] = None,
    lang: Language = "en",
):
    """
    Visualize the confidence interval for the estimated Type I error rate (alpha).

    This function draws a horizontal line representing the theoretical alpha level,
    fills the confidence interval around the empirical alpha, and marks the
    observed empirical alpha with a star. It is typically used in A/A testing
    to assess whether the statistical test maintains the nominal error rate.

    Parameters
    ----------
    real_alpha : float
        The empirically observed Type I error rate (proportion of rejections
        when null hypothesis is true).
    left_alpha : float
        Left boundary of the confidence interval for `real_alpha`.
    right_alpha : float
        Right boundary of the confidence interval for `real_alpha`.
    axes : plt.Axes
        Matplotlib Axes object to draw the plot on.
    alpha : float, optional
        The theoretical (nominal) significance level, by default 0.05.
    name : str, optional
        Name of the experiment or group, used as part of the plot title.
        If None, only the rejection rate is shown in the title.
    lang : Language, optional
        Language code for labels ('en' or 'ru'), by default 'en'.

    Returns
    -------
    None
    """
    with sns.axes_style("whitegrid"):
        # Set background grid color to light gray (like in screenshot)
        axes.grid(True)

        # Horizontal line (baseline)
        axes.hlines(0, 0, 1, color="black", lw=2, alpha=0.6)

        # Theoretical alpha line (vertical dashed red)
        axes.vlines(alpha, -0.2, 0.2, color="red", lw=3, linestyle="--", alpha=0.9)

        # Confidence interval (green fill)
        axes.fill_between(
            [left_alpha, right_alpha],
            [-0.1], [0.1],
            color="#20b2aa",
            alpha=0.4
        )

        # Star marker for empirical alpha
        axes.scatter(real_alpha, 0, s=300, marker="*", color="red")

        # X-axis limits
        margin = max(0.002, 0.05 * (right_alpha - left_alpha))
        axes.set_xlim((min(alpha, left_alpha) - margin, max(alpha, right_alpha) + margin))

        # Title
        if name is not None:
            title = f"{name} | {t('rejections', lang)} = {100 * real_alpha:.2f}%, ({100 * left_alpha:.2f}%, {100 * right_alpha:.2f}%)"
        else:
            title = f"{t('rejections', lang).capitalize()} = {100 * real_alpha:.2f}%, ({100 * left_alpha:.2f}%, {100 * right_alpha:.2f}%)"
        
        axes.set_title(title, fontsize=12, pad=4, fontweight='bold')


        axes.set_ylim((-0.25, 0.25))
        axes.set_yticks([])
        axes.set_xlabel("")
        axes.set_ylabel("")


def draw_pvalue_distribution(
    pvals: List[float],
    axes: plt.Axes,
    name: Optional[str] = None,
    lang: Language = "en",
):
    """
    Plot a histogram of p-values to assess their distribution.

    Under the null hypothesis (e.g., in an A/A test), p-values should be
    uniformly distributed between 0 and 1. This plot helps diagnose issues
    like p-hacking, selection bias, or violations of test assumptions.
    A flat histogram around the density of 1.0 indicates a well-calibrated test.

    Parameters
    ----------
    pvals : List[float]
        A list of p-values obtained from repeated statistical tests.
    axes : plt.Axes
        Matplotlib Axes object to draw the plot on.
    name : str, optional
        Name of the experiment or group, used as part of the plot title.
        If None, a generic title is used.
    lang : Language, optional
        Language code for labels ('en' or 'ru'), by default 'en'.

    Returns
    -------
    None
    """
    with sns.axes_style("whitegrid"):
        axes.grid(True)

        n, bins, patches = axes.hist(
            pvals,
            bins=50,  
            density=True,  
            alpha=0.4,
            color="#ec7c26",
            edgecolor="navy",
            linewidth=0.3,
        )

        axes.axhline(1.0, color="red", linestyle="--", linewidth=2, alpha=0.8)

        # Labels
        axes.set_xlabel(t('pvalue', lang), fontsize=10)
        axes.set_ylabel(t('density', lang), fontsize=10)
        title = t('pvalue_distribution', lang)
        if name:
            title = f"{name} | {title}"
        axes.set_title(title, fontsize=12, pad=4)

        # X-axis ticks every 0.2
        axes.set_xticks(np.arange(0, 1.1, 0.2))


def draw_ecdf(
    pvals: List[float],
    axes: plt.Axes,
    name: Optional[str] = None,
    lang: Language = "en",
):
    """
     Plot the Empirical Cumulative Distribution Function (ECDF) of p-values.

    The ECDF is another way to visualize the distribution of p-values.
    Under the null hypothesis, the ECDF should follow the diagonal line y=x.
    Deviations from this line can indicate departures from uniformity,
    suggesting that the null hypothesis may be false for some tests or
    that the test statistic is miscalibrated.

    Parameters
    ----------
    pvals : List[float]
        A list of p-values obtained from repeated statistical tests.
    axes : plt.Axes
        Matplotlib Axes object to draw the plot on.
    name : str, optional
        Name of the experiment or group, used as part of the plot title.
        If None, a generic title is used.
    lang : Language, optional
        Language code for labels ('en' or 'ru'), by default 'en'.

    Returns
    -------
    None
    """
    with sns.axes_style("darkgrid"):
        axes.grid(True)


        sorted_pvalues = np.sort(pvals)
        n = len(sorted_pvalues)
        empirical_cdf = np.arange(1, n + 1) / n

        # ECDF line: purple
        axes.plot(
            sorted_pvalues,
            empirical_cdf,
            linewidth=3,
            color="#8b00ff",
            label=t('ecdf', lang),
        )

        # Diagonal (H0): red dashed
        axes.plot(
            [0, 1],
            [0, 1],
            color="red",
            linestyle="--",
            alpha=0.9,
            linewidth=2,
            label=r"$H_0$",
        )

        # Labels
        axes.set_xlabel(t('pvalue', lang), fontsize=10)
        axes.set_ylabel(t('empirical_cdf', lang), fontsize=10)
        axes.set_title(t('ecdf_of_pvalues', lang), fontsize=12, pad=4)

        # Limits
        axes.set_xlim(0, 1)
        axes.set_ylim(0, 1)

        # Legend
        axes.legend(loc='lower right', fontsize=9, framealpha=0.9, fancybox=True, edgecolor='gray')

        # Title with group name
        if name is not None:
            axes.set_title(f"{name} | {t('ecdf_of_pvalues', lang)}", fontsize=12, pad=10)


def draw_comparison_scatter(
    group1_pvals: List[float],
    group2_pvals: List[float],
    axes: plt.Axes,
    group1_name: str = "Group 1",
    group2_name: str = "Group 2",
    alpha: float = 0.05,
    lang: Language = "en",
):
    """
    Scatter plot — kept for compatibility, but not used in main layout.
    """
    with sns.axes_style("whitegrid"):
        axes.patch.set_facecolor('#f5f5f5')
        axes.grid(True, alpha=0.3, linestyle='-', linewidth=0.5, color='#cccccc')

        min_len = min(len(group1_pvals), len(group2_pvals))
        g1, g2 = np.array(group1_pvals[:min_len]), np.array(group2_pvals[:min_len])

        both_sig = (g1 < alpha) & (g2 < alpha)
        only_g1_sig = (g1 < alpha) & (g2 >= alpha)
        only_g2_sig = (g1 >= alpha) & (g2 < alpha)
        neither_sig = (g1 >= alpha) & (g2 >= alpha)

        # Labels (localized)
        neither_label = t("neither_significant", lang)
        both_label = t("both_significant", lang)
        only1_label = t("only_group1", lang).format(group1_name)
        only2_label = t("only_group2", lang).format(group2_name)

        axes.scatter(g1[neither_sig], g2[neither_sig], alpha=0.5, color='gray', label=neither_label)
        axes.scatter(g1[both_sig], g2[both_sig], alpha=0.5, color='green', label=both_label)
        axes.scatter(g1[only_g1_sig], g2[only_g1_sig], alpha=0.5, color='blue', label=only1_label)
        axes.scatter(g1[only_g2_sig], g2[only_g2_sig], alpha=0.5, color='red', label=only2_label)

        axes.axhline(alpha, color='red', linestyle='--', alpha=0.5)
        axes.axvline(alpha, color='red', linestyle='--', alpha=0.5)
        axes.plot([0, 1], [0, 1], 'k--', alpha=0.3, label='y=x')

        axes.set_xlabel(f'{group1_name} {t("pvalue", lang)}')
        axes.set_ylabel(f'{group2_name} {t("pvalue", lang)}')
        axes.set_title(f'{group1_name} vs {group2_name}', fontsize=12)
        axes.set_xlim(-0.05, 1.05)
        axes.set_ylim(-0.05, 1.05)
        axes.legend(loc='best', fontsize=9, framealpha=0.9)
        

def draw_power_curve(
    effect_sizes: List[float],
    powers: List[float],
    ci_lower: List[float],
    ci_upper: List[float],
    figsize,
    target_power: Optional[float] = None,
    alpha: float = 0.05,
    lang: Language = "en",
):
    """
    Draw power curve: power vs effect size.

    Args:
        effect_sizes (List[float]): x-axis values (e.g., [0.0, 0.02, ..., 0.25])
        powers (List[float]): estimated power for each effect size
        ci_lower (List[float]): lower bound of 95% CI for power
        ci_upper (List[float]): upper bound of 95% CI for power
        target_power (float, optional): horizontal line for desired power (e.g., 0.8)
        alpha (float): significance level (used for legend only)
        axes (plt.Axes): matplotlib axes to plot on
        lang (str): 'en' or 'ru'
    """
    
    plt.figure(figsize=figsize)

    with sns.axes_style("whitegrid"):
        plt.grid(True)

        plt.plot(
            effect_sizes,
            powers,
            color="#8b00ff",
            linewidth=2,
            label=t('power', lang),
        )

        plt.fill_between(
            effect_sizes,
            ci_lower,
            ci_upper,
            color="#8b00ff",
            alpha=0.2,
            label=t('power_ci', lang),
        )

        if target_power is not None:
            plt.axhline(
                target_power,
                color="#20b2aa",
                linestyle="--",
                linewidth=2,
                label=t('target_power', lang),
            )

        plt.axhline(
            alpha,
            color="#ff2400",
            linestyle="--",
            linewidth=2,
            label=t('significance_level', lang),
        )

        plt.xlabel(t('effect_size', lang))
        plt.ylabel(t('power', lang))

        plt.xlim(left=0, right=max(effect_sizes) * 1.03)
        plt.ylim(bottom=0, top=1.03)

        plt.legend(loc='lower right')
        plt.title(t('power_curve_plot', lang), fontsize=14, fontweight='bold')
        
        plt.show()