"""
Description: Includes the analyses for the predictive machine learning framework.
Including train/test data splitting, model training and model performance metric calculations.
The framework is model-agnostic and the balanced random forest can be replaced with other
preferrend methods. Run through the main.py file
Author: Susanne Jauhiainen
Created: 2026-05-21
"""

import numpy as np
import sklearn
from sklearn.model_selection import StratifiedKFold
from sklearn.inspection import permutation_importance
from imblearn.ensemble import BalancedRandomForestClassifier
from sklearn.metrics import confusion_matrix
import warnings
from sklearn.impute import KNNImputer


#Train and test the classifier for the current repetition, rep
def train_and_test_rf(X, y, rep, folds, permutated, inner_folds):
    #Stratified k-fold data split with the repetition number as seed to ensure same divisions for true and random models
    kf = StratifiedKFold(n_splits=folds, shuffle=True, random_state=rep) 
    #Initialize lists for results saving
    res={"auc_test":[], "auc_train":[], "imps":[],"auc_pr_test":[],"auc_pr_train":[],"tpr":[],"tnr":[], "labels":[], "scores":[], "predicted":[], "precision":[]}
    #Go trought the folds, split train and test data
    for i, (train_index, test_index) in enumerate(kf.split(X, y)):
        X_train = X.iloc[train_index]
        y_train = y.iloc[train_index]
        X_test = X.iloc[test_index]
        y_test = y.iloc[test_index]
        
        #Impute missing data 
        imputer = KNNImputer(n_neighbors=5)
        imputer.fit(X_train)
        X_train = imputer.transform(X_train)
        X_test = imputer.transform(X_test)
        
        #Train the models, you can use any other methods here as well
        with warnings.catch_warnings(): #ignore future warnings from rf model
            warnings.filterwarnings("ignore")
            rf = BalancedRandomForestClassifier(random_state=rep)
            model = rf.fit(X_train,y_train)
                                
            
        #Get predictions and probabilities
        y_pred = model.predict(X_test)
        y_score = model.predict_proba(X_test)[:,1]
        y_score2 = model.predict_proba(X_train)[:,1]
        tn, fp, fn, tp = confusion_matrix(y_test, y_pred).ravel()
        tpr = tp/(tp+fn)
        tnr = tn/(tn+fp) 
        precision, recall, thresholds = sklearn.metrics.precision_recall_curve(y_test, y_score)

        r = permutation_importance(model, X_test, y_test, n_repeats=5, random_state=rep, scoring="recall") #"roc_auc"
        imps = r.importances_mean
        
        res["imps"].append(imps)
        res["auc_test"].append(sklearn.metrics.roc_auc_score(y_test,y_score))
        res["auc_train"].append(sklearn.metrics.roc_auc_score(y_train,y_score2))
        res["labels"].append(y_test)
        res["auc_pr_test"].append(sklearn.metrics.average_precision_score(y_test,y_score))
        res["auc_pr_train"].append(sklearn.metrics.average_precision_score(y_train,y_score2))
        res["tpr"].append(tpr)
        res["tnr"].append(tnr)
        res["predicted"].append(y_pred)
        res["scores"].append(y_score)
        res["precision"].append(precision)
    
    return res