"""
smartViz.py — EssentiaX Smart Visualization (Auto-Select)
=================================================
INTELLIGENT automatic variable selection and visualization

NEW FEATURES:
• auto_select=True: Automatically picks most important variables
• Selects based on correlation, variance, missing values, cardinality
• Smart feature ranking and selection
• Visualizes only the most relevant features
"""

import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
import warnings
warnings.filterwarnings('ignore')

# Set style
sns.set_style("whitegrid")
plt.rcParams['figure.figsize'] = (10, 6)


def _auto_select_features(df, target_column=None, max_features=10, verbose=True):
    """
    Automatically select most important features for visualization.
    
    Selection criteria:
    1. High correlation with target (if provided)
    2. High variance (informative features)
    3. Low missing values
    4. Reasonable cardinality for categorical
    
    Returns:
    --------
    list : Selected feature names
    """
    
    if verbose:
        print("\n🤖 AUTO-SELECTING MOST IMPORTANT FEATURES...")
        print("-"*60)
    
    scores = {}
    
    # Separate numeric and categorical
    numeric_cols = df.select_dtypes(include=[np.number]).columns.tolist()
    categorical_cols = df.select_dtypes(include=['object', 'category']).columns.tolist()
    
    # Remove target from consideration
    if target_column:
        if target_column in numeric_cols:
            numeric_cols.remove(target_column)
        if target_column in categorical_cols:
            categorical_cols.remove(target_column)
    
    # ═══════════════════════════════════════════════════════
    # SCORE NUMERIC COLUMNS
    # ═══════════════════════════════════════════════════════
    for col in numeric_cols:
        score = 0.0
        
        # 1. Correlation with target (if numeric target)
        if target_column and df[target_column].dtype in ['int64', 'float64']:
            try:
                corr = abs(df[col].corr(df[target_column]))
                if not np.isnan(corr):
                    score += corr * 100  # Weight: 0-100
            except:
                pass
        
        # 2. Variance (normalized)
        try:
            variance = df[col].var()
            if variance > 0:
                # Normalize by standard deviation
                std = df[col].std()
                if std > 0:
                    cv = std / (abs(df[col].mean()) + 1e-10)  # Coefficient of variation
                    score += min(cv * 10, 50)  # Weight: 0-50
        except:
            pass
        
        # 3. Completeness (less missing = better)
        missing_pct = df[col].isnull().sum() / len(df)
        score += (1 - missing_pct) * 30  # Weight: 0-30
        
        # 4. Uniqueness (more unique values = more informative)
        unique_ratio = df[col].nunique() / len(df)
        if 0.01 < unique_ratio < 0.95:  # Not too few, not too many
            score += 20  # Weight: 20
        
        scores[col] = score
    
    # ═══════════════════════════════════════════════════════
    # SCORE CATEGORICAL COLUMNS
    # ═══════════════════════════════════════════════════════
    for col in categorical_cols:
        score = 0.0
        
        # 1. Association with target (if categorical target)
        if target_column and df[target_column].dtype in ['object', 'category']:
            try:
                # Chi-square test would be ideal, but simple diversity check
                crosstab = pd.crosstab(df[col], df[target_column])
                # If categories have different target distributions, it's informative
                score += 40
            except:
                pass
        
        # 2. Cardinality (moderate is best)
        unique_count = df[col].nunique()
        if 2 <= unique_count <= 20:
            score += 50  # Weight: 50 (good for visualization)
        elif 20 < unique_count <= 50:
            score += 30  # Weight: 30 (okay)
        else:
            score += 10  # Weight: 10 (too many or too few)
        
        # 3. Completeness
        missing_pct = df[col].isnull().sum() / len(df)
        score += (1 - missing_pct) * 30  # Weight: 0-30
        
        # 4. Balance (not too skewed)
        try:
            value_counts = df[col].value_counts()
            max_freq = value_counts.iloc[0] / len(df)
            if max_freq < 0.95:  # Not dominated by one category
                score += 20  # Weight: 20
        except:
            pass
        
        scores[col] = score
    
    # ═══════════════════════════════════════════════════════
    # SELECT TOP FEATURES
    # ═══════════════════════════════════════════════════════
    sorted_features = sorted(scores.items(), key=lambda x: x[1], reverse=True)
    selected = [feat for feat, score in sorted_features[:max_features]]
    
    if verbose:
        print(f"📊 Selected {len(selected)} most important features:\n")
        print("   Rank | Feature              | Score  | Type")
        print("   " + "-"*56)
        
        for i, (feat, score) in enumerate(sorted_features[:max_features], 1):
            feat_type = "Numeric" if feat in numeric_cols else "Categorical"
            print(f"   {i:4d} | {feat:20s} | {score:6.1f} | {feat_type}")
        
        if len(sorted_features) > max_features:
            print(f"\n   ... and {len(sorted_features) - max_features} more features")
        
        print("-"*60)
    
    return selected


