Retention Analysis

"""
Basic Retention Analysis Example

This example demonstrates how to use the Retention Analysis skill
to perform basic survival analysis and retention calculations.
"""

import pandas as pd
import numpy as np
import sys
import os
sys.path.append(os.path.join(os.path.dirname(__file__), '..'))
from scripts.survival_analyzer import SurvivalAnalyzer
from scripts.retention_analyzer import RetentionAnalyzer
from scripts.visualizer import RetentionVisualizer

def create_sample_retention_data():
    """
    Create sample retention data for demonstration.
    """
    np.random.seed(42)
    n_users = 2000

    # Generate user data
    data = {
        'user_id': range(1, n_users + 1),
        'registration_date': pd.date_range('2023-01-01', periods=n_users, freq='H'),
        'tenure_days': np.random.exponential(scale=180, size=n_users).astype(int),
        'churned': np.random.choice([0, 1], n_users, p=[0.6, 0.4]),
        'age_group': np.random.choice(['18-25', '26-35', '36-45', '46+'], n_users),
        'subscription_type': np.random.choice(['Basic', 'Premium', 'Enterprise'], n_users, p=[0.5, 0.3, 0.2]),
        'payment_method': np.random.choice(['Credit Card', 'Bank Transfer', 'PayPal'], n_users, p=[0.6, 0.2, 0.2]),
        'usage_frequency': np.random.exponential(scale=10, size=n_users).astype(int),
        'support_tickets': np.random.poisson(lam=1, size=n_users),
        'monthly_revenue': np.random.normal(loc=50, scale=20, size=n_users).round(2)
    }

    df = pd.DataFrame(data)

    # Ensure logical consistency: users who haven't churned should have right-censored data
    active_mask = df['churned'] == 0
    max_days = (pd.Timestamp.now() - df.loc[active_mask, 'registration_date']).dt.days
    df.loc[active_mask, 'tenure_days'] = np.minimum(
        df.loc[active_mask, 'tenure_days'],
        max_days
    ).astype(int)

    # Add some realistic patterns
    # Premium users tend to stay longer
    premium_mask = df['subscription_type'] == 'Premium'
    df.loc[premium_mask, 'tenure_days'] = (df.loc[premium_mask, 'tenure_days'] * 1.5).astype(int)

    # Users with more support tickets are more likely to churn
    high_tickets_mask = df['support_tickets'] > 2
    df.loc[high_tickets_mask, 'churned'] = 1

    return df

def main():
    """
    Run basic retention analysis example.
    """
    print("=== Basic Retention Analysis Example ===\n")

    # Create sample data
    print("1. Creating sample retention data...")
    df = create_sample_retention_data()
    print(f"   Generated data for {len(df)} users")
    print(f"   Churn rate: {df['churned'].mean():.1%}")
    print(f"   Average tenure: {df['tenure_days'].mean():.1f} days")

    # Initialize survival analyzer
    print("\n2. Performing survival analysis...")
    survival_analyzer = SurvivalAnalyzer()
    survival_analyzer.load_data(df, time_col='tenure_days', event_col='churned')

    # Fit Kaplan-Meier curve
    km_fit = survival_analyzer.fit_kaplan_meier()

    # Plot survival curve
    print("   Creating survival curve visualization...")
    survival_fig = survival_analyzer.plot_survival_curve(
        title="User Survival Curve"
    )
    survival_fig.savefig('survival_curve.png', dpi=300, bbox_inches='tight')

    # Compare survival by subscription type
    print("   Comparing survival by subscription type...")
    groups = survival_analyzer.compare_survival_groups('subscription_type')
    comparison_fig = survival_analyzer.plot_group_comparison(
        'subscription_type',
        title="Survival Comparison by Subscription Type"
    )
    comparison_fig.savefig('survival_by_subscription.png', dpi=300, bbox_inches='tight')

    # Cox regression analysis
    print("\n3. Performing Cox regression analysis...")
    categorical_cols = ['age_group', 'subscription_type', 'payment_method']
    df_encoded = pd.get_dummies(df, columns=categorical_cols, drop_first=True)

    covariates = [col for col in df_encoded.columns
                 if col not in ['user_id', 'registration_date', 'tenure_days', 'churned']]

    # Load encoded data for Cox model
    survival_analyzer.load_data(df_encoded, time_col='tenure_days', event_col='churned')
    cox_fit = survival_analyzer.fit_cox_model(covariates[:10])  # Limit to 10 covariates for clarity

    # Plot Cox coefficients
    print("   Creating Cox coefficient visualization...")
    cox_fig = survival_analyzer.plot_cox_coefficients(
        title="Cox Model - Factors Influencing Churn"
    )
    cox_fig.savefig('cox_coefficients.png', dpi=300, bbox_inches='tight')

    # Generate insights
    print("\n4. Generating survival analysis insights...")
    insights = survival_analyzer.generate_insights()
    print("Key Insights:")
    for insight in insights:
        print(f"  {insight}")

    # Get summary statistics
    stats = survival_analyzer.get_summary_statistics()
    print(f"\nSummary Statistics:")
    print(f"  Total users: {stats['n_observations']:,}")
    print(f"  Events (churn): {stats['n_events']:,}")
    print(f"  Event rate: {stats['event_rate']:.1%}")
    if 'median_survival' in stats and stats['median_survival'] != np.inf:
        print(f"  Median survival: {stats['median_survival']:.1f} days")
    if 'survival_at_30_days' in stats:
        print(f"  30-day survival: {stats['survival_at_30_days']:.1%}")
        print(f"  90-day survival: {stats['survival_at_90_days']:.1%}")

    # Export survival report
    print("\n5. Exporting survival analysis report...")
    report_path = survival_analyzer.export_report('survival_analysis_report.html')

    # Initialize retention analyzer
    print("\n6. Initializing retention analyzer for additional analysis...")
    retention_analyzer = RetentionAnalyzer()
    retention_analyzer.load_data(df)

    # Prepare churn prediction data
    feature_cols = ['age_group', 'subscription_type', 'payment_method',
                   'usage_frequency', 'support_tickets', 'monthly_revenue']

    X, y = retention_analyzer.prepare_churn_data(
        target_col='churned',
        feature_cols=feature_cols
    )

    # Train churn prediction model
    print("   Training churn prediction model...")
    model_results = retention_analyzer.train_churn_model(
        X, y, model_type='random_forest'
    )

    print(f"   Model AUC Score: {model_results['auc_score']:.3f}")

    # Predict churn risk for all users
    print("   Predicting churn risk...")
    churn_risk_df = retention_analyzer.predict_churn_risk(df)

    # Visualize churn prediction results
    print("   Creating churn prediction visualizations...")
    visualizer = RetentionVisualizer()
    churn_results_fig = visualizer.plot_churn_prediction_results(
        model_results['y_true'],
        model_results['y_pred'],
        model_results['y_pred_proba'],
        retention_analyzer.feature_importance
    )
    churn_results_fig.savefig('churn_prediction_results.png', dpi=300, bbox_inches='tight')

    # Calculate customer lifetime value (simplified)
    print("\n7. Calculating customer lifetime value...")
    # Use tenure and monthly revenue to estimate simple CLV
    clv_estimates = df['tenure_days'] * df['monthly_revenue'] / 30  # Convert to months
    print(f"   Estimated CLV for {len(clv_estimates)} users")
    print(f"   Average CLV: ${clv_estimates.mean():.2f}")
    print(f"   Median CLV: ${clv_estimates.median():.2f}")

    # Generate retention insights
    print("\n8. Generating overall retention insights...")
    retention_insights = retention_analyzer.generate_retention_insights()
    print("Additional Insights:")
    for insight in retention_insights:
        print(f"  {insight}")

    # Export comprehensive report
    print("\n9. Exporting comprehensive retention report...")
    retention_report_path = retention_analyzer.export_retention_report(
        'comprehensive_retention_report.html'
    )

    print(f"\n=== Analysis Complete ===")
    print(f"Generated files:")
    print(f"  - survival_curve.png")
    print(f"  - survival_by_subscription.png")
    print(f"  - cox_coefficients.png")
    print(f"  - churn_prediction_results.png")
    print(f"  - survival_analysis_report.html")
    print(f"  - comprehensive_retention_report.html")

