import pandas as pd
import mlflow.catboost
import os
from catboost import CatBoostClassifier
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import LabelBinarizer
from load_data import uci_bank
from datetime import datetime

def main():

    data = uci_bank()
    df = data["df_train"]
    cat_col_idx = [df.columns.get_loc(c) for c in data["cat_cols"]]
    X = df.drop(data["target"], axis=1)
    y = LabelBinarizer().fit_transform(df[data["target"][0]])

    X_train, X_val, y_train, y_val = train_test_split(X, y, train_size=0.7, random_state=1337)
    model = CatBoostClassifier(iterations=100, learning_rate=0.05, depth=15, eval_metric="AUC")
    model.fit(X_train, y_train,
              cat_features=cat_col_idx,
              eval_set=(X_val, y_val))

    ts = datetime.now().strftime("%Y%m%d-%H%M%S")
    model_dir = "../dist/catboost-uci-bank-{}".format(ts)
    
    mlflow.catboost.save_model(model, model_dir)
    print("Model saved at {} !".format(os.path.abspath(model_dir)))

if __name__ == "__main__":
    main()
