# ===================================================================
# MODULE 7, LESSON 2: SYNTHETIC DATA GENERATION FOR FINANCE
# ===================================================================

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
from scipy.stats import ks_2samp, wasserstein_distance
from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import train_test_split
from sklearn.metrics import roc_auc_score, classification_report
import warnings
warnings.filterwarnings('ignore')

# Check if SDV is installed
try:
    from sdv.tabular import CTGAN
    from sdv.metadata import SingleTableMetadata
    SDV_AVAILABLE = True
except ImportError:
    SDV_AVAILABLE = False
    print("SDV not installed. Install with: pip install sdv")
    print("Using a simplified synthetic generation for demonstration.")

# Set style
sns.set_style("whitegrid")
np.random.seed(42)

print("="*70)
print("SYNTHETIC DATA GENERATION FOR FINANCIAL APPLICATIONS")
print("="*70)

# ----------------------------------------------------------------
# PART A: REAL DATA – CREDIT DATASET
# ----------------------------------------------------------------

print("\n" + "-"*60)
print("PART A: Real Credit Data")
print("-"*60)

# Generate synthetic "real" credit data (we'll treat this as our ground truth)
def generate_credit_data(n=5000):
    """Generate realistic credit data."""
    np.random.seed(42)
    df = pd.DataFrame({
        'age': np.random.normal(45, 12, n).clip(22, 75).astype(int),
        'income': np.random.gamma(5, 15, n) + 20,
        'dti': np.random.beta(2, 5, n) * 60,
        'credit_score': np.random.normal(700, 50, n).clip(550, 850).astype(int),
        'loan_amount': np.random.gamma(4, 50, n) + 30,
        'employment_years': np.random.gamma(3, 5, n).clip(0, 30).astype(int),
        'home_owner': np.random.binomial(1, 0.65, n),
        'marital_status': np.random.choice(['Single', 'Married', 'Other'], n, p=[0.4, 0.4, 0.2]),
        'education': np.random.choice(['High School', 'College', 'Bachelors', 'Post-Grad'], n, p=[0.2, 0.3, 0.3, 0.2]),
        'default': np.random.binomial(1, 0.05, n)  # 5% default rate
    })
    # Add correlation: higher credit_score -> lower default
    # We'll do this by adjusting default based on credit score
    prob_default = 0.1 - 0.0001 * (df['credit_score'] - 550)
    prob_default = prob_default.clip(0.01, 0.3)
    df['default'] = np.random.binomial(1, prob_default)
    return df

# Generate "real" data
real_df = generate_credit_data(10000)
print(f"Real data shape: {real_df.shape}")
print(f"Default rate: {real_df['default'].mean():.4f}")

# ----------------------------------------------------------------
# PART B: SYNTHETIC DATA GENERATION (CTGAN)
# ----------------------------------------------------------------

print("\n" + "-"*60)
print("PART B: Synthetic Data Generation")
print("-"*60)

if SDV_AVAILABLE:
    print("Using CTGAN for synthetic data generation...")
    
    # Create metadata
    metadata = SingleTableMetadata()
    metadata.detect_from_dataframe(real_df)
    
    # Train CTGAN
    ctgan = CTGAN(epochs=100, batch_size=500, verbose=False)
    ctgan.fit(real_df)
    
    # Generate synthetic data
    synthetic_df = ctgan.sample(num_rows=len(real_df))
    print(f"Synthetic data shape: {synthetic_df.shape}")
    
    # Save models
    ctgan.save('ctgan_model.pkl')
    print("CTGAN model saved.")
    
else:
    print("CTGAN not available. Using a simpler approach (bootstrapping with noise).")
    # Generate synthetic data by bootstrapping with noise
    synthetic_df = real_df.sample(n=len(real_df), replace=True).reset_index(drop=True)
    # Add noise to continuous features
    for col in ['age', 'income', 'dti', 'credit_score', 'loan_amount', 'employment_years']:
        noise = np.random.normal(0, synthetic_df[col].std() * 0.05, len(synthetic_df))
        synthetic_df[col] = synthetic_df[col] + noise
        synthetic_df[col] = synthetic_df[col].clip(synthetic_df[col].min(), synthetic_df[col].max())
    
    print("Synthetic data generated using bootstrapping with noise.")

