import pandas as pd
from catboost import CatBoostClassifier
from sklearn.metrics import confusion_matrix, accuracy_score, recall_score, precision_score, f1_score, roc_auc_score, roc_curve
import numpy as np
import shap

def calculate_metrics(Y_test, predictions, predictions_proba):
    # Calculate metrics
    auc = roc_auc_score(Y_test, predictions_proba)
    accuracy = accuracy_score(Y_test, predictions)
    recall = recall_score(Y_test, predictions)
    precision = precision_score(Y_test, predictions)
    f1 = f1_score(Y_test, predictions)
    
    # Calculate sensitivity (recall) and specificity
    tn, fp, fn, tp = confusion_matrix(Y_test, predictions).ravel()
    sensitivity = tp / (tp + fn)
    specificity = tn / (tn + fp)
    
    # Return metrics and confusion matrix elements in a dictionary
    return {"AUC": auc, "Accuracy": accuracy, "Sensitivity": sensitivity, 
            "Specificity": specificity, "Recall": recall, 
            "Precision": precision, "F1": f1, 
            "True Positives": tp, "True Negatives": tn, 
            "False Positives": fp, "False Negatives": fn}

def predict(df):
    # Prepare the data
    object_columns = []
    for col in df.columns:
        if df[col].dtype == 'object':
            object_columns.append(col)
            
    col_X = ['Model1','CityBranch',
         'HargaUnit','InstallmentContract','TLO Year',
         'All Risk Year', 'Tenor', 'DPPercentage', 'BPKBAmount', 'HargaOTR', 'RangeDPPercentage',
         'Assurance', 'TotalAP', 'DPTotal', 'ProgressStatus', 'AdditionalAmount', 'Premi', 'SalesType', 
         'LeasingType','Branch', 'Flag AA', 'Credit Year'
         ]
    
    categorical_columns = [col for col in object_columns if col in col_X]
    print(categorical_columns)
    X = df[col_X]
    Y = df['IsRepurchase']
    print(X.shape)
    
    # Initialize two CatBoostRegressor models
    model = CatBoostClassifier(
        iterations=150,
        depth=5,
        cat_features=categorical_columns,
        scale_pos_weight=len(Y[Y == 0]) / len(Y[Y == 1]),
        random_seed=0
    )
    
    # Train the model
    model.fit(X, Y)
    # Predict on both training and test data
    predictions_proba = model.predict_proba(X)[:, 1]
    fpr, tpr, thresholds = roc_curve(Y, predictions_proba)
    optimal_threshold = thresholds[np.argmax(tpr - fpr)]
    predictions = (model.predict_proba(X)[:, 1]>optimal_threshold).astype(int)

    # Create an explainer object
    explainer = shap.Explainer(model)

    # Calculate SHAP values for all features
    shap_values = explainer(X)
    
    # Don't visualize, just make table contains each feature and its shap value
    shap_df = pd.DataFrame(shap_values.values, columns=X.columns)
    shap_df = shap_df.abs().mean().sort_values(ascending=False).reset_index()
    shap_df.columns = ['feature', 'importance score']
    shap_df['importance score'] = shap_df['importance score']/shap_df['importance score'].sum() * 100
    print(shap_df)
    
    # Calculate and store metrics for this fold
    metrics = calculate_metrics(Y, predictions, predictions_proba)
    print(pd.DataFrame(metrics, index=[0]))
    
    # Create new df for predictions
    df['Prediction_IsRepurchase'] = predictions
    df['Prediction_ProbIsRepurchase'] = predictions_proba
    
    return(df)

def get_output_schema():
  return pd.DataFrame({
    'BillingNo' : prep_string(),
    'Prediction_IsRepurchase' : prep_int(),
    'Prediction_ProbIsRepurchase' : prep_decimal()
})