model.py

import pandas as pd
import numpy as np
from xgboost import XGBClassifier
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score, precision_score, recall_score
from data_generation import generate_synthetic_data
from feature_engineering import perform_feature_engineering

class FinancialRiskModel:
    def __init__(self):
        self.model = XGBClassifier(
            eval_metric='logloss',
            random_state=42,
            n_estimators=100,
            max_depth=4,
            learning_rate=0.1
        )
        self.feature_cols = None
        self.metrics = {}
        self.feature_importances = None
    
    def train(self, df_features):
        """
        Trains the XGBoost model on the provided dataset.
        df_features should not contain 'customer_id' but should contain 'missed_payment'
        """
        X = df_features.drop('missed_payment', axis=1)
        y = df_features['missed_payment']
        self.feature_cols = X.columns.tolist()
        
        X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42, stratify=y)
        
        self.model.fit(X_train, y_train)
        
        # Calculate metrics
        y_pred = self.model.predict(X_test)
        
        self.metrics = {
            'accuracy': accuracy_score(y_test, y_pred),
            'precision': precision_score(y_test, y_pred, zero_division=0),
            'recall': recall_score(y_test, y_pred, zero_division=0)
        }
        
        # Calculate feature importances
        importances = self.model.feature_importances_
        feature_importance_df = pd.DataFrame({
            'Feature': self.feature_cols,
            'Importance': importances
        }).sort_values(by='Importance', ascending=False)
        
        self.feature_importances = feature_importance_df
        
    def predict_risk(self, df_sample):
        """
        Predicts the probability of financial stress and assigns a risk tier.
        Returns probability and tier.
        """
        # Ensure only the correct features are passed
        X = df_sample[self.feature_cols]
        prob = self.model.predict_proba(X)[:, 1]
        
        return prob
        
    def get_risk_tier_and_action(self, probability, low_thresh=0.3, high_thresh=0.7, very_high_thresh=0.85):
        """
        Takes a probability and returns the tier and recommended intervention.
        """
        if probability < low_thresh:
            tier = "Low Risk"
            action = "No action required."
            message = "Good job maintaining your financial health! No actions needed at this time."
        elif probability < high_thresh:
            tier = "Medium Risk"
            action = "Send gentle reminder"
            message = "Hi! Just a friendly reminder about your upcoming payment. Let us know if you need any assistance managing your schedule."
        elif probability < very_high_thresh:
            tier = "High Risk"
            action = "Offer EMI restructuring or payment holiday"
            message = "Hello! We noticed times might be getting tight. We're here to help—would you like to explore EMI restructuring or a 1-month payment holiday?"
        else:
            tier = "Very High Risk"
            action = "Immediate Account Freeze & Escalation"
            message = "URGENT: Your account requires immediate attention to avoid severe default penalties. Please contact our support team immediately to resolve your outstanding balance."
            
        return tier, action, message

def get_trained_model(num_customers=1000):
    """
    Helper to instantly provide a trained model and the original dataset for the Streamlit app.
    """
    raw_df = generate_synthetic_data(num_customers)
    features_df = perform_feature_engineering(raw_df)
    
    # Needs customer ID back mapping, so we keep it separate
    training_df = features_df.copy()
    
    model = FinancialRiskModel()
    model.train(training_df)
    
    # Return the raw DataFrame (has customer_id) and the trained model
    return raw_df, features_df, model

if __name__ == "__main__":
    raw, features, mdl = get_trained_model(500)
    print("Metrics:", mdl.metrics)
    print("Top Features:\n", mdl.feature_importances.head())