"""
Healthcare Patient Dataset - Advanced Analysis with Machine Learning
====================================================================
This script performs advanced analysis including:
- Correlation analysis
- Predictive modeling
- Statistical tests
- Machine learning predictions
- Advanced visualizations

Author: RSK World
Website: https://rskworld.in
Email: help@rskworld.in
Phone: +91 93305 39277
"""

import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from sklearn.model_selection import train_test_split
from sklearn.ensemble import RandomForestClassifier, RandomForestRegressor
from sklearn.preprocessing import LabelEncoder
from sklearn.metrics import accuracy_score, classification_report, r2_score, mean_squared_error
from scipy import stats
import warnings
warnings.filterwarnings('ignore')

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

def load_data(file_path='healthcare_patients.csv'):
    """
    Load and preprocess the dataset.
    
    Parameters:
    file_path (str): Path to CSV file
    
    Returns:
    pd.DataFrame: Preprocessed dataset
    """
    print("Loading dataset...")
    df = pd.read_csv(file_path)
    df['admission_date'] = pd.to_datetime(df['admission_date'])
    df['discharge_date'] = pd.to_datetime(df['discharge_date'])
    print(f"[OK] Dataset loaded: {len(df)} patients, {len(df.columns)} features")
    return df

def correlation_analysis(df):
    """
    Perform correlation analysis on numerical features.
    
    Parameters:
    df (pd.DataFrame): Patient dataset
    """
    print("\n" + "="*60)
    print("CORRELATION ANALYSIS")
    print("="*60)
    
    # Select numerical columns
    numerical_cols = ['age', 'length_of_stay', 'charges']
    corr_data = df[numerical_cols]
    correlation_matrix = corr_data.corr()
    
    print("\nCorrelation Matrix:")
    print(correlation_matrix)
    
    # Create correlation heatmap
    plt.figure(figsize=(10, 8))
    sns.heatmap(correlation_matrix, annot=True, cmap='coolwarm', center=0,
                square=True, linewidths=1, cbar_kws={"shrink": 0.8})
    plt.title('Correlation Matrix - Numerical Features', fontsize=16, fontweight='bold', pad=20)
    plt.tight_layout()
    plt.savefig('visualizations/correlation_matrix.png', dpi=300, bbox_inches='tight')
    print("\n[OK] Saved: visualizations/correlation_matrix.png")
    plt.close()
    
    # Find strong correlations
    print("\nStrong Correlations (|r| > 0.3):")
    for i in range(len(correlation_matrix.columns)):
        for j in range(i+1, len(correlation_matrix.columns)):
            corr_value = correlation_matrix.iloc[i, j]
            if abs(corr_value) > 0.3:
                print(f"  {correlation_matrix.columns[i]} <-> {correlation_matrix.columns[j]}: {corr_value:.3f}")

def statistical_tests(df):
    """
    Perform statistical tests on the data.
    
    Parameters:
    df (pd.DataFrame): Patient dataset
    """
    print("\n" + "="*60)
    print("STATISTICAL TESTS")
    print("="*60)
    
    # T-test: Compare charges by gender
    male_charges = df[df['gender'] == 'Male']['charges']
    female_charges = df[df['gender'] == 'Female']['charges']
    
    t_stat, p_value = stats.ttest_ind(male_charges, female_charges)
    print(f"\nT-test: Charges by Gender")
    print(f"  T-statistic: {t_stat:.4f}")
    print(f"  P-value: {p_value:.4f}")
    print(f"  Significant: {'Yes' if p_value < 0.05 else 'No'} (p < 0.05)")
    
    # T-test: Length of stay by outcome
    recovered_stay = df[df['outcome'] == 'Recovered']['length_of_stay']
    improved_stay = df[df['outcome'] == 'Improved']['length_of_stay']
    
    t_stat2, p_value2 = stats.ttest_ind(recovered_stay, improved_stay)
    print(f"\nT-test: Length of Stay by Outcome (Recovered vs Improved)")
    print(f"  T-statistic: {t_stat2:.4f}")
    print(f"  P-value: {p_value2:.4f}")
    print(f"  Significant: {'Yes' if p_value2 < 0.05 else 'No'} (p < 0.05)")
    
    # Chi-square test: Outcome vs Treatment Type
    contingency_table = pd.crosstab(df['outcome'], df['treatment_type'])
    chi2, p_value3, dof, expected = stats.chi2_contingency(contingency_table)
    print(f"\nChi-square Test: Outcome vs Treatment Type")
    print(f"  Chi-square: {chi2:.4f}")
    print(f"  P-value: {p_value3:.4f}")
    print(f"  Degrees of freedom: {dof}")
    print(f"  Significant: {'Yes' if p_value3 < 0.05 else 'No'} (p < 0.05)")

