############################################
# imports
############################################

import pandas as pd
from sklearn.datasets import make_classification
from sklearn.ensemble import RandomForestClassifier
from fgclustering.forest_guided_clustering import *


############################################
# Tests
############################################


def test_forest_guided_clustering():
    # parameters
    max_K = 6
    method_clustering = "pam"
    init_clustering = "random"
    max_iter_clustering = 100
    discart_value_JI = 0.6
    bootstraps_JI = 100
    n_jobs = 3

    # test data
    X, y = make_classification(
        n_samples=300,
        n_features=10,
        n_informative=4,
        n_redundant=1,
        n_classes=2,
        n_clusters_per_class=1,
        random_state=42,
    )
    X = pd.DataFrame(X)

    model = RandomForestClassifier(
        max_depth=10,
        max_features="sqrt",
        max_samples=0.8,
        bootstrap=True,
        oob_score=True,
        random_state=42,
    )
    model.fit(X, y)

    data_classification = X
    data_classification["target"] = y

    # initialize and run fgclustering object
    fgc = FgClustering(model=model, data=data_classification, target_column="target")
    fgc.run(
        max_K=max_K,
        method_clustering=method_clustering,
        init_clustering=init_clustering,
        max_iter_clustering=max_iter_clustering,
        discart_value_JI=discart_value_JI,
        bootstraps_JI=bootstraps_JI,
        n_jobs=n_jobs,
    )

    # obtain optimal number of clusters
    assert fgc.k == 2, "error: wrong optimal k calculated"
