import numpy as np
import pandas as pd
from scipy import stats
from itertools import combinations
from typing import Union, List, Optional, Tuple


def _calculate_cohens_f(f_statistic: float, df_between: int, df_within: int) -> float:
    """
    Calculate Cohen's f effect size from F statistic.
    
    Cohen's f = sqrt(eta^2 / (1 - eta^2))
    where eta^2 = SS_between / SS_total = (F * df_between) / (F * df_between + df_within)
    
    Parameters
    ----------
    f_statistic : float
        F statistic from ANOVA
    df_between : int
        Degrees of freedom between groups
    df_within : int
        Degrees of freedom within groups
        
    Returns
    -------
    float
        Cohen's f effect size
    """
    if f_statistic <= 0 or df_between <= 0 or df_within <= 0:
        return 0.0
    
    # Calculate eta squared
    eta_squared = (f_statistic * df_between) / (f_statistic * df_between + df_within)
    
    # Calculate Cohen's f
    if eta_squared >= 1.0:
        return np.inf
    
    cohens_f = np.sqrt(eta_squared / (1 - eta_squared))
    return cohens_f


def _perform_two_way_anova(X: np.ndarray, y: np.ndarray, i: int, j: int) -> Tuple[float, float, float]:
    """
    Perform two-way ANOVA for feature pair interaction effect.
    
    This function discretizes continuous features into groups and performs
    a two-way ANOVA to test for interaction effects between features i and j.
    
    Parameters
    ----------
    X : np.ndarray
        Feature matrix of shape (n_samples, n_features)
    y : np.ndarray
        Target variable of shape (n_samples,)
    i : int
        Index of first feature
    j : int
        Index of second feature
        
    Returns
    -------
    Tuple[float, float, float]
        F statistic, p-value, and Cohen's f effect size for interaction
    """
    try:
        # Extract features
        feature_i = X[:, i]
        feature_j = X[:, j]
        
        # Discretize continuous features into groups (tertiles)
        # This is necessary for ANOVA which requires categorical grouping
        q_i = np.quantile(feature_i, [0.33, 0.67])
        q_j = np.quantile(feature_j, [0.33, 0.67])
        
        group_i = np.digitize(feature_i, q_i)  # 0, 1, 2
        group_j = np.digitize(feature_j, q_j)  # 0, 1, 2
        
        # Create interaction groups
        unique_combinations = []
        group_data = []
        
        for gi in range(3):  # 3 groups for feature i
            for gj in range(3):  # 3 groups for feature j
                mask = (group_i == gi) & (group_j == gj)
                if np.sum(mask) > 0:  # Only include non-empty groups
                    unique_combinations.append((gi, gj))
                    group_data.append(y[mask])
        
        # Need at least 2 groups for ANOVA
        if len(group_data) < 2:
            return 0.0, 1.0, 0.0
            
        # Perform one-way ANOVA on interaction groups
        f_stat, p_value = stats.f_oneway(*group_data)
        
        if np.isnan(f_stat) or np.isnan(p_value):
            return 0.0, 1.0, 0.0
            
        # Calculate degrees of freedom
        df_between = len(group_data) - 1
        df_within = len(y) - len(group_data)
        
        # Calculate Cohen's f effect size
        cohens_f = _calculate_cohens_f(f_stat, df_between, df_within)
        
        return float(f_stat), float(p_value), float(cohens_f)
        
    except Exception:
        # Return neutral values if calculation fails
        return 0.0, 1.0, 0.0


