import pandas as pd
import numpy as np
from scipy import stats
import seaborn as sns
import matplotlib.pyplot as plt
from scipy.stats import pearsonr, spearmanr, kendalltau

def prepare_correlation_data(df, method='pearson'):
    """
    Prepare data for correlation analysis with appropriate checks
    """
    # Remove non-numeric columns for numeric correlations
    numeric_df = df.select_dtypes(include=[np.number])

    # Check for sufficient data
    if len(numeric_df) < 3:
        print("Warning: Sample size < 3, results may be unreliable")

    # Handle missing values
    missing_pct = numeric_df.isnull().sum() / len(numeric_df)
    high_missing = missing_pct[missing_pct > 0.1]
    if not high_missing.empty:
        print(f"Variables with >10% missing: {high_missing.index.tolist()}")

    # Remove outliers for Pearson (optional)
    if method == 'pearson':
        Q1 = numeric_df.quantile(0.25)
        Q3 = numeric_df.quantile(0.75)
        IQR = Q3 - Q1
        outlier_mask = ~((numeric_df < (Q1 - 1.5 * IQR)) | (numeric_df > (Q3 + 1.5 * IQR))).any(axis=1)
        clean_df = numeric_df[outlier_mask]
        print(f"Removed {len(numeric_df) - len(clean_df)} outliers")
        return clean_df

    return numeric_df

def comprehensive_correlation_analysis(df, variables=None, methods=['pearson', 'spearman']):
    """
    Perform multiple correlation analyses with statistical significance
    """
    if variables:
        df = df[variables]

    results = {}

    for method in methods:
        if method == 'pearson':
            corr_matrix = df.corr(method='pearson')
            # Calculate p-values
            p_values = pd.DataFrame(index=df.columns, columns=df.columns)
            for i, col1 in enumerate(df.columns):
                for j, col2 in enumerate(df.columns):
                    if i != j:
                        stat, p = pearsonr(df[col1].dropna(), df[col2].dropna())
                        p_values.loc[col1, col2] = p
                    else:
                        p_values.loc[col1, col2] = 0

        elif method == 'spearman':
            corr_matrix = df.corr(method='spearman')
            p_values = pd.DataFrame(index=df.columns, columns=df.columns)
            for i, col1 in enumerate(df.columns):
                for j, col2 in enumerate(df.columns):
                    if i != j:
                        stat, p = spearmanr(df[col1].dropna(), df[col2].dropna())
                        p_values.loc[col1, col2] = p
                    else:
                        p_values.loc[col1, col2] = 0

        results[method] = {
            'correlation': corr_matrix,
            'p_values': p_values.astype(float),
            'significant': (p_values.astype(float) < 0.05) & (p_values.astype(float) > 0)
        }

    return results

def partial_correlation(df, x, y, control_vars):
    """
    Calculate partial correlation controlling for other variables
    """
    from sklearn.linear_model import LinearRegression

    # Residualize x and y against control variables
    lr = LinearRegression()

    # Get residuals for x
    lr.fit(df[control_vars], df[x])
    x_resid = df[x] - lr.predict(df[control_vars])

    # Get residuals for y
    lr.fit(df[control_vars], df[y])
    y_resid = df[y] - lr.predict(df[control_vars])

    # Correlate residuals
    partial_corr, p_value = pearsonr(x_resid, y_resid)
    return partial_corr, p_value

def rolling_correlation(df, var1, var2, window=30):
    """
    Calculate rolling correlation for time series data
    """
    return df[var1].rolling(window=window).corr(df[var2])

def create_correlation_heatmap(corr_matrix, p_values=None, figsize=(10, 8)):
    """
    Create publication-ready correlation heatmap
    """
    plt.figure(figsize=figsize)

    # Create mask for upper triangle
    mask = np.triu(np.ones_like(corr_matrix, dtype=bool))

    # Create significance mask if p-values provided
    if p_values is not None:
        sig_mask = p_values < 0.05
        annot_matrix = corr_matrix.copy()
        annot_matrix = annot_matrix.round(3).astype(str)
        annot_matrix[~sig_mask] = annot_matrix[~sig_mask] + '\n(ns)'
    else:
        annot_matrix = True

    sns.heatmap(corr_matrix, mask=mask, annot=annot_matrix, 
                cmap='RdBu_r', center=0, square=True, 
                cbar_kws={'label': 'Correlation Coefficient'},
                fmt='' if p_values is not None else '.3f')

    plt.title('Correlation Matrix with Statistical Significance')
    plt.tight_layout()
    return plt.gcf()

def correlation_strength_plot(corr_matrix):
    """
    Visualize correlation strength distribution
    """
    # Get upper triangle values
    mask = np.triu(np.ones_like(corr_matrix, dtype=bool), k=1)
    correlations = corr_matrix.values[mask]

    plt.figure(figsize=(10, 6))
    plt.hist(correlations, bins=20, alpha=0.7, color='skyblue', edgecolor='black')
    plt.axvline(0, color='red', linestyle='--', alpha=0.7)
    plt.xlabel('Correlation Coefficient')
    plt.ylabel('Frequency')
    plt.title('Distribution of Correlation Coefficients')
    plt.grid(True, alpha=0.3)
    return plt.gcf()

def correlation_confidence_interval(r, n, confidence=0.95):
    """
    Calculate confidence interval for Pearson correlation
    """
    z = np.arctanh(r)  # Fisher's z-transformation
    se = 1 / np.sqrt(n - 3)
    alpha = 1 - confidence
    z_critical = stats.norm.ppf(1 - alpha/2)

    ci_lower = np.tanh(z - z_critical * se)
    ci_upper = np.tanh(z + z_critical * se)

    return ci_lower, ci_upper