if __name__ == "__main__":
    main()

        <!DOCTYPE html>
        <html>
        <head>
            <title>Retention Analysis Report</title>
            <style>
                body { font-family: Arial, sans-serif; margin: 40px; }
                .header { background-color: #f0f0f0; padding: 20px; border-radius: 5px; }
                .metrics { margin: 20px 0; }
                .insights { background-color: #f9f9f9; padding: 20px; border-radius: 5px; }
                table { border-collapse: collapse; width: 100%; }
                th, td { border: 1px solid #ddd; padding: 8px; text-align: left; }
                th { background-color: #f2f2f2; }
            </style>
        </head>
        <body>
            <div class="header">
                <h1>Retention Analysis Report</h1>
                <p>Generated on 2025-12-19 21:58:01</p>
            </div>

            <div class="metrics">
                <h2>Key Metrics</h2>
                <ul>
                    <li><strong>Total Users:</strong> 2,000</li>
                    <li><strong>Cohort Analysis:</strong> Not Available</li>
                    <li><strong>Churn Model:</strong> Trained</li>
        <li><strong>Top Churn Predictors:</strong> monthly_revenue, usage_frequency, support_tickets</li>
                </ul>
            </div>

            <div class="insights">
                <h2>Key Insights</h2>
                <ul>
        <li>🎯 Top churn predictors:</li><li>   • monthly_revenue</li><li>   • usage_frequency</li><li>   • support_tickets</li><li>✅ Churn prediction model available for proactive retention</li>
                </ul>
            </div>
        </body>
        </html>
        

        <!DOCTYPE html>
        <html>
        <head>
            <title>Survival Analysis Report</title>
            <style>
                body { font-family: Arial, sans-serif; margin: 40px; }
                .header { background-color: #f0f0f0; padding: 20px; border-radius: 5px; }
                .metrics { margin: 20px 0; }
                .insights { background-color: #f9f9f9; padding: 20px; border-radius: 5px; }
                table { border-collapse: collapse; width: 100%; }
                th, td { border: 1px solid #ddd; padding: 8px; text-align: left; }
                th { background-color: #f2f2f2; }
            </style>
        </head>
        <body>
            <div class="header">
                <h1>Survival Analysis Report</h1>
                <p>Generated on 2025-12-19 21:57:59</p>
            </div>

            <div class="metrics">
                <h2>Key Metrics</h2>
                <ul>
                    <li><strong>Total Observations:</strong> 2,000</li>
                    <li><strong>Events (Churn):</strong> 892 (44.6%)</li>
                    <li><strong>Median Time:</strong> 138.0</li>
        <li><strong>Median Survival:</strong> 319.0</li><li><strong>30-Day Survival:</strong> 92.9%</li><li><strong>90-Day Survival:</strong> 81.3%</li><li><strong>Model Concordance:</strong> 0.587</li>
                </ul>
            </div>

            <div class="insights">
                <h2>Key Insights</h2>
                <ul>
        <li>📊 Median user lifetime: 319.0 days</li><li>📈 30-day retention: 92.9%</li><li>📈 90-day retention: 81.3%</li><li>🔺 Top risk factors for churn:</li><li>   • support_tickets: 1.3x higher risk</li><li>   • subscription_type_Enterprise: 1.1x higher risk</li><li>   • monthly_revenue: 1.0x higher risk</li><li>🔻 Factors that reduce churn risk:</li><li>   • subscription_type_Premium: 36% lower risk</li><li>   • payment_method_PayPal: 15% lower risk</li><li>   • age_group_36-45: 11% lower risk</li><li>⚠️ Weak predictive model - consider additional features</li>
                </ul>
            </div>
        </body>
        </html>
        

"""
Retention Analysis Core Functions

This module provides comprehensive retention analysis capabilities including
cohort analysis, churn prediction, and customer lifetime value calculation.
"""

import pandas as pd
import numpy as np
from typing import Dict, List, Optional, Tuple, Union
from sklearn.model_selection import train_test_split
from sklearn.ensemble import RandomForestClassifier, GradientBoostingClassifier
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import classification_report, roc_auc_score, precision_recall_curve
from sklearn.preprocessing import StandardScaler, LabelEncoder
import warnings

class RetentionAnalyzer:
    """
    Comprehensive retention analysis toolkit.

    Provides cohort analysis, churn prediction, and CLV calculation
    for customer retention analysis.
    """

    def __init__(self):
        self.data = None
        self.cohort_matrix = None
        self.churn_model = None
        self.feature_importance = None
        self.scaler = None

    def load_data(self, data: pd.DataFrame) -> None:
        """
        Load retention analysis data.

        Args:
            data: DataFrame with user activity and retention data
        """
        self.data = data.copy()
        print(f"Loaded retention data with {len(data)} records")

    def create_cohort_matrix(self,
                           user_id_col: str,
                           date_col: str,
                           activity_col: str,
                           period_type: str = 'monthly') -> pd.DataFrame:
        """
        Create cohort retention matrix.

        Args:
            user_id_col: Column name for user identifier
            date_col: Column name for activity date
            activity_col: Column name indicating active status (1/0)
            period_type: 'monthly' or 'weekly' periods

        Returns:
            Cohort retention matrix
        """
        if self.data is None:
            raise ValueError("No data loaded. Use load_data() first.")

        # Convert date column to datetime
        self.data[date_col] = pd.to_datetime(self.data[date_col])

        # Determine period
        if period_type == 'monthly':
            self.data['period'] = self.data[date_col].dt.to_period('M')
        elif period_type == 'weekly':
            self.data['period'] = self.data[date_col].dt.to_period('W')
        else:
            raise ValueError("period_type must be 'monthly' or 'weekly'")

        # Get user's first activity period (cohort)
        cohort_periods = self.data.groupby(user_id_col)['period'].min().reset_index()
        cohort_periods.columns = [user_id_col, 'cohort_period']

        # Merge cohort information back to main data
        self.data = self.data.merge(cohort_periods, on=user_id_col)

        # Calculate period number (0 = first period)
        self.data['period_number'] = (self.data['period'] - self.data['cohort_period']).apply(lambda x: x.n)

        # Create cohort matrix
        cohort_data = self.data.groupby(['cohort_period', 'period_number'])[user_id_col].nunique().reset_index()
        cohort_sizes = self.data.groupby('cohort_period')[user_id_col].nunique().reset_index()
        cohort_sizes.columns = ['cohort_period', 'cohort_size']

        # Merge cohort sizes
        cohort_data = cohort_data.merge(cohort_sizes, on='cohort_period')

        # Calculate retention rates
        cohort_data['retention_rate'] = cohort_data[user_id_col] / cohort_data['cohort_size']

        # Create pivot table
        self.cohort_matrix = cohort_data.pivot(
            index='cohort_period',
            columns='period_number',
            values='retention_rate'
        )

        return self.cohort_matrix

    def plot_cohort_heatmap(self, figsize: Tuple[int, int] = (12, 8)) -> 'plt.Figure':
        """
        Plot cohort retention heatmap.

        Args:
            figsize: Figure size tuple

        Returns:
            Matplotlib figure object
        """
        import matplotlib.pyplot as plt
        import seaborn as sns

        if self.cohort_matrix is None:
            raise ValueError("No cohort matrix created. Use create_cohort_matrix() first.")

        fig, ax = plt.subplots(figsize=figsize)

        # Create heatmap
        sns.heatmap(self.cohort_matrix,
                   annot=True,
                   fmt='.0%',
                   cmap='YlOrRd',
                   ax=ax)

        ax.set_title('Cohort Retention Matrix', fontsize=16, fontweight='bold', pad=20)
        ax.set_xlabel('Period Number', fontsize=14)
        ax.set_ylabel('Cohort Period', fontsize=14)

        plt.tight_layout()
        return fig

    def prepare_churn_data(self,
                          target_col: str,
                          feature_cols: List[str],
                          exclude_cols: List[str] = None) -> Tuple[pd.DataFrame, pd.Series]:
        """
        Prepare data for churn prediction.

        Args:
            target_col: Target variable column (churn indicator)
            feature_cols: List of feature columns to use
            exclude_cols: Columns to exclude from features

        Returns:
            Tuple of (X_features, y_target)
        """
        if self.data is None:
            raise ValueError("No data loaded. Use load_data() first.")

        # Prepare features
        if exclude_cols:
            feature_cols = [col for col in feature_cols if col not in exclude_cols]

        X = self.data[feature_cols].copy()
        y = self.data[target_col].copy()

        # Handle categorical variables
        categorical_cols = X.select_dtypes(include=['object']).columns
        for col in categorical_cols:
            le = LabelEncoder()
            X[col] = le.fit_transform(X[col].astype(str))

        # Handle missing values
        X = X.fillna(X.median())

        self.feature_cols = feature_cols
        return X, y

    def train_churn_model(self,
                         X: pd.DataFrame,
                         y: pd.Series,
                         model_type: str = 'random_forest',
                         test_size: float = 0.2) -> Dict:
        """
        Train churn prediction model.

        Args:
            X: Feature matrix
            y: Target variable
            model_type: 'logistic', 'random_forest', 'gradient_boosting'
            test_size: Test set proportion

        Returns:
            Dictionary with model results
        """
        # Split data
        X_train, X_test, y_train, y_test = train_test_split(
            X, y, test_size=test_size, random_state=42, stratify=y
        )

        # Scale features
        self.scaler = StandardScaler()
        X_train_scaled = self.scaler.fit_transform(X_train)
        X_test_scaled = self.scaler.transform(X_test)

        # Select model
        if model_type == 'logistic':
            self.churn_model = LogisticRegression(random_state=42)
        elif model_type == 'random_forest':
            self.churn_model = RandomForestClassifier(random_state=42, n_estimators=100)
        elif model_type == 'gradient_boosting':
            self.churn_model = GradientBoostingClassifier(random_state=42, n_estimators=100)
        else:
            raise ValueError("model_type must be 'logistic', 'random_forest', or 'gradient_boosting'")

        # Train model
        if model_type == 'logistic':
            self.churn_model.fit(X_train_scaled, y_train)
        else:
            self.churn_model.fit(X_train, y_train)

        # Make predictions
        if model_type == 'logistic':
            y_pred = self.churn_model.predict(X_test_scaled)
            y_pred_proba = self.churn_model.predict_proba(X_test_scaled)[:, 1]
        else:
            y_pred = self.churn_model.predict(X_test)
            y_pred_proba = self.churn_model.predict_proba(X_test)[:, 1]

        # Calculate metrics
        results = {
            'model': self.churn_model,
            'y_true': y_test,
            'y_pred': y_pred,
            'y_pred_proba': y_pred_proba,
            'auc_score': roc_auc_score(y_test, y_pred_proba),
            'classification_report': classification_report(y_test, y_pred)
        }

        # Feature importance
        if hasattr(self.churn_model, 'feature_importances_'):
            self.feature_importance = pd.DataFrame({
                'feature': X.columns,
                'importance': self.churn_model.feature_importances_
            }).sort_values('importance', ascending=False)
        elif hasattr(self.churn_model, 'coef_'):
            self.feature_importance = pd.DataFrame({
                'feature': X.columns,
                'importance': np.abs(self.churn_model.coef_[0])
            }).sort_values('importance', ascending=False)

        return results

    def predict_churn_risk(self, X: pd.DataFrame) -> pd.DataFrame:
        """
        Predict churn risk for new data.

        Args:
            X: Feature matrix for prediction

        Returns:
            DataFrame with churn probabilities and risk categories
        """
        if self.churn_model is None:
            raise ValueError("No model trained. Use train_churn_model() first.")

        # Prepare data
        X_processed = X[self.feature_cols].copy()

        # Handle categorical variables
        categorical_cols = X_processed.select_dtypes(include=['object']).columns
        for col in categorical_cols:
            le = LabelEncoder()
            X_processed[col] = le.fit_transform(X_processed[col].astype(str))

        # Handle missing values
        X_processed = X_processed.fillna(X_processed.median())

        # Scale features if needed
        if isinstance(self.churn_model, LogisticRegression):
            X_processed = self.scaler.transform(X_processed)

        # Predict
        churn_prob = self.churn_model.predict_proba(X_processed)[:, 1]

        # Create results DataFrame
        results = X.copy()
        results['churn_probability'] = churn_prob
        results['risk_category'] = pd.cut(churn_prob,
                                        bins=[0, 0.2, 0.5, 0.8, 1.0],
                                        labels=['Low', 'Medium', 'High', 'Very High'])

        return results

    def calculate_clv(self,
                     revenue_col: str,
                     user_id_col: str,
                     time_col: str,
                     discount_rate: float = 0.1) -> pd.DataFrame:
        """
        Calculate customer lifetime value (CLV).

        Args:
            revenue_col: Column name for revenue/purchase amount
            user_id_col: Column name for user identifier
            time_col: Column name for time period
            discount_rate: Discount rate for present value calculation

        Returns:
            DataFrame with CLV calculations
        """
        if self.data is None:
            raise ValueError("No data loaded. Use load_data() first.")

        # Calculate metrics per user
        user_metrics = self.data.groupby(user_id_col).agg({
            revenue_col: ['sum', 'mean', 'count'],
            time_col: ['min', 'max']
        }).reset_index()

        # Flatten column names
        user_metrics.columns = [user_id_col, 'total_revenue', 'avg_revenue',
                               'transaction_count', 'first_period', 'last_period']

        # Calculate customer lifespan (in periods)
        user_metrics['lifespan'] = user_metrics['last_period'] - user_metrics['first_period'] + 1

        # Calculate average purchase frequency
        user_metrics['purchase_frequency'] = user_metrics['transaction_count'] / user_metrics['lifespan']

        # Calculate CLV (simplified version)
        user_metrics['clv_simple'] = user_metrics['total_revenue']

        # Calculate CLV with discounting (more sophisticated)
        user_metrics['clv_discounted'] = 0
        for _, user in user_metrics.iterrows():
            periods = user['lifespan']
            avg_revenue = user['avg_revenue']
            freq = user['purchase_frequency']

            # Present value of future cash flows
            clv = 0
            for period in range(1, int(periods) + 1):
                clv += (avg_revenue * freq) / ((1 + discount_rate) ** period)

            user_metrics.loc[user.name, 'clv_discounted'] = clv

        return user_metrics.sort_values('clv_discounted', ascending=False)

    def get_retention_summary(self) -> Dict:
        """
        Get comprehensive retention summary statistics.

        Returns:
            Dictionary with retention metrics
        """
        if self.data is None:
            raise ValueError("No data loaded. Use load_data() first.")

        summary = {
            'total_users': self.data.nunique()[self.data.columns[0]] if len(self.data.columns) > 0 else 0,
            'cohort_matrix_available': self.cohort_matrix is not None,
            'model_trained': self.churn_model is not None,
            'feature_importance_available': self.feature_importance is not None
        }

        if self.cohort_matrix is not None:
            # Cohort analysis metrics
            summary['first_period_retention'] = self.cohort_matrix.iloc[:, 0].mean()
            summary['third_period_retention'] = self.cohort_matrix.iloc[:, 2].mean() if len(self.cohort_matrix.columns) > 2 else None

        if self.feature_importance is not None:
            # Top features
            summary['top_5_features'] = self.feature_importance.head(5)['feature'].tolist()

        return summary

    def generate_retention_insights(self) -> List[str]:
        """
        Generate actionable retention insights.

        Returns:
            List of insight strings
        """
        insights = []

        if self.cohort_matrix is not None:
            # Cohort insights
            first_period_retention = self.cohort_matrix.iloc[:, 0].mean()
            insights.append(f"📊 Initial period retention: {first_period_retention:.1%}")

            if len(self.cohort_matrix.columns) > 2:
                third_period_retention = self.cohort_matrix.iloc[:, 2].mean()
                retention_drop = first_period_retention - third_period_retention
                insights.append(f"📈 Retention drop from period 1 to 3: {retention_drop:.1%}")

                if retention_drop > 0.3:
                    insights.append("⚠️ High early churn - focus on onboarding and initial engagement")

            # Cohort comparison
            if len(self.cohort_matrix) > 1:
                latest_cohort = self.cohort_matrix.iloc[-1, 0]
                earliest_cohort = self.cohort_matrix.iloc[0, 0]
                if latest_cohort > earliest_cohort:
                    insights.append("📈 Improvement in initial retention over time")
                else:
                    insights.append("📉 Decline in initial retention - investigate recent changes")

        if self.feature_importance is not None:
            # Feature importance insights
            top_features = self.feature_importance.head(3)
            insights.append("🎯 Top churn predictors:")
            for _, row in top_features.iterrows():
                insights.append(f"   • {row['feature']}")

        if self.churn_model is not None:
            insights.append("✅ Churn prediction model available for proactive retention")

        return insights

    def export_retention_report(self, filename: str = 'retention_analysis_report.html') -> str:
        """
        Export comprehensive retention analysis report.

        Args:
            filename: Output filename

        Returns:
            Path to generated report file
        """
        # Generate insights
        insights = self.generate_retention_insights()
        summary = self.get_retention_summary()

        # Create HTML report
        html_content = f"""
        <!DOCTYPE html>
        <html>
        <head>
            <title>Retention Analysis Report</title>
            <style>
                body {{ font-family: Arial, sans-serif; margin: 40px; }}
                .header {{ background-color: #f0f0f0; padding: 20px; border-radius: 5px; }}
                .metrics {{ margin: 20px 0; }}
                .insights {{ background-color: #f9f9f9; padding: 20px; border-radius: 5px; }}
                table {{ border-collapse: collapse; width: 100%; }}
                th, td {{ border: 1px solid #ddd; padding: 8px; text-align: left; }}
                th {{ background-color: #f2f2f2; }}
            </style>
        </head>
        <body>
            <div class="header">
                <h1>Retention Analysis Report</h1>
                <p>Generated on {pd.Timestamp.now().strftime('%Y-%m-%d %H:%M:%S')}</p>
            </div>

            <div class="metrics">
                <h2>Key Metrics</h2>
                <ul>
                    <li><strong>Total Users:</strong> {summary['total_users']:,}</li>
                    <li><strong>Cohort Analysis:</strong> {'Available' if summary['cohort_matrix_available'] else 'Not Available'}</li>
                    <li><strong>Churn Model:</strong> {'Trained' if summary['model_trained'] else 'Not Trained'}</li>
        """

        if 'first_period_retention' in summary:
            html_content += f"<li><strong>Initial Retention:</strong> {summary['first_period_retention']:.1%}</li>"

        if 'third_period_retention' in summary and summary['third_period_retention']:
            html_content += f"<li><strong>Third Period Retention:</strong> {summary['third_period_retention']:.1%}</li>"

        if 'top_5_features' in summary:
            html_content += f"<li><strong>Top Churn Predictors:</strong> {', '.join(summary['top_5_features'][:3])}</li>"

        html_content += """
                </ul>
            </div>

            <div class="insights">
                <h2>Key Insights</h2>
                <ul>
        """

        for insight in insights:
            html_content += f"<li>{insight}</li>"

        html_content += """
                </ul>
            </div>
        </body>
        </html>
        """

        # Write to file
        with open(filename, 'w', encoding='utf-8') as f:
            f.write(html_content)

        print(f"Report saved to {filename}")
        return filename

"""
Survival Analysis Tools

This module provides comprehensive survival analysis capabilities including
Kaplan-Meier curves, Cox regression, and related statistical methods
for retention and churn analysis.
"""

import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from typing import Dict, List, Optional, Tuple, Union
from lifelines import KaplanMeierFitter, CoxPHFitter
from lifelines.statistics import logrank_test
import warnings
import matplotlib.patches as mpatches

class SurvivalAnalyzer:
    """
    Comprehensive survival analysis toolkit for retention analysis.

    Provides Kaplan-Meier survival analysis, Cox proportional hazards modeling,
    and statistical testing for retention data.
    """

    def __init__(self):
        self.data = None
        self.time_col = None
        self.event_col = None
        self.km_fit = None
        self.cox_fit = None
        self.groups = {}

        # Set up matplotlib for better plots
        plt.rcParams['font.size'] = 12
        plt.rcParams['axes.labelsize'] = 14
        plt.rcParams['axes.titlesize'] = 16
        plt.rcParams['xtick.labelsize'] = 12
        plt.rcParams['ytick.labelsize'] = 12
        plt.rcParams['legend.fontsize'] = 12

    def load_data(self, data: pd.DataFrame, time_col: str, event_col: str) -> None:
        """
        Load survival analysis data.

        Args:
            data: DataFrame with survival data
            time_col: Column name for time-to-event or censoring
            event_col: Column name for event indicator (1=event, 0=censored)
        """
        self.data = data.copy()
        self.time_col = time_col
        self.event_col = event_col

        # Validate data
        if time_col not in data.columns:
            raise ValueError(f"Time column '{time_col}' not found in data")
        if event_col not in data.columns:
            raise ValueError(f"Event column '{event_col}' not found in data")

        print(f"Loaded survival data with {len(data)} observations")
        print(f"Event rate: {data[event_col].mean():.1%}")
        print(f"Median {time_col}: {data[time_col].median():.1f}")

    def fit_kaplan_meier(self, data: Optional[pd.DataFrame] = None) -> KaplanMeierFitter:
        """
        Fit Kaplan-Meier survival curve.

        Args:
            data: Optional data to use (default: self.data)

        Returns:
            Fitted KaplanMeierFitter object
        """
        if data is None:
            data = self.data

        if data is None:
            raise ValueError("No data available. Use load_data() first.")

        self.km_fit = KaplanMeierFitter()
        self.km_fit.fit(data[self.time_col], data[self.event_col])

        return self.km_fit

    def plot_survival_curve(self, title: str = "Survival Curve",
                          show_ci: bool = True, figsize: Tuple[int, int] = (10, 6)) -> plt.Figure:
        """
        Plot Kaplan-Meier survival curve.

        Args:
            title: Plot title
            show_ci: Whether to show confidence intervals
            figsize: Figure size tuple

        Returns:
            Matplotlib figure object
        """
        if self.km_fit is None:
            raise ValueError("No Kaplan-Meier fit available. Use fit_kaplan_meier() first.")

        fig, ax = plt.subplots(figsize=figsize)

        # Plot survival curve
        self.km_fit.plot_survival_function(ax=ax, ci_show=show_ci,
                                          label='Survival Probability')

        # Customize plot
        ax.set_title(title, fontsize=16, fontweight='bold', pad=20)
        ax.set_xlabel(self.time_col.capitalize(), fontsize=14)
        ax.set_ylabel('Survival Probability', fontsize=14)
        ax.grid(True, alpha=0.3)
        ax.legend(loc='best')

        # Add median survival time if available
        if self.km_fit.median_survival_time_ != np.inf:
            ax.axvline(x=self.km_fit.median_survival_time_,
                      color='red', linestyle='--', alpha=0.7,
                      label=f'Median: {self.km_fit.median_survival_time_:.1f}')
            ax.legend()

        plt.tight_layout()
        return fig

    def compare_survival_groups(self, group_col: str, data: Optional[pd.DataFrame] = None) -> Dict[str, KaplanMeierFitter]:
        """
        Compare survival curves across different groups.

        Args:
            group_col: Column name for grouping variable
            data: Optional data to use (default: self.data)

        Returns:
            Dictionary with group names as keys and fitted KM objects as values
        """
        if data is None:
            data = self.data

        if group_col not in data.columns:
            raise ValueError(f"Group column '{group_col}' not found in data")

        self.groups = {}
        group_names = data[group_col].unique()

        for group in group_names:
            group_data = data[data[group_col] == group]
            km = KaplanMeierFitter()
            km.fit(group_data[self.time_col], group_data[self.event_col])
            self.groups[group] = km

        # Perform log-rank test if there are exactly 2 groups
        if len(group_names) == 2:
            group1_data = data[data[group_col] == group_names[0]]
            group2_data = data[data[group_col] == group_names[1]]

            results = logrank_test(group1_data[self.time_col], group2_data[self.time_col],
                                 group1_data[self.event_col], group2_data[self.event_col])

            print(f"Log-rank test p-value: {results.p_value:.4f}")
            if results.p_value < 0.05:
                print("Significant difference in survival between groups (p < 0.05)")
            else:
                print("No significant difference in survival between groups")

        return self.groups

    def plot_group_comparison(self, group_col: str, title: str = None,
                            figsize: Tuple[int, int] = (12, 8)) -> plt.Figure:
        """
        Plot survival curves for multiple groups.

        Args:
            group_col: Column name for grouping variable
            title: Plot title
            figsize: Figure size tuple

        Returns:
            Matplotlib figure object
        """
        if not self.groups:
            self.compare_survival_groups(group_col)

        fig, ax = plt.subplots(figsize=figsize)

        colors = plt.cm.Set3(np.linspace(0, 1, len(self.groups)))

        for i, (group, km_fit) in enumerate(self.groups.items()):
            km_fit.plot_survival_function(ax=ax, ci_show=False,
                                        label=str(group), color=colors[i])

        if title is None:
            title = f"Survival Curves by {group_col}"

        ax.set_title(title, fontsize=16, fontweight='bold', pad=20)
        ax.set_xlabel(self.time_col.capitalize(), fontsize=14)
        ax.set_ylabel('Survival Probability', fontsize=14)
        ax.grid(True, alpha=0.3)
        ax.legend(loc='best', bbox_to_anchor=(1.05, 1))

        plt.tight_layout()
        return fig

    def fit_cox_model(self, covariates: List[str], data: Optional[pd.DataFrame] = None) -> CoxPHFitter:
        """
        Fit Cox proportional hazards model.

        Args:
            covariates: List of covariate column names
            data: Optional data to use (default: self.data)

        Returns:
            Fitted CoxPHFitter object
        """
        if data is None:
            data = self.data

        # Prepare data for Cox model
        model_data = data[[self.time_col, self.event_col] + covariates].copy()

        # Handle categorical variables
        for col in covariates:
            if model_data[col].dtype == 'object':
                model_data = pd.get_dummies(model_data, columns=[col], drop_first=True)

        self.cox_fit = CoxPHFitter()
        self.cox_fit.fit(model_data, duration_col=self.time_col, event_col=self.event_col)

        return self.cox_fit

    def plot_cox_coefficients(self, title: str = "Cox Model Coefficients",
                            figsize: Tuple[int, int] = (10, 8)) -> plt.Figure:
        """
        Plot Cox regression coefficients with confidence intervals.

        Args:
            title: Plot title
            figsize: Figure size tuple

        Returns:
            Matplotlib figure object
        """
        if self.cox_fit is None:
            raise ValueError("No Cox model fitted. Use fit_cox_model() first.")

        # Get coefficients and confidence intervals
        coef_df = self.cox_fit.confidence_intervals_.copy()
        coef_df['coef'] = self.cox_fit.params_

        # Sort by coefficient value
        coef_df = coef_df.sort_values('coef')

        fig, ax = plt.subplots(figsize=figsize)

        # Plot coefficients
        y_pos = range(len(coef_df))

        # Error bars for confidence intervals
        lower_col = coef_df.columns[0]  # First column is lower bound
        upper_col = coef_df.columns[1]  # Second column is upper bound
        ax.errorbar(coef_df['coef'], y_pos,
                   xerr=[coef_df['coef'] - coef_df[lower_col],
                         coef_df[upper_col] - coef_df['coef']],
                   fmt='o', markersize=8, capsize=5)

        # Reference line at 0
        ax.axvline(x=0, color='red', linestyle='--', alpha=0.7)

        # Labels and formatting
        ax.set_yticks(y_pos)
        ax.set_yticklabels(coef_df.index)
        ax.set_xlabel('Hazard Ratio (log scale)', fontsize=14)
        ax.set_title(title, fontsize=16, fontweight='bold', pad=20)
        ax.grid(True, alpha=0.3, axis='x')

        # Add hazard ratio values
        for i, (index, row) in enumerate(coef_df.iterrows()):
            hr = np.exp(row['coef'])
            ax.text(row['coef'] + 0.05, i, f'HR: {hr:.2f}',
                   verticalalignment='center', fontsize=10)

        plt.tight_layout()
        return fig

    def get_summary_statistics(self) -> Dict:
        """
        Get comprehensive summary statistics.

        Returns:
            Dictionary with summary statistics
        """
        if self.data is None:
            raise ValueError("No data available. Use load_data() first.")

        stats = {
            'n_observations': len(self.data),
            'n_events': self.data[self.event_col].sum(),
            'event_rate': self.data[self.event_col].mean(),
            'median_time': self.data[self.time_col].median(),
            'mean_time': self.data[self.time_col].mean()
        }

        if self.km_fit:
            stats['median_survival'] = self.km_fit.median_survival_time_
            stats['survival_at_30_days'] = self.km_fit.predict(30).iloc[0] if hasattr(self.km_fit.predict(30), 'iloc') else self.km_fit.predict(30)
            stats['survival_at_90_days'] = self.km_fit.predict(90).iloc[0] if hasattr(self.km_fit.predict(90), 'iloc') else self.km_fit.predict(90)

        if self.cox_fit:
            stats['model_concordance'] = self.cox_fit.concordance_index_
            stats['log_likelihood'] = self.cox_fit.log_likelihood_

        return stats

    def generate_insights(self) -> List[str]:
        """
        Generate actionable insights from survival analysis.

        Returns:
            List of insight strings
        """
        insights = []

        if self.km_fit:
            # Overall retention insights
            if self.km_fit.median_survival_time_ != np.inf:
                insights.append(f"📊 Median user lifetime: {self.km_fit.median_survival_time_:.1f} days")
            else:
                insights.append("📊 More than 50% of users are still active (median survival not reached)")

            # Early retention insights
            surv_30 = self.km_fit.predict(30).iloc[0] if hasattr(self.km_fit.predict(30), 'iloc') else self.km_fit.predict(30)
            surv_90 = self.km_fit.predict(90).iloc[0] if hasattr(self.km_fit.predict(90), 'iloc') else self.km_fit.predict(90)

            insights.append(f"📈 30-day retention: {surv_30:.1%}")
            insights.append(f"📈 90-day retention: {surv_90:.1%}")

            if surv_30 < 0.5:
                insights.append("⚠️ Low early retention (<50% at 30 days) - focus on onboarding")

            if surv_90 < 0.2:
                insights.append("⚠️ Poor long-term retention (<20% at 90 days) - review engagement strategies")

        if self.cox_fit:
            # Identify top risk and protective factors
            coef = self.cox_fit.params_
            hr = np.exp(coef)

            # Top risk factors (HR > 1)
            risk_factors = hr[hr > 1].sort_values(ascending=False).head(3)
            if not risk_factors.empty:
                insights.append("🔺 Top risk factors for churn:")
                for factor, hr_val in risk_factors.items():
                    insights.append(f"   • {factor}: {hr_val:.1f}x higher risk")

            # Top protective factors (HR < 1)
            protective_factors = hr[hr < 1].sort_values().head(3)
            if not protective_factors.empty:
                insights.append("🔻 Factors that reduce churn risk:")
                for factor, hr_val in protective_factors.items():
                    reduction = (1 - hr_val) * 100
                    insights.append(f"   • {factor}: {reduction:.0f}% lower risk")

            # Model performance
            if self.cox_fit.concordance_index_ > 0.7:
                insights.append("✅ Strong predictive model (concordance > 0.7)")
            elif self.cox_fit.concordance_index_ > 0.6:
                insights.append("📊 Moderate predictive model (concordance > 0.6)")
            else:
                insights.append("⚠️ Weak predictive model - consider additional features")

        # Group comparison insights
        if len(self.groups) == 2:
            groups_list = list(self.groups.keys())
            surv_30_group1 = self.groups[groups_list[0]].predict(30).iloc[0] if hasattr(self.groups[groups_list[0]].predict(30), 'iloc') else self.groups[groups_list[0]].predict(30)
            surv_30_group2 = self.groups[groups_list[1]].predict(30).iloc[0] if hasattr(self.groups[groups_list[1]].predict(30), 'iloc') else self.groups[groups_list[1]].predict(30)

            if surv_30_group1 > surv_30_group2:
                better_group = groups_list[0]
                diff = surv_30_group1 - surv_30_group2
                insights.append(f"🏆 {better_group} shows {diff:.1%} better 30-day retention")
            else:
                better_group = groups_list[1]
                diff = surv_30_group2 - surv_30_group1
                insights.append(f"🏆 {better_group} shows {diff:.1%} better 30-day retention")

        return insights

    def export_report(self, filename: str = 'survival_analysis_report.html') -> str:
        """
        Export comprehensive survival analysis report.

        Args:
            filename: Output filename

        Returns:
            Path to generated report file
        """
        # Generate insights
        insights = self.generate_insights()
        stats = self.get_summary_statistics()

        # Create HTML report
        html_content = f"""
        <!DOCTYPE html>
        <html>
        <head>
            <title>Survival Analysis Report</title>
            <style>
                body {{ font-family: Arial, sans-serif; margin: 40px; }}
                .header {{ background-color: #f0f0f0; padding: 20px; border-radius: 5px; }}
                .metrics {{ margin: 20px 0; }}
                .insights {{ background-color: #f9f9f9; padding: 20px; border-radius: 5px; }}
                table {{ border-collapse: collapse; width: 100%; }}
                th, td {{ border: 1px solid #ddd; padding: 8px; text-align: left; }}
                th {{ background-color: #f2f2f2; }}
            </style>
        </head>
        <body>
            <div class="header">
                <h1>Survival Analysis Report</h1>
                <p>Generated on {pd.Timestamp.now().strftime('%Y-%m-%d %H:%M:%S')}</p>
            </div>

            <div class="metrics">
                <h2>Key Metrics</h2>
                <ul>
                    <li><strong>Total Observations:</strong> {stats['n_observations']:,}</li>
                    <li><strong>Events (Churn):</strong> {stats['n_events']:,} ({stats['event_rate']:.1%})</li>
                    <li><strong>Median Time:</strong> {stats['median_time']:.1f}</li>
        """

        if 'median_survival' in stats and stats['median_survival'] != np.inf:
            html_content += f"<li><strong>Median Survival:</strong> {stats['median_survival']:.1f}</li>"

        if 'survival_at_30_days' in stats:
            html_content += f"<li><strong>30-Day Survival:</strong> {stats['survival_at_30_days']:.1%}</li>"
            html_content += f"<li><strong>90-Day Survival:</strong> {stats['survival_at_90_days']:.1%}</li>"

        if 'model_concordance' in stats:
            html_content += f"<li><strong>Model Concordance:</strong> {stats['model_concordance']:.3f}</li>"

        html_content += """
                </ul>
            </div>

            <div class="insights">
                <h2>Key Insights</h2>
                <ul>
        """

        for insight in insights:
            html_content += f"<li>{insight}</li>"

        html_content += """
                </ul>
            </div>
        </body>
        </html>
        """

        # Write to file
        with open(filename, 'w', encoding='utf-8') as f:
            f.write(html_content)

        print(f"Report saved to {filename}")
        return filename

"""
Retention Analysis Visualization Tools

This module provides comprehensive visualization capabilities for retention analysis
including survival curves, cohort heatmaps, and churn prediction visualizations.
"""

import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
import plotly.graph_objects as go
import plotly.express as px
from plotly.subplots import make_subplots
from typing import Dict, List, Optional, Tuple

class RetentionVisualizer:
    """
    Visualization toolkit for retention analysis.

    Provides interactive and static visualizations for survival analysis,
    cohort analysis, and churn prediction results.
    """

    def __init__(self):
        # Set up matplotlib for better plots
        plt.rcParams['font.size'] = 12
        plt.rcParams['axes.labelsize'] = 14
        plt.rcParams['axes.titlesize'] = 16
        plt.rcParams['xtick.labelsize'] = 12
        plt.rcParams['ytick.labelsize'] = 12
        plt.rcParams['legend.fontsize'] = 12

        # Color palettes
        self.qualitative_colors = px.colors.qualitative.Set3
        self.sequential_colors = px.colors.sequential.Viridis

    def plot_retention_dashboard(self,
                               cohort_matrix: pd.DataFrame,
                               survival_data: Dict = None,
                               figsize: Tuple[int, int] = (16, 12)) -> plt.Figure:
        """
        Create comprehensive retention dashboard.

        Args:
            cohort_matrix: Cohort retention matrix
            survival_data: Dictionary with survival curve data
            figsize: Figure size tuple

        Returns:
            Matplotlib figure object
        """
        fig = plt.figure(figsize=figsize)

        # Create subplots
        gs = fig.add_gridspec(2, 3, hspace=0.3, wspace=0.3)

        # 1. Cohort heatmap
        ax1 = fig.add_subplot(gs[0, :2])
        sns.heatmap(cohort_matrix, annot=True, fmt='.0%', cmap='YlOrRd', ax=ax1)
        ax1.set_title('Cohort Retention Matrix', fontsize=16, fontweight='bold')
        ax1.set_xlabel('Period Number')
        ax1.set_ylabel('Cohort Period')

        # 2. Average retention by period
        ax2 = fig.add_subplot(gs[0, 2])
        avg_retention = cohort_matrix.mean()
        avg_retention.plot(kind='bar', ax=ax2, color=self.qualitative_colors[0])
        ax2.set_title('Average Retention by Period', fontsize=14, fontweight='bold')
        ax2.set_xlabel('Period Number')
        ax2.set_ylabel('Retention Rate')
        ax2.tick_params(axis='x', rotation=45)

        # 3. First period retention trend
        ax3 = fig.add_subplot(gs[1, 0])
        first_period = cohort_matrix.iloc[:, 0]
        first_period.plot(kind='line', marker='o', ax=ax3, color=self.qualitative_colors[1])
        ax3.set_title('Initial Retention Trend', fontsize=14, fontweight='bold')
        ax3.set_xlabel('Cohort Period')
        ax3.set_ylabel('Initial Retention')
        ax3.tick_params(axis='x', rotation=45)

        # 4. Retention decay curve
        ax4 = fig.add_subplot(gs[1, 1])
        avg_retention_decay = cohort_matrix.mean(axis=0)
        avg_retention_decay.plot(kind='line', marker='s', ax=ax4, color=self.qualitative_colors[2])
        ax4.set_title('Average Retention Decay', fontsize=14, fontweight='bold')
        ax4.set_xlabel('Period Number')
        ax4.set_ylabel('Retention Rate')

        # 5. Survival curve (if available)
        ax5 = fig.add_subplot(gs[1, 2])
        if survival_data and 'km_fit' in survival_data:
            survival_data['km_fit'].plot_survival_function(ax=ax5, ci_show=False)
            ax5.set_title('Survival Curve', fontsize=14, fontweight='bold')
            ax5.set_xlabel('Time')
            ax5.set_ylabel('Survival Probability')
        else:
            ax5.text(0.5, 0.5, 'No Survival\nData Available',
                    ha='center', va='center', transform=ax5.transAxes)
            ax5.set_title('Survival Curve', fontsize=14, fontweight='bold')

        plt.suptitle('Retention Analysis Dashboard', fontsize=20, fontweight='bold', y=0.98)
        return fig

    def plot_churn_prediction_results(self,
                                    y_true: np.array,
                                    y_pred: np.array,
                                    y_pred_proba: np.array,
                                    feature_importance: pd.DataFrame = None,
                                    figsize: Tuple[int, int] = (15, 10)) -> plt.Figure:
        """
        Visualize churn prediction results.

        Args:
            y_true: True labels
            y_pred: Predicted labels
            y_pred_proba: Prediction probabilities
            feature_importance: Feature importance DataFrame
            figsize: Figure size tuple

        Returns:
            Matplotlib figure object
        """
        fig = plt.figure(figsize=figsize)

        # Create subplots
        if feature_importance is not None:
            gs = fig.add_gridspec(2, 3, hspace=0.3, wspace=0.3)
        else:
            gs = fig.add_gridspec(2, 2, hspace=0.3, wspace=0.3)

        # 1. Confusion Matrix
        ax1 = fig.add_subplot(gs[0, 0])
        from sklearn.metrics import confusion_matrix
        cm = confusion_matrix(y_true, y_pred)
        sns.heatmap(cm, annot=True, fmt='d', cmap='Blues', ax=ax1)
        ax1.set_title('Confusion Matrix', fontsize=14, fontweight='bold')
        ax1.set_xlabel('Predicted')
        ax1.set_ylabel('Actual')

        # 2. ROC Curve
        ax2 = fig.add_subplot(gs[0, 1])
        from sklearn.metrics import roc_curve, auc
        fpr, tpr, _ = roc_curve(y_true, y_pred_proba)
        roc_auc = auc(fpr, tpr)
        ax2.plot(fpr, tpr, color='darkorange', lw=2, label=f'ROC curve (AUC = {roc_auc:.2f})')
        ax2.plot([0, 1], [0, 1], color='navy', lw=2, linestyle='--')
        ax2.set_xlim([0.0, 1.0])
        ax2.set_ylim([0.0, 1.05])
        ax2.set_xlabel('False Positive Rate')
        ax2.set_ylabel('True Positive Rate')
        ax2.set_title('ROC Curve', fontsize=14, fontweight='bold')
        ax2.legend(loc="lower right")

        # 3. Precision-Recall Curve
        ax3 = fig.add_subplot(gs[1, 0])
        from sklearn.metrics import precision_recall_curve, average_precision_score
        precision, recall, _ = precision_recall_curve(y_true, y_pred_proba)
        avg_precision = average_precision_score(y_true, y_pred_proba)
        ax3.plot(recall, precision, color='blue', lw=2,
                label=f'PR curve (AP = {avg_precision:.2f})')
        ax3.set_xlabel('Recall')
        ax3.set_ylabel('Precision')
        ax3.set_title('Precision-Recall Curve', fontsize=14, fontweight='bold')
        ax3.legend()

        # 4. Prediction Probability Distribution
        ax4 = fig.add_subplot(gs[1, 1])
        ax4.hist(y_pred_proba[y_true == 0], bins=30, alpha=0.7, label='Non-Churn', color='green')
        ax4.hist(y_pred_proba[y_true == 1], bins=30, alpha=0.7, label='Churn', color='red')
        ax4.set_xlabel('Predicted Churn Probability')
        ax4.set_ylabel('Frequency')
        ax4.set_title('Prediction Probability Distribution', fontsize=14, fontweight='bold')
        ax4.legend()

        # 5. Feature Importance (if available)
        if feature_importance is not None:
            ax5 = fig.add_subplot(gs[:, 2])
            top_features = feature_importance.head(10)
            ax5.barh(range(len(top_features)), top_features['importance'],
                    color='steelblue')
            ax5.set_yticks(range(len(top_features)))
            ax5.set_yticklabels(top_features['feature'])
            ax5.set_xlabel('Importance')
            ax5.set_title('Top Feature Importance', fontsize=14, fontweight='bold')

        plt.suptitle('Churn Prediction Results', fontsize=18, fontweight='bold')
        return fig

    def create_interactive_survival_curves(self,
                                        survival_groups: Dict[str, any],
                                        title: str = "Interactive Survival Curves") -> go.Figure:
        """
        Create interactive survival curves using Plotly.

        Args:
            survival_groups: Dictionary with group names and KM fits
            title: Chart title

        Returns:
            Plotly figure object
        """
        fig = go.Figure()

        colors = self.qualitative_colors[:len(survival_groups)]

        for i, (group_name, km_fit) in enumerate(survival_groups.items()):
            # Get survival function data
            survival_df = km_fit.survival_function_
            confidence_df = km_fit.confidence_interval_

            # Add main survival curve
            fig.add_trace(go.Scatter(
                x=survival_df.index,
                y=survival_df.iloc[:, 0],
                mode='lines',
                name=str(group_name),
                line=dict(color=colors[i], width=3)
            ))

            # Add confidence intervals
            fig.add_trace(go.Scatter(
                x=survival_df.index,
                y=confidence_df.iloc[:, 0],
                mode='lines',
                line=dict(width=0),
                showlegend=False,
                hoverinfo='skip'
            ))

            fig.add_trace(go.Scatter(
                x=survival_df.index,
                y=confidence_df.iloc[:, 1],
                mode='lines',
                line=dict(width=0),
                fill='tonexty',
                fillcolor=f'rgba{colors[i][4:-1]},0.2',
                showlegend=False,
                hoverinfo='skip'
            ))

        fig.update_layout(
            title={
                'text': title,
                'x': 0.5,
                'xanchor': 'center',
                'font': {'size': 20}
            },
            xaxis_title='Time',
            yaxis_title='Survival Probability',
            yaxis=dict(range=[0, 1]),
            hovermode='x unified',
            height=600
        )

        return fig

    def create_interactive_cohort_heatmap(self,
                                        cohort_matrix: pd.DataFrame,
                                        title: str = "Interactive Cohort Heatmap") -> go.Figure:
        """
        Create interactive cohort heatmap using Plotly.

        Args:
            cohort_matrix: Cohort retention matrix
            title: Chart title

        Returns:
            Plotly figure object
        """
        # Convert to percentage format for display
        display_matrix = cohort_matrix.round(3) * 100

        fig = go.Figure(data=go.Heatmap(
            z=display_matrix.values,
            x=[f"Period {i}" for i in display_matrix.columns],
            y=[str(cohort) for cohort in display_matrix.index],
            colorscale='YlOrRd',
            text=display_matrix.round(1).astype(str) + '%',
            texttemplate='%{text}',
            textfont={"size": 10},
            hoverongaps=False
        ))

        fig.update_layout(
            title={
                'text': title,
                'x': 0.5,
                'xanchor': 'center',
                'font': {'size': 20}
            },
            xaxis_title='Period Number',
            yaxis_title='Cohort Period',
            height=600
        )

        return fig

    def create_churn_risk_segments(self,
                                 churn_data: pd.DataFrame,
                                 title: str = "Churn Risk Segmentation") -> go.Figure:
        """
        Create churn risk segmentation visualization.

        Args:
            churn_data: DataFrame with churn predictions and risk categories
            title: Chart title

        Returns:
            Plotly figure object
        """
        # Count users by risk category
        risk_counts = churn_data['risk_category'].value_counts()

        fig = go.Figure(data=[
            go.Funnel(
                y=risk_counts.index,
                x=risk_counts.values,
                textinfo="value+percent initial",
                marker=dict(color=px.colors.sequential.RdYlBu_r[:len(risk_counts)])
            )
        ])

        fig.update_layout(
            title={
                'text': title,
                'x': 0.5,
                'xanchor': 'center',
                'font': {'size': 20}
            },
            height=500
        )

        return fig

    def create_retention_trend_analysis(self,
                                      retention_data: pd.DataFrame,
                                      time_col: str,
                                      retention_col: str,
                                      group_col: str = None,
                                      title: str = "Retention Trend Analysis") -> go.Figure:
        """
        Create retention trend analysis over time.

        Args:
            retention_data: DataFrame with retention metrics over time
            time_col: Time period column
            retention_col: Retention rate column
            group_col: Optional grouping column
            title: Chart title

        Returns:
            Plotly figure object
        """
        fig = go.Figure()

        if group_col:
            # Create separate lines for each group
            for group in retention_data[group_col].unique():
                group_data = retention_data[retention_data[group_col] == group]
                fig.add_trace(go.Scatter(
                    x=group_data[time_col],
                    y=group_data[retention_col],
                    mode='lines+markers',
                    name=str(group),
                    line=dict(width=3)
                ))
        else:
            # Single line for overall retention
            fig.add_trace(go.Scatter(
                x=retention_data[time_col],
                y=retention_data[retention_col],
                mode='lines+markers',
                name='Overall Retention',
                line=dict(width=3, color=self.qualitative_colors[0])
            ))

        fig.update_layout(
            title={
                'text': title,
                'x': 0.5,
                'xanchor': 'center',
                'font': {'size': 20}
            },
            xaxis_title=time_col.replace('_', ' ').title(),
            yaxis_title='Retention Rate',
            yaxis=dict(tickformat='.0%'),
            height=500
        )

        return fig

    def create_feature_importance_plot(self,
                                     feature_importance: pd.DataFrame,
                                     top_n: int = 15,
                                     title: str = "Feature Importance for Churn Prediction") -> go.Figure:
        """
        Create interactive feature importance plot.

        Args:
            feature_importance: DataFrame with features and importance scores
            top_n: Number of top features to display
            title: Chart title

        Returns:
            Plotly figure object
        """
        top_features = feature_importance.head(top_n)

        fig = go.Figure(data=[
            go.Bar(
                x=top_features['importance'],
                y=top_features['feature'],
                orientation='h',
                marker=dict(color=self.sequential_colors[::-1][:top_n])
            )
        ])

        fig.update_layout(
            title={
                'text': title,
                'x': 0.5,
                'xanchor': 'center',
                'font': {'size': 20}
            },
            xaxis_title='Importance Score',
            yaxis_title='Features',
            height=max(400, top_n * 30),
            yaxis={'categoryorder': 'total ascending'}
        )

        return fig

    def save_figure(self, fig, filename: str, format: str = 'html') -> str:
        """
        Save figure to file.

        Args:
            fig: Plotly or matplotlib figure object
            filename: Output filename
            format: Output format ('html', 'png', 'pdf', 'svg')

        Returns:
            Path to saved file
        """
        if hasattr(fig, 'write_html'):  # Plotly figure
            if format.lower() == 'html':
                fig.write_html(filename)
            elif format.lower() == 'png':
                fig.write_image(filename)
            elif format.lower() == 'pdf':
                fig.write_image(filename)
            elif format.lower() == 'svg':
                fig.write_image(filename)
        else:  # Matplotlib figure
            if format.lower() == 'png':
                fig.savefig(filename, dpi=300, bbox_inches='tight')
            elif format.lower() == 'pdf':
                fig.savefig(filename, bbox_inches='tight')
            elif format.lower() == 'svg':
                fig.savefig(filename, format='svg', bbox_inches='tight')

        print(f"Figure saved to {filename}")
        return filename