# ----------------------------------------------------------------
# PART C: STATISTICAL EVALUATION
# ----------------------------------------------------------------

print("\n" + "-"*60)
print("PART C: Statistical Evaluation")
print("-"*60)

# Compare distributions
def compare_distributions(real, synthetic, columns):
    """Compare real and synthetic distributions."""
    results = []
    for col in columns:
        if col in real.columns and col in synthetic.columns:
            # Kolmogorov-Smirnov test
            ks_stat, ks_p = ks_2samp(real[col].dropna(), synthetic[col].dropna())
            # Wasserstein distance
            wass_dist = wasserstein_distance(real[col].dropna(), synthetic[col].dropna())
            results.append({
                'Feature': col,
                'KS Statistic': ks_stat,
                'KS P-value': ks_p,
                'Wasserstein Distance': wass_dist
            })
    return pd.DataFrame(results)

# Select continuous features
continuous_features = ['age', 'income', 'dti', 'credit_score', 'loan_amount', 'employment_years']
comparison_df = compare_distributions(real_df, synthetic_df, continuous_features)
print("\nDistribution Comparison (Real vs Synthetic):")
print(comparison_df.to_string(index=False))

# Visualise distributions
fig, axes = plt.subplots(2, 3, figsize=(15, 10))

for i, col in enumerate(continuous_features):
    ax = axes[i // 3, i % 3]
    ax.hist(real_df[col], bins=30, alpha=0.5, label='Real', color='blue')
    ax.hist(synthetic_df[col], bins=30, alpha=0.5, label='Synthetic', color='green')
    ax.set_xlabel(col)
    ax.set_ylabel('Frequency')
    ax.set_title(f'{col}')
    ax.legend()
    ax.grid(True, alpha=0.3)

plt.tight_layout()
plt.savefig('distribution_comparison.png', dpi=300)
plt.show()

# Correlation comparison
fig, axes = plt.subplots(1, 2, figsize=(14, 6))

# Real correlation
ax = axes[0]
corr_real = real_df[continuous_features + ['default']].corr()
sns.heatmap(corr_real, annot=True, fmt='.2f', cmap='coolwarm', ax=ax, cbar=False)
ax.set_title('Real Data Correlation')

# Synthetic correlation
ax = axes[1]
corr_synth = synthetic_df[continuous_features + ['default']].corr()
sns.heatmap(corr_synth, annot=True, fmt='.2f', cmap='coolwarm', ax=ax, cbar=False)
ax.set_title('Synthetic Data Correlation')

plt.tight_layout()
plt.savefig('correlation_comparison.png', dpi=300)
plt.show()

# Correlation difference
corr_diff = np.abs(corr_real - corr_synth).mean().mean()
print(f"\nAverage correlation difference: {corr_diff:.4f}")

# ----------------------------------------------------------------
# PART D: UTILITY EVALUATION – TRAIN-ON-SYNTHETIC, TEST-ON-REAL
# ----------------------------------------------------------------

print("\n" + "-"*60)
print("PART D: Utility Evaluation (Train-on-Synthetic, Test-on-Real)")
print("-"*60)

# Prepare features
features = continuous_features + ['home_owner', 'marital_status', 'education']
# Encode categorical variables
def encode_categorical(df):
    df_enc = df.copy()
    # One-hot encode categorical variables
    categorical_cols = ['marital_status', 'education']
    for col in categorical_cols:
        dummies = pd.get_dummies(df[col], prefix=col)
        df_enc = pd.concat([df_enc, dummies], axis=1)
        df_enc = df_enc.drop(columns=[col])
    # Home_owner is already binary
    return df_enc

# Encode both datasets
real_enc = encode_categorical(real_df)
synthetic_enc = encode_categorical(synthetic_df)

# Get feature columns (after encoding)
X_cols = [col for col in real_enc.columns if col != 'default']
y_col = 'default'

# Split real data for final testing
X_real, X_test, y_real, y_test = train_test_split(real_enc[X_cols], real_enc[y_col], test_size=0.3, random_state=42)

# Train on synthetic, test on real
print("Training on Synthetic, Testing on Real...")
model_syn = RandomForestClassifier(n_estimators=100, max_depth=8, random_state=42)
model_syn.fit(synthetic_enc[X_cols], synthetic_enc[y_col])
y_pred_syn = model_syn.predict_proba(X_test)[:, 1]
auc_syn = roc_auc_score(y_test, y_pred_syn)
print(f"AUC (Train Synthetic, Test Real): {auc_syn:.4f}")

# Train on real, test on real (baseline)
print("\nTraining on Real, Testing on Real (Baseline)...")
model_real = RandomForestClassifier(n_estimators=100, max_depth=8, random_state=42)
model_real.fit(X_real, y_real)
y_pred_real = model_real.predict_proba(X_test)[:, 1]
auc_real = roc_auc_score(y_test, y_pred_real)
print(f"AUC (Train Real, Test Real): {auc_real:.4f}")

# Utility ratio
utility_ratio = auc_syn / auc_real
print(f"\nUtility Ratio (Synthetic/Real): {utility_ratio:.4f}")
if utility_ratio > 0.9:
    print("  ✓ Excellent utility – synthetic data is a good substitute.")
elif utility_ratio > 0.7:
    print("  ✓ Good utility – synthetic data can be used for development.")
else:
    print("  ⚠ Limited utility – synthetic data may need improvement.")

# ----------------------------------------------------------------
# PART E: PRIVACY EVALUATION – MEMBER INFERENCE ATTACK
# ----------------------------------------------------------------

print("\n" + "-"*60)
print("PART E: Privacy Evaluation – Member Inference Attack")
print("-"*60)

def member_inference_attack(real_df, synthetic_df, model):
    """
    Simplified member inference attack.
    If the model can distinguish between real and synthetic data, privacy may be compromised.
    """
    # Create combined dataset with labels: 0 = synthetic, 1 = real
    combined = pd.concat([
        synthetic_df.assign(label=0),
        real_df.assign(label=1)
    ]).sample(frac=1, random_state=42).reset_index(drop=True)
    
    # Use the same features as before
    X_attack = combined[X_cols]
    y_attack = combined['label']
    
    # Train-test split
    X_attack_train, X_attack_test, y_attack_train, y_attack_test = train_test_split(
        X_attack, y_attack, test_size=0.3, random_state=42
    )
    
    # Train a classifier to distinguish real from synthetic
    attack_model = RandomForestClassifier(n_estimators=100, random_state=42)
    attack_model.fit(X_attack_train, y_attack_train)
    
    # Evaluate
    y_attack_pred = attack_model.predict_proba(X_attack_test)[:, 1]
    auc_attack = roc_auc_score(y_attack_test, y_attack_pred)
    
    return auc_attack

# Run attack
auc_attack = member_inference_attack(real_enc, synthetic_enc, None)
print(f"Member Inference Attack AUC: {auc_attack:.4f}")
print(f"Interpretation: {auc_attack - 0.5:.4f} above random (0.5)")
if auc_attack < 0.6:
    print("  ✓ Good privacy – synthetic data does not leak real data.")
elif auc_attack < 0.7:
    print("  ⚠ Moderate privacy – some risk of re-identification.")
else:
    print("  ✗ Poor privacy – synthetic data resembles real data too closely.")

# ----------------------------------------------------------------
# PART F: SYNTHETIC DATA FOR IMBALANCE MITIGATION
# ----------------------------------------------------------------

print("\n" + "-"*60)
print("PART F: Using Synthetic Data to Address Class Imbalance")
print("-"*60)

# Check class imbalance in real data
default_rate = real_df['default'].mean()
print(f"Default rate in real data: {default_rate:.4f} ({default_rate*100:.2f}%)")

# Generate more synthetic samples for the minority class (default)
print("Generating additional synthetic data for the minority class...")

if SDV_AVAILABLE:
    # Generate more synthetic data
    extra_synthetic = ctgan.sample(num_rows=5000)
    # Identify the default class
    extra_defaults = extra_synthetic[extra_synthetic['default'] == 1]
    print(f"Generated {len(extra_defaults)} additional default samples.")
else:
    # Use bootstrapped approach
    default_samples = real_df[real_df['default'] == 1]
    extra_defaults = default_samples.sample(n=5000, replace=True).reset_index(drop=True)
    # Add noise
    for col in continuous_features:
        noise = np.random.normal(0, extra_defaults[col].std() * 0.05, len(extra_defaults))
        extra_defaults[col] = extra_defaults[col] + noise
    print(f"Generated {len(extra_defaults)} additional default samples (bootstrapped).")

# Combine with original data
augmented_df = pd.concat([real_df, extra_defaults], ignore_index=True)
new_default_rate = augmented_df['default'].mean()
print(f"New default rate: {new_default_rate:.4f} ({new_default_rate*100:.2f}%)")

# Train on augmented data
augmented_enc = encode_categorical(augmented_df)
X_aug = augmented_enc[X_cols]
y_aug = augmented_enc[y_col]

model_aug = RandomForestClassifier(n_estimators=100, max_depth=8, random_state=42)
model_aug.fit(X_aug, y_aug)
y_pred_aug = model_aug.predict_proba(X_test)[:, 1]
auc_aug = roc_auc_score(y_test, y_pred_aug)

print(f"\nAUC with augmented data: {auc_aug:.4f}")
print(f"Improvement over real-only: {auc_aug - auc_real:.4f}")

# ----------------------------------------------------------------
# PART G: REGULATORY PERSPECTIVE
# ----------------------------------------------------------------

print("\n" + "-"*60)
print("PART G: Regulatory Perspective on Synthetic Data")
print("-"*60)

print("""
Regulatory Stance on Synthetic Data:

1. GDPR (EU):
   - Synthetic data is NOT considered personal data if it cannot be re-identified.
   - Benefits: Can be shared freely, reducing compliance burden.
   - Risks: If synthetic data can be linked back to individuals, it becomes personal data.

2. EU AI Act:
   - Synthetic data is encouraged for model training (reduces privacy risks).
   - Must be properly validated and documented.

3. CCPA (California):
   - Similar to GDPR: Synthetic data is not personal data if it cannot be re-identified.

4. FDA Guidance:
   - Synthetic data can be used in clinical trials (drug discovery, device testing).
   - Must demonstrate that synthetic data is representative.

5. Banking Regulators (SR 11-7):
   - Synthetic data can be used for model development and testing.
   - Must be validated against real data.
   - Models trained on synthetic data still require validation on real data.

Best Practices:
  - Always evaluate synthetic data quality (statistical and utility).
  - Test for re-identification risk.
  - Document the generation process.
  - Use differential privacy for sensitive applications.
  - Combine synthetic and real data for best results.
""")

# ----------------------------------------------------------------
# PART H: SUMMARY AND RECOMMENDATIONS
# ----------------------------------------------------------------

print("\n" + "="*70)
print("PART H: Summary and Recommendations")
print("="*70)

print("""
Synthetic Data Generation – Key Takeaways:

1. Synthetic data solves privacy, scarcity, and imbalance challenges.
2. Methods: GANs (CTGAN), VAEs, Diffusion models, Statistical models.
3. Evaluate using: Statistical similarity (KS, Wasserstein), utility (AUC), privacy (member inference).
4. Utility Ratio = AUC(synthetic-trained) / AUC(real-trained); > 0.9 is excellent.
5. Synthetic data can significantly improve model performance on rare events.
6. CTGAN is the leading library for tabular financial data.
7. Differential privacy should be considered for highly sensitive data.
8. Regulatory acceptance is growing; document generation process thoroughly.

Recommendations:
  - Start with CTGAN for tabular financial data.
  - Evaluate synthetic data quality rigorously.
  - Use synthetic data for internal development and testing.
  - For regulatory submissions, validate models on real data.
  - Consider differential privacy for high-risk applications.
""")

print("="*70)
print("END OF LESSON 2 – MODULE 7")
print("="*70)