from sklearn.datasets import make_classification
from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split
import xgboost as xgb
import swanlab


def get_roc_pr_curve():
    X, y = make_classification(n_samples=1000, n_features=20, n_informative=2, n_redundant=10, random_state=42)

    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.3, random_state=42)

    model = xgb.XGBClassifier(use_label_encoder=False, eval_metric='logloss')
    model.fit(X_train, y_train)

    y_pred_proba = model.predict_proba(X_test)[:, 1]
    return y_test, y_pred_proba


def get_cm_mock():
    iris_data = load_iris()
    X = iris_data.data
    y = iris_data.target
    class_names = iris_data.target_names.tolist()

    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

    dtrain = xgb.DMatrix(X_train, label=y_train)
    dtest = xgb.DMatrix(X_test, label=y_test)

    param = {
        'objective': 'multi:softmax',
        'num_class': len(class_names),
        'eta': 0.1,
        'max_depth': 3,
        'eval_metric': 'mlogloss',
    }
    num_round = 100

    bst = xgb.train(param, dtrain, num_round)
    preds = bst.predict(dtest)
    return y_test, preds, class_names


if __name__ == '__main__':

    y_test, y_pred_proba = get_roc_pr_curve()
    cm_test, cm_pred, class_names = get_cm_mock()
    print(cm_test, cm_pred, class_names)

    swanlab.login(api_key="<API_KEY>", host="<DEV_HOST>")
    swanlab.init(project="Echarts-Metrics-Demo", experiment_name="PR&ROC Curve")

    swanlab.log(
        {
            "pr_curve": swanlab.pr_curve(y_test, y_pred_proba),
            "roc_curve": swanlab.roc_curve(y_test, y_pred_proba),
        }
    )

    swanlab.log({"confusion_matrix": swanlab.confusion_matrix(cm_test, cm_pred, class_names)})

    swanlab.finish()