def predict_readmission(df):
    """
    Predict 30-day readmission using machine learning.
    
    Parameters:
    df (pd.DataFrame): Patient dataset
    """
    print("\n" + "="*60)
    print("READMISSION PREDICTION MODEL")
    print("="*60)
    
    # Prepare features
    feature_cols = ['age', 'length_of_stay', 'charges']
    X = df[feature_cols].copy()
    
    # Encode target
    le = LabelEncoder()
    y = le.fit_transform(df['readmission_30_days'])
    
    # Split data
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.3, random_state=42)
    
    # Train model
    model = RandomForestClassifier(n_estimators=100, random_state=42)
    model.fit(X_train, y_train)
    
    # Predictions
    y_pred = model.predict(X_test)
    accuracy = accuracy_score(y_test, y_pred)
    
    print(f"\nModel Performance:")
    print(f"  Accuracy: {accuracy:.4f} ({accuracy*100:.2f}%)")
    print(f"\nClassification Report:")
    # Get unique classes present in test set
    unique_test_classes = np.unique(y_test)
    unique_pred_classes = np.unique(y_pred)
    all_unique = np.unique(np.concatenate([y_test, y_pred]))
    target_names_list = [le.classes_[int(cls)] for cls in all_unique]
    print(classification_report(y_test, y_pred, labels=all_unique, target_names=target_names_list, zero_division=0))
    
    # Feature importance
    feature_importance = pd.DataFrame({
        'feature': feature_cols,
        'importance': model.feature_importances_
    }).sort_values('importance', ascending=False)
    
    print("\nFeature Importance:")
    for _, row in feature_importance.iterrows():
        print(f"  {row['feature']}: {row['importance']:.4f}")
    
    # Visualize feature importance
    plt.figure(figsize=(10, 6))
    plt.barh(feature_importance['feature'], feature_importance['importance'], color='steelblue')
    plt.xlabel('Importance', fontsize=12)
    plt.title('Feature Importance - Readmission Prediction', fontsize=16, fontweight='bold')
    plt.tight_layout()
    plt.savefig('visualizations/feature_importance.png', dpi=300, bbox_inches='tight')
    print("\n[OK] Saved: visualizations/feature_importance.png")
    plt.close()
    
    return model, feature_importance

def predict_charges(df):
    """
    Predict patient charges using regression.
    
    Parameters:
    df (pd.DataFrame): Patient dataset
    """
    print("\n" + "="*60)
    print("CHARGES PREDICTION MODEL")
    print("="*60)
    
    # Prepare features
    feature_cols = ['age', 'length_of_stay']
    X = df[feature_cols].copy()
    y = df['charges']
    
    # Split data
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.3, random_state=42)
    
    # Train model
    model = RandomForestRegressor(n_estimators=100, random_state=42)
    model.fit(X_train, y_train)
    
    # Predictions
    y_pred = model.predict(X_test)
    r2 = r2_score(y_test, y_pred)
    rmse = np.sqrt(mean_squared_error(y_test, y_pred))
    
    print(f"\nModel Performance:")
    print(f"  R² Score: {r2:.4f}")
    print(f"  RMSE: ${rmse:,.2f}")
    
    # Visualize predictions
    plt.figure(figsize=(12, 6))
    plt.scatter(y_test, y_pred, alpha=0.6, color='purple')
    plt.plot([y_test.min(), y_test.max()], [y_test.min(), y_test.max()], 'r--', lw=2)
    plt.xlabel('Actual Charges ($)', fontsize=12)
    plt.ylabel('Predicted Charges ($)', fontsize=12)
    plt.title(f'Charges Prediction (R² = {r2:.3f})', fontsize=16, fontweight='bold')
    plt.grid(True, alpha=0.3)
    plt.tight_layout()
    plt.savefig('visualizations/charges_prediction.png', dpi=300, bbox_inches='tight')
    print("\n[OK] Saved: visualizations/charges_prediction.png")
    plt.close()
    
    return model

