import os
import tempfile
from collections import namedtuple
from typing import Optional, Tuple, Union

import numpy as np
import pytest

import xgboost as xgb
from xgboost import testing as tm
from xgboost.testing.callbacks import (
    run_eta_decay,
    run_eta_decay_leaf_output,
    tree_methods_objs,
)

# We use the dataset for tests.
pytestmark = pytest.mark.skipif(**tm.no_sklearn())


BreastCancer = namedtuple("BreastCancer", ["full", "tr", "va"])


@pytest.fixture
def breast_cancer() -> BreastCancer:
    from sklearn.datasets import load_breast_cancer

    X, y = load_breast_cancer(return_X_y=True)

    split = int(X.shape[0] * 0.8)
    return BreastCancer(
        full=(X, y),
        tr=(X[:split, ...], y[:split, ...]),
        va=(X[split:, ...], y[split:, ...]),
    )


def eval_error_metric(predt: np.ndarray, dtrain: xgb.DMatrix) -> Tuple[str, np.float64]:
    # No custom objective, recieve transformed output
    return tm.eval_error_metric(predt, dtrain, rev_link=False)


class TestCallbacks:
    def run_evaluation_monitor(
        self,
        D_train: xgb.DMatrix,
        D_valid: xgb.DMatrix,
        rounds: int,
        verbose_eval: Union[bool, int],
    ):
        def check_output(output: str) -> None:
            if int(verbose_eval) == 1:
                # Should print each iteration info
                assert len(output.split("\n")) == rounds
            elif int(verbose_eval) > rounds:
                # Should print first and latest iteration info
                assert len(output.split("\n")) == 2
            else:
                # Should print info by each period additionaly to first and latest
                # iteration
                num_periods = rounds // int(verbose_eval)
                # Extra information is required for latest iteration
                is_extra_info_required = num_periods * int(verbose_eval) < (rounds - 1)
                assert len(output.split("\n")) == (
                    1 + num_periods + int(is_extra_info_required)
                )

        evals_result: xgb.callback.TrainingCallback.EvalsLog = {}
        params = {"objective": "binary:logistic", "eval_metric": "error"}
        with tm.captured_output() as (out, err):
            xgb.train(
                params,
                D_train,
                evals=[(D_train, "Train"), (D_valid, "Valid")],
                num_boost_round=rounds,
                evals_result=evals_result,
                verbose_eval=verbose_eval,
            )
            output: str = out.getvalue().strip()
            check_output(output)

        with tm.captured_output() as (out, err):
            xgb.cv(params, D_train, num_boost_round=rounds, verbose_eval=verbose_eval)
            output = out.getvalue().strip()
            check_output(output)

    def test_evaluation_monitor(self, breast_cancer: BreastCancer) -> None:
        D_train = xgb.DMatrix(breast_cancer.tr[0], breast_cancer.tr[1])
        D_valid = xgb.DMatrix(breast_cancer.va[0], breast_cancer.va[1])
        evals_result = {}
        rounds = 10
        xgb.train(
            {"objective": "binary:logistic", "eval_metric": "error"},
            D_train,
            evals=[(D_train, "Train"), (D_valid, "Valid")],
            num_boost_round=rounds,
            evals_result=evals_result,
            verbose_eval=True,
        )
        assert len(evals_result["Train"]["error"]) == rounds
        assert len(evals_result["Valid"]["error"]) == rounds

        self.run_evaluation_monitor(D_train, D_valid, rounds, True)
        self.run_evaluation_monitor(D_train, D_valid, rounds, 2)
        self.run_evaluation_monitor(D_train, D_valid, rounds, 4)
        self.run_evaluation_monitor(D_train, D_valid, rounds, rounds + 1)

    def test_early_stopping(self, breast_cancer: BreastCancer) -> None:
        D_train = xgb.DMatrix(breast_cancer.tr[0], breast_cancer.tr[1])
        D_valid = xgb.DMatrix(breast_cancer.va[0], breast_cancer.va[1])
        evals_result = {}
        rounds = 30
        early_stopping_rounds = 5
        booster = xgb.train(
            {"objective": "binary:logistic", "eval_metric": "error"},
            D_train,
            evals=[(D_train, "Train"), (D_valid, "Valid")],
            num_boost_round=rounds,
            evals_result=evals_result,
            verbose_eval=True,
            early_stopping_rounds=early_stopping_rounds,
        )
        dump = booster.get_dump(dump_format="json")
        assert len(dump) - booster.best_iteration == early_stopping_rounds + 1

    def test_early_stopping_custom_eval(self, breast_cancer: BreastCancer) -> None:
        D_train = xgb.DMatrix(breast_cancer.tr[0], breast_cancer.tr[1])
        D_valid = xgb.DMatrix(breast_cancer.va[0], breast_cancer.va[1])
        early_stopping_rounds = 5
        booster = xgb.train(
            {
                "objective": "binary:logistic",
                "eval_metric": "error",
                "tree_method": "hist",
            },
            D_train,
            evals=[(D_train, "Train"), (D_valid, "Valid")],
            custom_metric=eval_error_metric,
            num_boost_round=1000,
            early_stopping_rounds=early_stopping_rounds,
            verbose_eval=False,
        )
        dump = booster.get_dump(dump_format="json")
        assert len(dump) - booster.best_iteration == early_stopping_rounds + 1

    def test_early_stopping_customize(self, breast_cancer: BreastCancer) -> None:
        D_train = xgb.DMatrix(breast_cancer.tr[0], breast_cancer.tr[1])
        D_valid = xgb.DMatrix(breast_cancer.va[0], breast_cancer.va[1])
        early_stopping_rounds = 5
        early_stop = xgb.callback.EarlyStopping(
            rounds=early_stopping_rounds, metric_name="CustomErr", data_name="Train"
        )
        # Specify which dataset and which metric should be used for early stopping.
        booster = xgb.train(
            {
                "objective": "binary:logistic",
                "eval_metric": ["error", "rmse"],
                "tree_method": "hist",
            },
            D_train,
            evals=[(D_train, "Train"), (D_valid, "Valid")],
            custom_metric=eval_error_metric,
            num_boost_round=1000,
            callbacks=[early_stop],
            verbose_eval=False,
        )
        dump = booster.get_dump(dump_format="json")
        assert len(dump) - booster.best_iteration == early_stopping_rounds + 1
        assert len(early_stop.stopping_history["Train"]["CustomErr"]) == len(dump)

        rounds = 100
        early_stop = xgb.callback.EarlyStopping(
            rounds=early_stopping_rounds,
            metric_name="CustomErr",
            data_name="Train",
            min_delta=100,
            save_best=True,
        )
        booster = xgb.train(
            {
                "objective": "binary:logistic",
                "eval_metric": ["error", "rmse"],
                "tree_method": "hist",
            },
            D_train,
            evals=[(D_train, "Train"), (D_valid, "Valid")],
            # No custom objective, transformed output
            custom_metric=eval_error_metric,
            num_boost_round=rounds,
            callbacks=[early_stop],
            verbose_eval=False,
        )
        # No iteration can be made with min_delta == 100
        assert booster.best_iteration == 0
        assert booster.num_boosted_rounds() == 1

    def test_early_stopping_skl(self, breast_cancer: BreastCancer) -> None:
        X, y = breast_cancer.full
        early_stopping_rounds = 5
        cls = xgb.XGBClassifier(
            early_stopping_rounds=early_stopping_rounds, eval_metric="error"
        )
        cls.fit(X, y, eval_set=[(X, y)])
        booster = cls.get_booster()
        dump = booster.get_dump(dump_format="json")
        assert len(dump) - booster.best_iteration == early_stopping_rounds + 1

    def test_early_stopping_custom_eval_skl(self, breast_cancer: BreastCancer) -> None:
        X, y = breast_cancer.full
        early_stopping_rounds = 5
        early_stop = xgb.callback.EarlyStopping(rounds=early_stopping_rounds)
        cls = xgb.XGBClassifier(
            eval_metric=tm.eval_error_metric_skl, callbacks=[early_stop]
        )
        cls.fit(X, y, eval_set=[(X, y)])
        booster = cls.get_booster()
        dump = booster.get_dump(dump_format="json")
        assert len(dump) - booster.best_iteration == early_stopping_rounds + 1

    def test_early_stopping_save_best_model(self, breast_cancer: BreastCancer) -> None:
        X, y = breast_cancer.full
        n_estimators = 100
        early_stopping_rounds = 5
        early_stop = xgb.callback.EarlyStopping(
            rounds=early_stopping_rounds, save_best=True
        )
        cls = xgb.XGBClassifier(
            n_estimators=n_estimators,
            eval_metric=tm.eval_error_metric_skl,
            callbacks=[early_stop],
        )
        cls.fit(X, y, eval_set=[(X, y)])
        booster = cls.get_booster()
        dump = booster.get_dump(dump_format="json")
        assert len(dump) == booster.best_iteration + 1

        early_stop = xgb.callback.EarlyStopping(
            rounds=early_stopping_rounds, save_best=True
        )
        cls = xgb.XGBClassifier(
            booster="gblinear",
            n_estimators=10,
            eval_metric=tm.eval_error_metric_skl,
            callbacks=[early_stop],
        )
        with pytest.raises(ValueError):
            cls.fit(X, y, eval_set=[(X, y)])

        # No error
        early_stop = xgb.callback.EarlyStopping(
            rounds=early_stopping_rounds, save_best=False
        )
        xgb.XGBClassifier(
            booster="gblinear",
            n_estimators=10,
            eval_metric=tm.eval_error_metric_skl,
            callbacks=[early_stop],
        ).fit(X, y, eval_set=[(X, y)])

    def test_early_stopping_continuation(self, breast_cancer: BreastCancer) -> None:
        X, y = breast_cancer.full

        early_stopping_rounds = 5
        early_stop = xgb.callback.EarlyStopping(
            rounds=early_stopping_rounds, save_best=True
        )
        cls = xgb.XGBClassifier(
            eval_metric=tm.eval_error_metric_skl, callbacks=[early_stop]
        )
        cls.fit(X, y, eval_set=[(X, y)])

        booster = cls.get_booster()
        assert booster.num_boosted_rounds() == booster.best_iteration + 1

        with tempfile.TemporaryDirectory() as tmpdir:
            path = os.path.join(tmpdir, "model.json")
            cls.save_model(path)
            cls = xgb.XGBClassifier()
            cls.load_model(path)
            assert cls._Booster is not None
            early_stopping_rounds = 3
            cls.set_params(
                eval_metric=tm.eval_error_metric_skl,
                early_stopping_rounds=early_stopping_rounds,
            )
            cls.fit(X, y, eval_set=[(X, y)])
            booster = cls.get_booster()
            assert (
                booster.num_boosted_rounds()
                == booster.best_iteration + early_stopping_rounds + 1
            )

    def test_early_stopping_multiple_metrics(self):
        from sklearn.datasets import make_classification

        X, y = make_classification(random_state=1994)
        # AUC approaches 1.0 real quick.
        clf = xgb.XGBClassifier(eval_metric=["logloss", "auc"], early_stopping_rounds=2)
        clf.fit(X, y, eval_set=[(X, y)])
        assert clf.best_iteration < 8
        assert clf.evals_result()["validation_0"]["auc"][-1] > 0.99

        clf = xgb.XGBClassifier(eval_metric=["auc", "logloss"], early_stopping_rounds=2)
        clf.fit(X, y, eval_set=[(X, y)])

        assert clf.best_iteration > 50
        assert clf.evals_result()["validation_0"]["auc"][-1] > 0.99

    @pytest.mark.parametrize("tree_method", ["hist", "approx", "exact"])
    def test_eta_decay(self, tree_method: str) -> None:
        dtrain, dtest = tm.load_agaricus(__file__)
        run_eta_decay(tree_method, dtrain, dtest, "cpu")

    @pytest.mark.parametrize("tree_method,objective", tree_methods_objs())
    def test_eta_decay_leaf_output(self, tree_method: str, objective: str) -> None:
        dtrain, dtest = tm.load_agaricus(__file__)
        run_eta_decay_leaf_output(tree_method, objective, dtrain, dtest, "cpu")

    def test_check_point(self, breast_cancer: BreastCancer) -> None:
        X, y = breast_cancer.full
        m = xgb.DMatrix(X, y)
        with tempfile.TemporaryDirectory() as tmpdir:
            check_point = xgb.callback.TrainingCheckPoint(
                directory=tmpdir, interval=1, name="model"
            )
            xgb.train(
                {"objective": "binary:logistic"},
                m,
                num_boost_round=10,
                verbose_eval=False,
                callbacks=[check_point],
            )
            for i in range(1, 10):
                assert os.path.exists(
                    os.path.join(
                        tmpdir,
                        f"model_{i}.{xgb.callback.TrainingCheckPoint.default_format}",
                    )
                )

            check_point = xgb.callback.TrainingCheckPoint(
                directory=tmpdir, interval=1, as_pickle=True, name="model"
            )
            xgb.train(
                {"objective": "binary:logistic"},
                m,
                num_boost_round=10,
                verbose_eval=False,
                callbacks=[check_point],
            )
            for i in range(1, 10):
                assert os.path.exists(os.path.join(tmpdir, "model_" + str(i) + ".pkl"))

    def test_callback_list(self) -> None:
        X, y = tm.data.get_california_housing()
        m = xgb.DMatrix(X, y)
        callbacks = [xgb.callback.EarlyStopping(rounds=10)]
        for i in range(4):
            xgb.train(
                {"objective": "reg:squarederror", "eval_metric": "rmse"},
                m,
                evals=[(m, "Train")],
                num_boost_round=1,
                verbose_eval=True,
                callbacks=callbacks,
            )
        assert len(callbacks) == 1

    def test_attribute_error(self, breast_cancer: BreastCancer) -> None:
        X, y = breast_cancer.full

        clf = xgb.XGBClassifier(n_estimators=8)
        clf.fit(X, y, eval_set=[(X, y)])

        with pytest.raises(AttributeError, match="early stopping is used"):
            clf.best_iteration

        with pytest.raises(AttributeError, match="early stopping is used"):
            clf.best_score

        booster = clf.get_booster()
        with pytest.raises(AttributeError, match="early stopping is used"):
            booster.best_iteration

        with pytest.raises(AttributeError, match="early stopping is used"):
            booster.best_score

    def test_preserve_order(self) -> None:
        """Test the ordering of the callbacks is preserved."""
        X, y, w = tm.make_regression(256, 16, False)
        fst_call: Optional[int] = None

        # If we use Python `set`, Cb1 is ordered before Cb2. This test makes sure Cb2 is
        # called before Cb1.
        class Cb2(xgb.callback.TrainingCallback):
            def before_iteration(self, model, epoch: int, evals_log) -> bool:
                nonlocal fst_call
                assert fst_call is None or fst_call == 2
                fst_call = 2
                return False

        class Cb1(xgb.callback.TrainingCallback):
            def before_iteration(self, model, epoch: int, evals_log) -> bool:
                assert fst_call == 2
                return False

        callbacks = [Cb2(), Cb1()]
        xgb.train({}, dtrain=xgb.QuantileDMatrix(X, y, weight=w), callbacks=callbacks)
