import os

import numpy as np
import pandas as pd
from perpetual import PerpetualBooster


def test_multi_output():
    X = pd.read_csv(
        os.path.join(
            os.path.dirname(__file__), "../../../resources", "cover_types_train.csv"
        ),
        index_col=False,
    )
    X = X.sample(n=1000, random_state=0)
    y = np.array(X.pop("Cover_Type"))
    X_test = pd.read_csv(
        os.path.join(
            os.path.dirname(__file__), "../../../resources", "cover_types_test.csv"
        ),
        index_col=False,
    )
    y_test = np.array(X_test.pop("Cover_Type"))
    model = PerpetualBooster(iteration_limit=100, memory_limit=0.1, budget=0.1)
    model.fit(X, y)
    pred_test = model.predict(X_test)
    proba_test = model.predict_proba(X_test)
    log_odds_test = model.predict_log_proba(X_test)
    assert not np.isnan(pred_test).any()
    assert not np.isnan(proba_test).any()
    assert not np.isnan(log_odds_test).any()
    assert np.allclose(np.sum(proba_test, axis=1), np.ones(proba_test.shape[0]))
    assert np.allclose(proba_test.shape, (len(X_test), len(np.unique(y_test))))
    assert set(pred_test).issubset(set(y_test))
    assert len(set(pred_test)) >= 6

    feat_imp = model.calculate_feature_importance()
    assert isinstance(feat_imp, dict)
    assert np.allclose(np.sum(list(feat_imp.values())), 1.0)
