//  [python]
def pre_train(self, lithology_numbers, model_name, fe=False, train_all=False, add_seq_features=True,
                  add_poly_features=False, add_log_transformation=False):
        data = pd.read_csv('train.csv', sep=';')
        data = data.drop(columns=['SGR'])

        if fe:
            X = data[list(data.columns[:-2]) + ['FORCE_2020_LITHOFACIES_CONFIDENCE']].copy()
            X = self._preprocess(X, train=True, add_seq_features=add_seq_features, add_poly_features=add_poly_features,
                                 add_log_transformation=add_log_transformation)
        else:
            X = pd.read_csv('data_fe_train.zip', compression='gzip')

        y = data['FORCE_2020_LITHOFACIES_LITHOLOGY'].copy()

        print(X.shape, y.shape)
        y = y.map(lithology_numbers)
        category_to_lithology = {y: x for x, y in lithology_numbers.items()}
        model = model_name.split('.')[0]

        clusterlist = pd.DataFrame(columns=['WELL', 'X_LOC', 'Y_LOC'])
        for i in X.WELL.unique():
            clusterlist = clusterlist.append(
                X[(X['WELL'] == i) & (X['X_LOC'] > 0) & (X['Y_LOC'] > 0)].iloc[0][['WELL', 'X_LOC', 'Y_LOC']])
        clusterlist = clusterlist.reset_index(drop=True)
        gmm = mixture.GMM(n_components=3, covariance_type='full', random_state=2).fit(
            clusterlist[['X_LOC', 'Y_LOC']])
        joblib.dump(gmm, 'cluster_model.joblib.gz', compress=('gzip', 3))

        labels = gmm.predict(clusterlist[['X_LOC', 'Y_LOC']]) + 1
        clusterlist['cluster'] = labels

        score_list = []
        A = np.load('penalty_matrix.npy')

        for i in range(1, 4):
            well_names = clusterlist[clusterlist['cluster'] == i].WELL
            X_tmp = X[X.WELL.isin(well_names)]
            y_tmp = y[X_tmp.index]
            X_tmp = X_tmp.drop(columns=['FORCE_2020_LITHOFACIES_CONFIDENCE', 'WELL'])
            if i == 1:
                X_tmp = X_tmp.drop(columns=['MUDWEIGHT', 'DCAL'])
            elif i == 2:
                X_tmp = X_tmp.drop(columns=['RMIC'])
            elif i == 3:
                X_tmp = X_tmp.drop(columns=['MUDWEIGHT', 'DCAL', 'RMIC'])

            if model == 'xgboost':
                X_tmp = X_tmp.drop(columns=['GROUP', 'FORMATION', 'BS'])

                if train_all:
                    weights = class_weight.compute_sample_weight(class_weight='balanced', y=y_tmp)
                    print('Training cluster {} using {} with input shape {} {}'.format(i, model, X_tmp.shape, '-' * 20))

                    clf = XGBClassifier(tree_method='gpu_hist', learning_rate=0.12,
                                        max_depth=10, min_child_weight=10, n_estimators=1000,
                                        seed=43, eval_metric='auc', objective="multi:softmax",
                                        colsample_bytree=0.9).fit(X_tmp, y_tmp, sample_weight=weights, verbose=False)
                else:
                    X_train, X_test, y_train, y_test = train_test_split(X_tmp, y_tmp, test_size=0.2,
                                                                        random_state=42, stratify=y_tmp)
                    print(
                        'Training cluster {} using {} with input shape {} {}'.format(i, model, X_train.shape, '-' * 20))

                    weights = class_weight.compute_sample_weight(class_weight='balanced', y=y_train)
                    clf = XGBClassifier(tree_method='gpu_hist', learning_rate=0.12,
                                        max_depth=10, min_child_weight=10, n_estimators=1000,
                                        seed=43, eval_metric='auc', objective="multi:softmax",
                                        colsample_bytree=0.9).fit(X_train, y_train, sample_weight=weights, verbose=True)
                    y_pred = clf.predict(X_test)
                    score = get_score(A, y_test.values, y_pred)
                    score_list.append(score)
                    print("\n score: {}".format(score))
                    print(classification_report(y_test, y_pred))
                    report = classification_report(y_test, y_pred, output_dict=True)
                    report = pd.DataFrame(report).transpose()
                    report.to_csv(model + '_classification_report' + str(i) + '.csv')
                    plot_confusion_matrix(y_true=np.vectorize(category_to_lithology.get)(y_test),
                                          y_pred=np.vectorize(category_to_lithology.get)(y_pred),
                                          classes=list(category_to_lithology.values()),
                                          title=model + '_confusion_matrix_' + str(i))
                ax = plot_importance(clf, title=model + '_feature_importance_' + str(i), max_num_features=15)
                ax.figure.tight_layout()
                ax.figure.savefig(model + '_feature_importance_' + str(i) + '.png')

            elif model == 'catboost':
                if i == 1:
                    X_tmp = X_tmp.drop(columns=['GR_rollingmean_30', 'GR_rollingmax_30', 'GR_rollingmin_30',
                                                'DTC_rollingmean_30', 'DTC_rollingmax_30', 'DTC_rollingmin_30'])

                best = {'eval_metric': 'TotalF1', 'depth': 10, 'l2_leaf_reg': 5, 'learning_rate': 0.15,
                        'loss_function': 'MultiClass', 'od_type': 'Iter', 'early_stopping_rounds': 100,
                        'task_type': 'GPU', 'verbose': 500, 'random_seed': 42}

                if i == 1:
                    best = {'eval_metric': 'TotalF1:use_weights=False', 'depth': 11, 'l2_leaf_reg': 2,
                            'learning_rate': 0.11679185273824155, 'fold_len_multiplier': 1.669990271921796,
                            'loss_function': 'MultiClass', 'od_type': 'Iter', 'od_wait': 25,
                            'task_type': 'GPU', 'verbose': 500, 'random_seed': 42}
                if i == 2:
                    best = {'eval_metric': 'TotalF1:use_weights=False', 'depth': 11, 'l2_leaf_reg': 7,
                            'learning_rate': 0.13207049403753848, 'fold_len_multiplier': 2.0345037506716057,
                            'loss_function': 'MultiClass', 'od_type': 'Iter', 'od_wait': 25, 'one_hot_max_size': 7,
                            'task_type': 'GPU', 'verbose': 500, 'random_seed': 42}
                # if i == 3:
                #     best = {'eval_metric': 'TotalF1:use_weights=False', 'depth': 12, 'l2_leaf_reg': 3,
                #             'learning_rate': 0.20557260217142548, 'fold_len_multiplier': 1.9159673725003985,
                #             'loss_function': 'MultiClass', 'od_type': 'Iter', 'od_wait': 25, 'one_hot_max_size': 11,
                #             'task_type': 'GPU', 'verbose': 500, 'random_seed': 42}

                cat_features = [s for s in X_tmp.columns if ('GROUP' in s) or ('BS' in s) or ('FORMATION' in s)]
                X_tmp[cat_features] = X_tmp[cat_features].astype(str)
                cat_features_index = [X_tmp.columns.get_loc(c) for c in cat_features if c in X_tmp]

                if train_all:
                    class_names = y_tmp.unique()
                    print('Training cluster {} using {} with input shape {} {}'.format(i, model, X_tmp.shape, '-' * 20))

                    weights = class_weight.compute_class_weight(class_weight='balanced', classes=class_names, y=y_tmp)
                    best['cat_features'] = cat_features_index
                    best['class_weights'] = weights
                    clf = CatBoostClassifier(**best).fit(X_tmp, y_tmp)
                else:
                    X_train, X_test, y_train, y_test = train_test_split(X_tmp, y_tmp, test_size=0.2,
                                                                        random_state=42, stratify=y_tmp)
                    class_names = y_train.unique()
                    print(
                        'Training cluster {} using {} with input shape {} {}'.format(i, model, X_train.shape, '-' * 20))

                    weights = class_weight.compute_class_weight(class_weight='balanced', classes=class_names, y=y_train)
                    best['cat_features'] = cat_features_index
                    best['class_weights'] = weights
                    clf = CatBoostClassifier(**best).fit(X_train, y_train)

                    y_pred = clf.predict(X_test)

                    tmp = X_test.join(pd.DataFrame(np.vectorize(category_to_lithology.get)(y_pred),
                                                   columns=['prediction'], index=X_test.index))
                    tmp = tmp.join(pd.DataFrame(np.vectorize(category_to_lithology.get)(y_test),
                                                columns=['FORCE_2020_LITHOFACIES_LITHOLOGY'], index=X_test.index))
                    tmp.to_csv('test_pred_' + str(i) + '.csv')

                    score = get_score(A, y_test.values, y_pred)
                    score_list.append(score)
                    print("\n score: {}".format(score))
                    print(classification_report(y_test, y_pred))
                    report = classification_report(y_test, y_pred, output_dict=True)
                    report = pd.DataFrame(report).transpose()
                    report.to_csv(model + '_classification_report' + str(i) + '.csv')
                    plot_confusion_matrix(y_true=np.vectorize(category_to_lithology.get)(y_test),
                                          y_pred=np.vectorize(category_to_lithology.get)(y_pred),
                                          classes=list(category_to_lithology.values()),
                                          title=model + '_confusion_matrix_' + str(i))
                plot_feature_importance(clf, title=model + '_feature_importance_' + str(i))

            joblib.dump(clf, model + '_' + str(i) + '.joblib.gz', compress=('gzip', 3))
            np.savetxt("scorelist.txt", score_list, delimiter=",", fmt='%s')