def get_anova_interactions(
    X: Union[np.ndarray, pd.DataFrame],
    y: Union[np.ndarray, List],
    feature_names: Optional[List[str]] = None
) -> pd.DataFrame:
    """
    Perform two-way ANOVA analysis for all feature pairs to detect interaction effects.
    
    This function performs ANOVA analysis on all pairs of features to identify
    significant interaction effects. Features are discretized into tertiles to
    enable ANOVA analysis on continuous variables.
    
    Parameters
    ----------
    X : array-like of shape (n_samples, n_features)
        Input features.
    y : array-like of shape (n_samples,)
        Target values.
    feature_names : list of str or None, default=None
        Feature names. If None, will use DataFrame columns or generate default names.
        
    Returns
    -------
    pd.DataFrame
        DataFrame with columns:
        - i: index of first feature
        - j: index of second feature  
        - feature_i: name of first feature
        - feature_j: name of second feature
        - F_statistic: F statistic from two-way ANOVA
        - p_value: p-value from ANOVA test
        - cohens_f: Cohen's f effect size
        Sorted by cohens_f in descending order.
        
    Examples
    --------
    >>> X = np.random.randn(100, 5)
    >>> y = X[:, 0] * X[:, 1] + np.random.randn(100) * 0.1  # interaction effect
    >>> df = get_anova_interactions(X, y)
    >>> print(df.head())
    """
    # Convert inputs to numpy arrays
    X = np.asarray(X, dtype=float)
    y = np.asarray(y, dtype=float).reshape(-1)
    
    # Validate input shapes
    if X.shape[0] != len(y):
        raise ValueError(f"X and y have incompatible shapes: X has {X.shape[0]} samples, y has {len(y)}")
    
    n_samples, n_features = X.shape
    
    # Handle feature names
    if feature_names is None:
        if hasattr(X, 'columns'):
            feature_names = list(X.columns)
        else:
            feature_names = [f"x_{i}" for i in range(n_features)]
    elif len(feature_names) != n_features:
        raise ValueError(f"feature_names length ({len(feature_names)}) does not match "
                        f"number of features ({n_features})")
    
    # Calculate ANOVA statistics for all feature pairs
    results = []
    
    for i in range(n_features):
        for j in range(i + 1, n_features):
            # Perform two-way ANOVA for interaction
            f_stat, p_val, cohens_f = _perform_two_way_anova(X, y, i, j)
            
            results.append({
                'i': i,
                'j': j,
                'feature_i': feature_names[i],
                'feature_j': feature_names[j],
                'F_statistic': f_stat,
                'p_value': p_val,
                'cohens_f': cohens_f
            })
    
    # Convert to DataFrame and sort by Cohen's f (descending)
    results_df = pd.DataFrame(results)
    results_df = results_df.sort_values('cohens_f', ascending=False).reset_index(drop=True)
    
    return results_df


def get_anova_interactions_filtered(
    X: Union[np.ndarray, pd.DataFrame],
    y: Union[np.ndarray, List],
    feature_names: Optional[List[str]] = None,
    p_threshold: float = 0.05,
    cohens_f_threshold: float = 0.1
) -> pd.DataFrame:
    """
    Get ranked interaction dataframe with only statistically significant interactions.
    
    This is a convenience function that filters the results from get_anova_interactions()
    to only include interactions that meet significance and effect size thresholds.
    
    Parameters
    ----------
    X : array-like of shape (n_samples, n_features)
        Input features.
    y : array-like of shape (n_samples,)
        Target values.
    feature_names : list of str or None, default=None
        Feature names. If None, will use DataFrame columns or generate default names.
    p_threshold : float, default=0.05
        Maximum p-value for statistical significance.
    cohens_f_threshold : float, default=0.1
        Minimum Cohen's f effect size to include in results.
        
    Returns
    -------
    pd.DataFrame
        Filtered DataFrame with only significant interactions, sorted by Cohen's f.
    """
    # Get all ANOVA results
    df = get_anova_interactions(X, y, feature_names=feature_names)
    
    # Filter by significance and effect size
    filtered_df = df[
        (df['p_value'] <= p_threshold) & 
        (df['cohens_f'] >= cohens_f_threshold)
    ].reset_index(drop=True)
    
    return filtered_df


if __name__ == "__main__":
    # Test the ANOVA interaction function
    print("Testing ANOVA interaction analysis...")
    
    # Generate synthetic data with known interaction
    np.random.seed(42)
    n_samples, n_features = 200, 5
    
    X = np.random.randn(n_samples, n_features)
    # Create target with interaction between features 0 and 1
    y = (X[:, 0] * X[:, 1] + 
         0.5 * X[:, 2] + 
         np.random.randn(n_samples) * 0.5)
    
    # Test the function
    df = get_anova_interactions(X, y)
    print("\nTop 5 feature pairs by Cohen's f:")
    print(df.head())
    
    # Test filtered version
    df_filtered = get_anova_interactions_filtered(X, y)
    print(f"\nSignificant interactions (p < 0.05, Cohen's f >= 0.1): {len(df_filtered)}")
    if len(df_filtered) > 0:
        print(df_filtered.head())
