from typing import List, Tuple, Optional
from tqdm.auto import tqdm
import pandas as pd
import numpy as np
from aboba.base.base_test import BaseTest, TestResult
from aboba.pipeline import Pipeline
from aboba.base import EffectModifier
from statsmodels.stats.proportion import proportion_confint


def simulate_power_for_effect(
    test: BaseTest,
    data: pd.DataFrame,
    pipeline: Pipeline,
    effect_modifier: Optional[EffectModifier],
    n_iter: int,
    alpha: float,
    effect_size: float,
    progress_desc: str = "Simulating"
) -> Tuple[int, List[float]]:
    """
    Simulates the experiment `n_iter` times for a given `effect_size` to estimate power.

    Uses the provided `test`, `data`, and `pipeline`. Applies the `effect_modifier`
    configured for the given `effect_size` during each simulation iteration.

    Parameters
    ----------
    test : BaseTest
        The statistical test to use (e.g., AbsoluteIndependentTTest).
    data : pd.DataFrame
        The source data to sample from using the pipeline.
    pipeline : Pipeline
        The pipeline to sample and potentially preprocess data.
    effect_modifier : EffectModifier, optional
        The modifier to apply the effect. If None, no effect is applied (useful for alpha).
        This modifier should be pre-configured to apply `effect_size` to the correct group/column.
    n_iter : int
        Number of simulation iterations.
    alpha : float
        Significance level for determining rejections.
    effect_size : float
        The magnitude of the effect being simulated (passed to effect_modifier).
    progress_desc : str, optional
        Description for the progress bar.

    Returns
    -------
    Tuple[int, List[float]]
        n_rejects : int
            Number of times the null hypothesis was rejected (p < alpha).
        p_values : List[float]
            List of all p-values obtained from the simulations.
    """
    n_rejects = 0
    p_values = []

    for _ in tqdm(range(n_iter), desc=progress_desc, leave=False):
        sampled_groups: List[pd.DataFrame] = pipeline.transform(data)

        if effect_modifier is not None:
            modified_groups = effect_modifier(sampled_groups)
        else:
            modified_groups = sampled_groups

        test_result: TestResult = test.test(modified_groups)

        p_val = test_result.pvalue
        p_values.append(p_val)
        if p_val < alpha:
            n_rejects += 1

    return n_rejects, p_values