def smart_viz(
    df: pd.DataFrame,
    target_column: str = None,
    specific_columns: list = None,
    auto_select: bool = False,
    max_auto_features: int = 8,
    sample_size: int = 10000,
    max_cols_distribution: int = 8,
    max_categories_plot: int = 15,
    show_plots: bool = True,
    plot_types: list = None
) -> None:
    """
    Smart visualization with AUTOMATIC feature selection.
    
    Parameters:
    -----------
    df : pd.DataFrame
        Input dataset
    target_column : str, optional
        Target variable to analyze against
    specific_columns : list, optional
        Manually specify columns to visualize
    auto_select : bool, default=False
        🤖 AUTOMATICALLY select most important features
    max_auto_features : int, default=8
        Maximum features to auto-select
    sample_size : int
        Maximum rows for plotting
    max_cols_distribution : int
        Maximum columns for distribution plots
    max_categories_plot : int
        Maximum categories to show
    show_plots : bool
        Display plots immediately
    plot_types : list, optional
        Specific plots to generate
    
    Examples:
    ---------
    # 🤖 AUTO-SELECT: Let EssentiaX choose best features
    smart_viz(df, auto_select=True)
    
    # 🤖 AUTO-SELECT with target: Choose features most related to target
    smart_viz(df, target_column='Class', auto_select=True)
    
    # 🤖 AUTO-SELECT top 5 features
    smart_viz(df, target_column='Price', auto_select=True, max_auto_features=5)
    
    # Manual selection (original behavior)
    smart_viz(df, specific_columns=['Age', 'Salary'])
    """
    
    print("\n📊 **EssentiaX Smart Visualization**")
    print("============================================================")
    
    # ═══════════════════════════════════════════════════════
    # FEATURE SELECTION
    # ═══════════════════════════════════════════════════════
    if auto_select:
        # 🤖 AUTOMATIC SELECTION
        if specific_columns:
            print("⚠️  Warning: Both auto_select=True and specific_columns provided.")
            print("   Using auto_select (ignoring specific_columns)\n")
        
        cols_to_analyze = _auto_select_features(
            df,
            target_column=target_column,
            max_features=max_auto_features,
            verbose=True
        )
        
        print(f"\n✅ Auto-selected {len(cols_to_analyze)} features for visualization")
        
    elif specific_columns:
        # Manual selection
        missing_cols = [col for col in specific_columns if col not in df.columns]
        if missing_cols:
            raise ValueError(f"❌ Columns not found: {missing_cols}")
        cols_to_analyze = specific_columns
        print(f"📌 Using {len(specific_columns)} manually specified columns")
        
    else:
        # Use all columns
        cols_to_analyze = df.columns.tolist()
        print(f"📊 Using all {len(cols_to_analyze)} columns")
    
    # Remove target from analysis
    if target_column and target_column in cols_to_analyze:
        cols_to_analyze = [col for col in cols_to_analyze if col != target_column]
    
    # ═══════════════════════════════════════════════════════
    # SAMPLING
    # ═══════════════════════════════════════════════════════
    original_rows = len(df)
    if len(df) > sample_size:
        print(f"⚡ Using sample: {sample_size:,} rows (Original: {original_rows:,})")
        df_sample = df.sample(n=sample_size, random_state=42)
    else:
        df_sample = df.copy()
        print(f"📊 Processing: {original_rows:,} rows")
    
    # Column types
    numeric_cols = df_sample[cols_to_analyze].select_dtypes(include=[np.number]).columns.tolist()
    categorical_cols = df_sample[cols_to_analyze].select_dtypes(include=['object', 'category']).columns.tolist()
    
    print(f"\n📋 Visualization Plan:")
    if target_column:
        print(f"   🎯 Target:       {target_column}")
    print(f"   • Numeric:       {len(numeric_cols)} columns")
    print(f"   • Categorical:   {len(categorical_cols)} columns")
    print()
    
    # Determine plots
    if plot_types is None:
        plot_types = ['correlation', 'distribution', 'boxplot', 'categorical', 'scatter', 'pairplot']
    
    plot_count = 0
    
    # ═══════════════════════════════════════════════════════
    # 1. CORRELATION HEATMAP
    # ═══════════════════════════════════════════════════════
    if 'correlation' in plot_types and len(numeric_cols) >= 2:
        print("[1] 🔥 Correlation Heatmap...")
        
        cols_for_corr = numeric_cols[:15]  # Max 15 for readability
        if target_column and df[target_column].dtype in ['int64', 'float64']:
            cols_for_corr = [target_column] + cols_for_corr
        
        plt.figure(figsize=(12, 10))
        corr_matrix = df_sample[cols_for_corr].corr()
        
        # Highlight target correlations if present
        if target_column and target_column in cols_for_corr:
            # Sort by correlation with target
            target_corr = corr_matrix[target_column].abs().sort_values(ascending=False)
            sorted_cols = target_corr.index.tolist()
            corr_matrix = corr_matrix.loc[sorted_cols, sorted_cols]
        
        mask = np.triu(np.ones_like(corr_matrix, dtype=bool))
        
        sns.heatmap(
            corr_matrix,
            mask=mask,
            annot=len(cols_for_corr) <= 10,
            cmap='coolwarm',
            center=0,
            square=True,
            linewidths=0.5,
            cbar_kws={"shrink": 0.8},
            fmt='.2f'
        )
        
        title = 'Correlation Heatmap'
        if auto_select:
            title += ' (Auto-Selected Features)'
        if target_column and df[target_column].dtype in ['int64', 'float64']:
            title += f' - Sorted by {target_column}'
        
        plt.title(title, fontsize=14, fontweight='bold', pad=20)
        plt.tight_layout()
        if show_plots:
            plt.show()
        plot_count += 1
        print("   ✅ Done\n")
    
    # ═══════════════════════════════════════════════════════
    # 2. DISTRIBUTION PLOTS
    # ═══════════════════════════════════════════════════════
    if 'distribution' in plot_types and len(numeric_cols) > 0:
        print("[2] 📊 Distribution Plots...")
        
        cols_to_plot = numeric_cols[:max_cols_distribution]
        
        if target_column and df[target_column].nunique() <= 10:
            # Grouped by target
            print(f"   🎯 Grouped by: {target_column}")
            
            for col in cols_to_plot:
                plt.figure(figsize=(12, 5))
                
                target_values = df_sample[target_column].unique()[:5]
                
                for val in target_values:
                    subset = df_sample[df_sample[target_column] == val][col]
                    if len(subset) > 0:
                        sns.kdeplot(subset, label=f'{target_column}={val}', alpha=0.6, linewidth=2)
                
                plt.title(f'{col} by {target_column}', fontweight='bold', fontsize=12)
                plt.xlabel(col)
                plt.ylabel('Density')
                plt.legend()
                plt.grid(alpha=0.3)
                plt.tight_layout()
                if show_plots:
                    plt.show()
                plot_count += 1
        else:
            # Regular distributions
            n_cols = 2
            n_rows = (len(cols_to_plot) + 1) // 2
            
            fig, axes = plt.subplots(n_rows, n_cols, figsize=(15, 4*n_rows))
            axes = axes.flatten() if isinstance(axes, np.ndarray) else [axes]
            
            for idx, col in enumerate(cols_to_plot):
                sns.histplot(df_sample[col], kde=True, color='skyblue', edgecolor='black', alpha=0.7, ax=axes[idx])
                axes[idx].set_title(f'{col}', fontweight='bold')
                axes[idx].grid(alpha=0.3)
            
            for idx in range(len(cols_to_plot), len(axes)):
                axes[idx].axis('off')
            
            plt.tight_layout()
            if show_plots:
                plt.show()
            plot_count += 1
        
        print("   ✅ Done\n")
    
    # ═══════════════════════════════════════════════════════
    # 3. BOX PLOTS
    # ═══════════════════════════════════════════════════════
    if 'boxplot' in plot_types and len(numeric_cols) > 0:
        print("[3] 📦 Box Plots...")
        
        cols_to_plot = numeric_cols[:max_cols_distribution]
        
        if target_column and df[target_column].dtype in ['object', 'category', 'int64'] and df[target_column].nunique() <= 10:
            # By target
            print(f"   🎯 Grouped by: {target_column}")
            
            n_cols = 2
            n_rows = (len(cols_to_plot) + 1) // 2
            
            fig, axes = plt.subplots(n_rows, n_cols, figsize=(15, 5*n_rows))
            axes = axes.flatten() if isinstance(axes, np.ndarray) else [axes]
            
            for idx, col in enumerate(cols_to_plot):
                sns.boxplot(x=target_column, y=col, data=df_sample, palette='Set2', ax=axes[idx])
                axes[idx].set_title(f'{col} by {target_column}', fontweight='bold')
                axes[idx].grid(alpha=0.3, axis='y')
            
            for idx in range(len(cols_to_plot), len(axes)):
                axes[idx].axis('off')
            
            plt.tight_layout()
            if show_plots:
                plt.show()
            plot_count += 1
        else:
            # Regular box plots
            n_cols = 2
            n_rows = (len(cols_to_plot) + 1) // 2
            
            fig, axes = plt.subplots(n_rows, n_cols, figsize=(15, 4*n_rows))
            axes = axes.flatten() if isinstance(axes, np.ndarray) else [axes]
            
            for idx, col in enumerate(cols_to_plot):
                sns.boxplot(x=df_sample[col], color='lightgreen', ax=axes[idx])
                axes[idx].set_title(f'{col}', fontweight='bold')
                axes[idx].grid(alpha=0.3)
            
            for idx in range(len(cols_to_plot), len(axes)):
                axes[idx].axis('off')
            
            plt.tight_layout()
            if show_plots:
                plt.show()
            plot_count += 1
        
        print("   ✅ Done\n")
    
    # ═══════════════════════════════════════════════════════
    # 4. CATEGORICAL PLOTS
    # ═══════════════════════════════════════════════════════
    if 'categorical' in plot_types and len(categorical_cols) > 0:
        print("[4] 🏷️  Categorical Plots...")
        
        for col in categorical_cols[:5]:
            unique_count = df_sample[col].nunique()
            
            if unique_count > max_categories_plot:
                print(f"   ⏭️  Skipping '{col}' ({unique_count} categories)")
                continue
            
            plt.figure(figsize=(12, 6))
            
            if target_column and df[target_column].nunique() <= 10:
                ct = pd.crosstab(df_sample[col], df_sample[target_column])
                ct.plot(kind='bar', stacked=True, colormap='viridis', ax=plt.gca())
                plt.title(f'{col} by {target_column}', fontsize=14, fontweight='bold', pad=20)
                plt.legend(title=target_column)
            else:
                value_counts = df_sample[col].value_counts().head(max_categories_plot)
                sns.countplot(y=value_counts.index, data=df_sample[df_sample[col].isin(value_counts.index)], palette='viridis')
                plt.title(f'{col}', fontsize=14, fontweight='bold', pad=20)
            
            plt.xlabel('Count')
            plt.grid(axis='x', alpha=0.3)
            plt.tight_layout()
            if show_plots:
                plt.show()
            plot_count += 1
        
        print("   ✅ Done\n")
    
    # ═══════════════════════════════════════════════════════
    # 5. SCATTER PLOTS (vs target)
    # ═══════════════════════════════════════════════════════
    if 'scatter' in plot_types and target_column and df[target_column].dtype in ['int64', 'float64'] and len(numeric_cols) >= 1:
        print(f"[5] 📍 Scatter Plots (vs {target_column})...")
        
        cols_to_plot = numeric_cols[:6]
        
        n_cols = 2
        n_rows = (len(cols_to_plot) + 1) // 2
        
        fig, axes = plt.subplots(n_rows, n_cols, figsize=(15, 5*n_rows))
        axes = axes.flatten() if isinstance(axes, np.ndarray) else [axes]
        
        for idx, col in enumerate(cols_to_plot):
            axes[idx].scatter(df_sample[col], df_sample[target_column], alpha=0.5, s=20, c='steelblue')
            axes[idx].set_xlabel(col, fontweight='bold')
            axes[idx].set_ylabel(target_column, fontweight='bold')
            axes[idx].set_title(f'{col} vs {target_column}')
            axes[idx].grid(alpha=0.3)
        
        for idx in range(len(cols_to_plot), len(axes)):
            axes[idx].axis('off')
        
        plt.tight_layout()
        if show_plots:
            plt.show()
        plot_count += 1
        print("   ✅ Done\n")
    
    # ═══════════════════════════════════════════════════════
    # 6. PAIRPLOT
    # ═══════════════════════════════════════════════════════
    if 'pairplot' in plot_types and 2 <= len(numeric_cols) <= 5:
        print("[6] 🔗 Pairplot...")
        
        pairplot_sample = min(1000, len(df_sample))
        cols_for_pair = numeric_cols[:5]
        
        if target_column and df[target_column].nunique() <= 10:
            df_pair = df_sample[cols_for_pair + [target_column]].sample(n=pairplot_sample, random_state=42)
            sns.pairplot(df_pair, hue=target_column, diag_kind='kde', plot_kws={'alpha': 0.6, 's': 30}, corner=True)
            title = f'Pairplot (colored by {target_column})'
        else:
            df_pair = df_sample[cols_for_pair].sample(n=pairplot_sample, random_state=42)
            sns.pairplot(df_pair, diag_kind='kde', plot_kws={'alpha': 0.6, 's': 30}, corner=True)
            title = 'Pairplot'
        
        if auto_select:
            title += ' - Auto-Selected'
        
        plt.suptitle(title, y=1.02, fontsize=14, fontweight='bold')
        if show_plots:
            plt.show()
        plot_count += 1
        print(f"   ✅ Done\n")
    
    # ═══════════════════════════════════════════════════════
    # SUMMARY
    # ═══════════════════════════════════════════════════════
    print("============================================================")
    print("✅ **VISUALIZATION COMPLETED!**")
    print("============================================================")
    print(f"\n📊 Summary:")
    print(f"   • Plots generated:    {plot_count}")
    if auto_select:
        print(f"   • Auto-selected:      {len(cols_to_analyze)} features")
    if target_column:
        print(f"   • Target variable:    {target_column}")
    print(f"   • Sample size:        {len(df_sample):,} rows")
    
    print("\n💡 Tips:")
    print("   • Use auto_select=True to let EssentiaX choose best features")
    print("   • Combine with target_column for focused analysis")
    print("   • Adjust max_auto_features to select more/fewer features")
    print("============================================================\n")


# ═══════════════════════════════════════════════════════
# EXAMPLE USAGE
# ═══════════════════════════════════════════════════════
if __name__ == "__main__":
    np.random.seed(42)
    
    # Create sample data
    df = pd.DataFrame({
        'Age': np.random.randint(18, 80, 5000),
        'Salary': np.random.gamma(2, 20000, 5000) + 30000,
        'Experience': np.random.randint(0, 30, 5000),
        'Score': np.random.normal(75, 15, 5000),
        'Rating': np.random.randint(1, 6, 5000),
        'City': np.random.choice(['Mumbai', 'Delhi', 'Bangalore'], 5000),
        'Department': np.random.choice(['IT', 'HR', 'Finance'], 5000),
        'Promoted': np.random.choice([0, 1], 5000, p=[0.7, 0.3])
    })
    
    print("\n🤖 AUTO-SELECT EXAMPLES:\n")
    
    # Example 1: Auto-select features
    # smart_viz(df, auto_select=True)
    
    # Example 2: Auto-select with target
    smart_viz(df, target_column='Promoted', auto_select=True, max_auto_features=5)