def advanced_visualizations(df):
    """
    Create advanced visualizations.
    
    Parameters:
    df (pd.DataFrame): Patient dataset
    """
    print("\n" + "="*60)
    print("ADVANCED VISUALIZATIONS")
    print("="*60)
    
    # 1. Age vs Charges by Outcome
    plt.figure(figsize=(12, 6))
    for outcome in df['outcome'].unique():
        data = df[df['outcome'] == outcome]
        plt.scatter(data['age'], data['charges'], label=outcome, alpha=0.6, s=100)
    plt.xlabel('Age (years)', fontsize=12)
    plt.ylabel('Charges ($)', fontsize=12)
    plt.title('Age vs Charges by Outcome', fontsize=16, fontweight='bold')
    plt.legend()
    plt.grid(True, alpha=0.3)
    plt.tight_layout()
    plt.savefig('visualizations/age_charges_outcome.png', dpi=300, bbox_inches='tight')
    print("[OK] Saved: visualizations/age_charges_outcome.png")
    plt.close()
    
    # 2. Box plot: Charges by Treatment Type
    plt.figure(figsize=(12, 6))
    df.boxplot(column='charges', by='treatment_type', ax=plt.gca())
    plt.title('Charges Distribution by Treatment Type', fontsize=16, fontweight='bold')
    plt.suptitle('')  # Remove default title
    plt.xlabel('Treatment Type', fontsize=12)
    plt.ylabel('Charges ($)', fontsize=12)
    plt.xticks(rotation=45, ha='right')
    plt.tight_layout()
    plt.savefig('visualizations/charges_by_treatment.png', dpi=300, bbox_inches='tight')
    print("[OK] Saved: visualizations/charges_by_treatment.png")
    plt.close()
    
    # 3. Length of Stay Distribution by Outcome
    plt.figure(figsize=(12, 6))
    df.boxplot(column='length_of_stay', by='outcome', ax=plt.gca())
    plt.title('Length of Stay Distribution by Outcome', fontsize=16, fontweight='bold')
    plt.suptitle('')
    plt.xlabel('Outcome', fontsize=12)
    plt.ylabel('Length of Stay (days)', fontsize=12)
    plt.tight_layout()
    plt.savefig('visualizations/stay_by_outcome.png', dpi=300, bbox_inches='tight')
    print("[OK] Saved: visualizations/stay_by_outcome.png")
    plt.close()
    
    # 4. Heatmap: Average Charges by Diagnosis and Treatment
    pivot_data = df.pivot_table(values='charges', index='diagnosis', 
                                columns='treatment_type', aggfunc='mean')
    plt.figure(figsize=(14, 10))
    sns.heatmap(pivot_data, annot=True, fmt='.0f', cmap='YlOrRd', cbar_kws={'label': 'Average Charges ($)'})
    plt.title('Average Charges: Diagnosis vs Treatment Type', fontsize=16, fontweight='bold', pad=20)
    plt.xlabel('Treatment Type', fontsize=12)
    plt.ylabel('Diagnosis', fontsize=12)
    plt.xticks(rotation=45, ha='right')
    plt.yticks(rotation=0)
    plt.tight_layout()
    plt.savefig('visualizations/charges_heatmap.png', dpi=300, bbox_inches='tight')
    print("[OK] Saved: visualizations/charges_heatmap.png")
    plt.close()
    
    # 5. Time Series: Admissions over time
    df['admission_month'] = df['admission_date'].dt.to_period('M')
    monthly_admissions = df.groupby('admission_month').size()
    
    plt.figure(figsize=(14, 6))
    monthly_admissions.plot(kind='line', marker='o', linewidth=2, markersize=8, color='steelblue')
    plt.title('Monthly Patient Admissions', fontsize=16, fontweight='bold')
    plt.xlabel('Month', fontsize=12)
    plt.ylabel('Number of Admissions', fontsize=12)
    plt.grid(True, alpha=0.3)
    plt.xticks(rotation=45)
    plt.tight_layout()
    plt.savefig('visualizations/monthly_admissions.png', dpi=300, bbox_inches='tight')
    print("[OK] Saved: visualizations/monthly_admissions.png")
    plt.close()

def generate_advanced_report(df):
    """
    Generate comprehensive advanced analysis report.
    
    Parameters:
    df (pd.DataFrame): Patient dataset
    """
    import sys
    from io import StringIO
    
    old_stdout = sys.stdout
    sys.stdout = report = StringIO()
    
    correlation_analysis(df)
    statistical_tests(df)
    predict_readmission(df)
    predict_charges(df)
    
    report_content = report.getvalue()
    sys.stdout = old_stdout
    
    with open('advanced_analysis_report.txt', 'w', encoding='utf-8') as f:
        f.write("="*60 + "\n")
        f.write("ADVANCED ANALYSIS REPORT - HEALTHCARE PATIENT DATASET\n")
        f.write("="*60 + "\n")
        f.write(f"Generated on: {pd.Timestamp.now().strftime('%Y-%m-%d %H:%M:%S')}\n")
        f.write(f"Author: RSK World (https://rskworld.in)\n")
        f.write("="*60 + "\n\n")
        f.write(report_content)
    
    print("[OK] Advanced analysis report saved to: advanced_analysis_report.txt")

def main():
    """
    Main function to run advanced analysis.
    """
    print("\n" + "="*60)
    print("HEALTHCARE PATIENT DATASET - ADVANCED ANALYSIS")
    print("="*60)
    print("Author: RSK World")
    print("Website: https://rskworld.in")
    print("Email: help@rskworld.in")
    print("Phone: +91 93305 39277")
    print("="*60)
    
    # Load data
    df = load_data()
    
    # Create visualizations directory
    import os
    os.makedirs('visualizations', exist_ok=True)
    
    # Run analyses
    correlation_analysis(df)
    statistical_tests(df)
    predict_readmission(df)
    predict_charges(df)
    advanced_visualizations(df)
    generate_advanced_report(df)
    
    print("\n" + "="*60)
    print("ADVANCED ANALYSIS COMPLETE!")
    print("="*60)
    print("\nGenerated Files:")
    print("  - visualizations/ (advanced charts)")
    print("  - advanced_analysis_report.txt (comprehensive report)")
    print("\n")

if __name__ == "__main__":
    main()