//  [python]
def train(self, seed_xgb, seed_tuned_cat,
              lithology_numbers, fe=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()
        y = y.map(lithology_numbers)

        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.GaussianMixture(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

        for i in range(1, 4):
            print('Training cluster {}'.format(i))

            cat_tuned_list = []
            xgb_list = []

            well_names = clusterlist[clusterlist['cluster'] == i].WELL
            X_train = X[X.WELL.isin(well_names)]
            X_train = X_train.drop(columns=['FORCE_2020_LITHOFACIES_CONFIDENCE', 'WELL'])
            y_train = y[X_train.index]

            label_encoder = LabelEncoder()
            label_encoder.fit(y_train)
            joblib.dump(label_encoder, 'label_encoder_' + str(i) + '.joblib.gz', compress=('gzip', 3))

            if i == 1:
                X_train = X_train.drop(columns=['MUDWEIGHT', 'DCAL'])
                best_xgb = {'colsample_bytree': 0.6457635688501269, 'gamma': 1,
                            'lambda': 0.9, 'learning_rate': 0.20408792946960846, 'max_depth': 15,
                            'min_child_weight': 4, 'subsample': 0.7,
                            'tree_method': 'gpu_hist', 'eval_metric': 'mlogloss', 'objective': 'multi:softmax',
                            'seed': 42}

                best_tuned_cat = {'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}

            elif i == 2:
                X_train = X_train.drop(columns=['RMIC'])

                best_xgb = {'colsample_bytree': 0.9231871602710936, 'eval_metric': 'mlogloss', 'gamma': 1,
                            'lambda': 0.9,
                            'learning_rate': 0.21381650119853535, 'max_depth': 16, 'min_child_weight': 5,
                            'num_class': 12,
                            'objective': 'multi:softmax', 'subsample': 0.7, 'tree_method': 'gpu_hist', 'seed': 42}

                best_tuned_cat = {'custom_metric': 'TotalF1:use_weights=True', 'depth': 12,
                                  'fold_len_multiplier': 1.4597904591216981,
                                  'l2_leaf_reg': 4, 'learning_rate': 0.2154607024950838, 'loss_function': 'MultiClass',
                                  'od_type': 'Iter',
                                  'od_wait': 25, 'one_hot_max_size': 7, 'task_type': 'GPU', 'verbose': 500,
                                  'random_seed': 42}


            elif i == 3:
                X_train = X_train.drop(columns=['MUDWEIGHT', 'DCAL', 'RMIC'])

                best_xgb = {'colsample_bytree': 0.986555257438249, 'eval_metric': 'mlogloss', 'gamma': 1,
                            'lambda': 0.93,
                            'learning_rate': 0.12228232369801341, 'max_depth': 19, 'min_child_weight': 4,
                            'num_class': 9,
                            'objective': 'multi:softmax', 'subsample': 0.7, 'tree_method': 'gpu_hist', 'seed': 42}

                best_tuned_cat = {'custom_metric': 'TotalF1:use_weights=True', 'depth': 12,
                                  'fold_len_multiplier': 1.0351256868651992, 'l2_leaf_reg': 3,
                                  'learning_rate': 0.24476582315993148, 'loss_function': 'MultiClass',
                                  'od_type': 'Iter', 'od_wait': 25, 'one_hot_max_size': 9,
                                  'task_type': 'GPU', 'verbose': 500, 'random_seed': 42}

            X_train_xgb = X_train.copy()
            X_train_cat = X_train.copy()
            X_train_xgb = X_train_xgb.drop(columns=['GROUP', 'FORMATION', 'BS'])

            sample_weights = class_weight.compute_sample_weight(class_weight='balanced', y=y_train)
            cat_features = [s for s in X_train_cat.columns if ('GROUP' in s) or ('BS' in s) or ('FORMATION' in s)]
            X_train_cat[cat_features] = X_train_cat[cat_features].astype(str)
            cat_features_index = [X_train_cat.columns.get_loc(c) for c in cat_features if c in X_train_cat]
            class_names = y_train.unique()
            class_weights = class_weight.compute_class_weight(class_weight='balanced', classes=class_names, y=y_train)

            version = 1
            for j in range(1, seed_xgb + 1):
                start_time = time.time()
                best_xgb['seed'] = j + 1000000
                best_xgb['num_class'] = len(y_train.unique())
                clf_xgb = XGBClassifier(**best_xgb).fit(X_train_xgb, y_train, sample_weight=sample_weights)
                xgb_list.append(clf_xgb)
                print('Training seed {} using xgboost with input shape {}, it took {:05.2f} seconds {}'.format(
                    j, X_train_xgb.shape, (time.time() - start_time), '-' * 20))
                if j % 5 == 0:
                    joblib.dump(xgb_list, "xgb_models_" + str(i) + "_" + str(version) + ".joblib.gz",
                                compress=('gzip', 5))
                    version = version + 1
                    xgb_list = []

            version = 1
            for p in range(1, seed_tuned_cat + 1):
                start_time = time.time()
                best_tuned_cat['cat_features'] = cat_features_index
                best_tuned_cat['class_weights'] = class_weights
                best_tuned_cat['random_seed'] = 100000000 + p
                clf_cat = CatBoostClassifier(**best_tuned_cat).fit(X_train_cat, y_train,
                                                                   cat_features=cat_features_index)
                cat_tuned_list.append(clf_cat)
                print('Training seed {} using tuned catboost with input shape {}, it took {:05.2f} seconds {}'.format(
                    1000000 + p, X_train_cat.shape, (time.time() - start_time), '-' * 20))
                if p % 5 == 0:
                    joblib.dump(cat_tuned_list, "cat_tuned_models_" + str(i) + "_" + str(version) + ".joblib.gz",
                                compress=('gzip', 5))
                    version = version + 1
                    cat_tuned_list = []