import itertools
import json
import os
import shutil
from typing import Optional

import numpy as np
import pytest
import xgboost
from hypothesis import given, note, settings
from scipy.sparse import csr_matrix
from xgboost import testing as tm
from xgboost.testing.data import RelDataCV, simulate_clicks, sort_ltr_samples
from xgboost.testing.params import lambdarank_parameter_strategy
from xgboost.testing.ranking import run_normalization, run_score_normalization


def test_ndcg_custom_gain():
    def ndcg_gain(y: np.ndarray) -> np.ndarray:
        return np.exp2(y.astype(np.float64)) - 1.0

    X, y, q, w = tm.make_ltr(n_samples=1024, n_features=4, n_query_groups=3, max_rel=3)
    y_gain = ndcg_gain(y)

    byxgb = xgboost.XGBRanker(tree_method="hist", ndcg_exp_gain=True, n_estimators=10)
    byxgb.fit(
        X,
        y,
        qid=q,
        sample_weight=w,
        eval_set=[(X, y)],
        eval_qid=(q,),
        sample_weight_eval_set=(w,),
        verbose=True,
    )
    byxgb_json = json.loads(byxgb.get_booster().save_raw(raw_format="json"))

    bynp = xgboost.XGBRanker(tree_method="hist", ndcg_exp_gain=False, n_estimators=10)
    bynp.fit(
        X,
        y_gain,
        qid=q,
        sample_weight=w,
        eval_set=[(X, y_gain)],
        eval_qid=(q,),
        sample_weight_eval_set=(w,),
        verbose=True,
    )
    bynp_json = json.loads(bynp.get_booster().save_raw(raw_format="json"))

    # Remove the difference in parameter for comparison
    byxgb_json["learner"]["objective"]["lambdarank_param"]["ndcg_exp_gain"] = "0"
    assert byxgb.evals_result() == bynp.evals_result()
    assert byxgb_json == bynp_json

    # test pairwise can handle max_rel > 31, while ndcg metric is using custom gain
    X, y, q, w = tm.make_ltr(n_samples=1024, n_features=4, n_query_groups=3, max_rel=33)
    ranknet = xgboost.XGBRanker(
        tree_method="hist",
        ndcg_exp_gain=False,
        n_estimators=10,
        objective="rank:pairwise",
    )
    ranknet.fit(X, y, qid=q, eval_set=[(X, y)], eval_qid=[q])
    history = ranknet.evals_result()
    assert (
        history["validation_0"]["ndcg@32"][0] < history["validation_0"]["ndcg@32"][-1]
    )


def test_ndcg_non_exp() -> None:
    # NDCG exp gain must have label smaller than 32
    X, y, q, w = tm.make_ltr(n_samples=1024, n_features=4, n_query_groups=3, max_rel=44)

    def fit(ltr: xgboost.XGBRanker):
        ltr.fit(
            X,
            y,
            qid=q,
            sample_weight=w,
            eval_set=[(X, y)],
            eval_qid=(q,),
            sample_weight_eval_set=(w,),
        )

    ltr = xgboost.XGBRanker(tree_method="hist", ndcg_exp_gain=True, n_estimators=2)
    with pytest.raises(ValueError, match="Set `ndcg_exp_gain`"):
        fit(ltr)
    ltr = xgboost.XGBRanker(tree_method="hist", ndcg_exp_gain=False, n_estimators=2)
    fit(ltr)


def test_ranking_with_unweighted_data():
    # fmt: off
    Xrow = np.array([1, 2, 6, 8, 11, 14, 16, 17])
    Xcol = np.array([0, 0, 1, 1, 2, 2, 3, 3])
    y = np.array([
        0.0, 1.0, 1.0, 0.0, 0.0,
        0.0, 1.0, 0.0, 1.0, 0.0,
        0.0, 1.0, 0.0, 0.0, 1.0,
        0.0, 1.0, 1.0, 0.0, 0.0
    ])
    # fmt: on
    X = csr_matrix((np.ones(shape=8), (Xrow, Xcol)), shape=(20, 4))

    group = np.array([5, 5, 5, 5], dtype=np.uint)
    dtrain = xgboost.DMatrix(X, label=y)
    dtrain.set_group(group)

    params = {
        "eta": 1,
        "tree_method": "exact",
        "objective": "rank:pairwise",
        "eval_metric": ["auc", "aucpr"],
        "max_depth": 1,
    }
    evals_result = {}
    bst = xgboost.train(
        params, dtrain, 10, evals=[(dtrain, "train")], evals_result=evals_result
    )
    auc_rec = evals_result["train"]["auc"]
    assert all(p <= q for p, q in zip(auc_rec, auc_rec[1:]))
    auc_rec = evals_result["train"]["aucpr"]
    assert all(p <= q for p, q in zip(auc_rec, auc_rec[1:]))


def test_ranking_with_weighted_data():
    # fmt: off
    Xrow = np.array([1, 2, 6, 8, 11, 14, 16, 17])
    Xcol = np.array([0, 0, 1, 1, 2, 2, 3, 3])
    y = np.array([
        0.0, 1.0, 1.0, 0.0, 0.0,
        0.0, 1.0, 0.0, 1.0, 0.0,
        0.0, 1.0, 0.0, 0.0, 1.0,
        0.0, 1.0, 1.0, 0.0, 0.0
    ])
    # fmt: on
    X = csr_matrix((np.ones(shape=8), (Xrow, Xcol)), shape=(20, 4))
    weights = np.array([1.0, 2.0, 3.0, 4.0])

    group = np.array([5, 5, 5, 5], dtype=np.uint)
    dtrain = xgboost.DMatrix(X, label=y, weight=weights)
    dtrain.set_group(group)

    params = {
        "eta": 1,
        "tree_method": "exact",
        "objective": "rank:pairwise",
        "eval_metric": ["auc", "aucpr"],
        "max_depth": 1,
    }
    evals_result = {}
    bst = xgboost.train(
        params, dtrain, 10, evals=[(dtrain, "train")], evals_result=evals_result
    )
    auc_rec = evals_result["train"]["auc"]
    assert all(p <= q for p, q in zip(auc_rec, auc_rec[1:]))
    auc_rec = evals_result["train"]["aucpr"]
    assert all(p <= q for p, q in zip(auc_rec, auc_rec[1:]))

    for i in range(1, 11):
        pred = bst.predict(dtrain, iteration_range=(0, i))
        # is_sorted[i]: is i-th group correctly sorted by the ranking predictor?
        is_sorted = []
        for k in range(0, 20, 5):
            ind = np.argsort(-pred[k : k + 5])
            z = y[ind + k]
            is_sorted.append(all(i >= j for i, j in zip(z, z[1:])))
        # Since we give weights 1, 2, 3, 4 to the four query groups,
        # the ranking predictor will first try to correctly sort the last query group
        # before correctly sorting other groups.
        assert all(p <= q for p, q in zip(is_sorted, is_sorted[1:]))


def test_error_msg() -> None:
    X, y, qid, w = tm.make_ltr(10, 2, 2, 2)
    ranker = xgboost.XGBRanker()
    with pytest.raises(ValueError, match=r"equal to the number of query groups"):
        ranker.fit(X, y, qid=qid, sample_weight=y)


@given(lambdarank_parameter_strategy)
@settings(deadline=None, print_blob=True, max_examples=10)
def test_lambdarank_parameters(params):
    if params["objective"] == "rank:map":
        rel = 1
    else:
        rel = 4
    X, y, q, w = tm.make_ltr(4096, 3, 13, rel)
    ranker = xgboost.XGBRanker(tree_method="hist", n_estimators=64, **params)
    ranker.fit(X, y, qid=q, sample_weight=w, eval_set=[(X, y)], eval_qid=[q])
    for k, v in ranker.evals_result()["validation_0"].items():
        note(v)
        assert v[-1] >= v[0]
        assert ranker.n_features_in_ == 3


@pytest.mark.skipif(**tm.no_pandas())
@pytest.mark.skipif(**tm.no_sklearn())
def test_unbiased() -> None:
    import pandas as pd
    from sklearn.model_selection import train_test_split

    X, y, q, w = tm.make_ltr(8192, 2, n_query_groups=6, max_rel=4)
    X, Xe, y, ye, q, qe = train_test_split(X, y, q, test_size=0.2, random_state=3)
    X = csr_matrix(X)
    Xe = csr_matrix(Xe)
    data = RelDataCV((X, y, q), (Xe, ye, qe), max_rel=4)

    train, _ = simulate_clicks(data)
    x, c, y, q = sort_ltr_samples(train.X, train.y, train.qid, train.click, train.pos)
    df: Optional[pd.DataFrame] = None

    class Position(xgboost.callback.TrainingCallback):
        def after_training(self, model) -> bool:
            nonlocal df
            config = json.loads(model.save_config())
            ti_plus = np.array(config["learner"]["objective"]["ti+"])
            tj_minus = np.array(config["learner"]["objective"]["tj-"])
            df = pd.DataFrame({"ti+": ti_plus, "tj-": tj_minus})
            return model

    ltr = xgboost.XGBRanker(
        n_estimators=8,
        tree_method="hist",
        lambdarank_unbiased=True,
        lambdarank_num_pair_per_sample=12,
        lambdarank_pair_method="topk",
        objective="rank:ndcg",
        callbacks=[Position()],
        base_score=0.5,
    )
    ltr.fit(x, c, qid=q, eval_set=[(x, c)], eval_qid=[q])

    assert df is not None
    # normalized
    np.testing.assert_allclose(df["ti+"].iloc[0], 1.0)
    np.testing.assert_allclose(df["tj-"].iloc[0], 1.0)
    assert np.isfinite(df["ti+"]).all()
    assert np.isfinite(df["tj-"]).all()
    assert (df["ti+"] >= 0.0).all()
    assert (df["tj-"] >= 0.0).all()

    # Training continuation
    ltr.fit(x, c, qid=q, eval_set=[(x, c)], eval_qid=[q], xgb_model=ltr)
    # normalized
    np.testing.assert_allclose(df["ti+"].iloc[0], 1.0)
    np.testing.assert_allclose(df["tj-"].iloc[0], 1.0)


def test_normalization() -> None:
    run_normalization("cpu")


@pytest.mark.parametrize("objective", ["rank:pairwise", "rank:ndcg", "rank:map"])
def test_score_normalization(objective: str) -> None:
    run_score_normalization("cpu", objective)


class TestRanking:
    @classmethod
    def setup_class(cls):
        """
        Download and setup the test fixtures
        """
        cls.dpath = "demo/"
        (
            x_train,
            y_train,
            qid_train,
            x_test,
            y_test,
            qid_test,
            x_valid,
            y_valid,
            qid_valid,
        ) = tm.data.get_mq2008(cls.dpath)

        # instantiate the matrices
        cls.dtrain = xgboost.DMatrix(x_train, y_train)
        cls.dvalid = xgboost.DMatrix(x_valid, y_valid)
        cls.dtest = xgboost.DMatrix(x_test, y_test)
        # set the group counts from the query IDs
        cls.dtrain.set_group(
            [len(list(items)) for _key, items in itertools.groupby(qid_train)]
        )
        cls.dtest.set_group(
            [len(list(items)) for _key, items in itertools.groupby(qid_test)]
        )
        cls.dvalid.set_group(
            [len(list(items)) for _key, items in itertools.groupby(qid_valid)]
        )
        # save the query IDs for testing
        cls.qid_train = qid_train
        cls.qid_test = qid_test
        cls.qid_valid = qid_valid

        # model training parameters
        cls.params = {
            "objective": "rank:pairwise",
            "booster": "gbtree",
            "eval_metric": ["ndcg"],
        }

    @classmethod
    def teardown_class(cls):
        """
        Cleanup test artifacts from download and unpacking
        :return:
        """
        zip_f = cls.dpath + "MQ2008.zip"
        if os.path.exists(zip_f):
            os.remove(zip_f)
        directory = cls.dpath + "MQ2008"
        if os.path.exists(directory):
            shutil.rmtree(directory)

    def test_training(self):
        """
        Train an XGBoost ranking model
        """
        # specify validations set to watch performance
        watchlist = [(self.dtest, "eval"), (self.dtrain, "train")]
        bst = xgboost.train(
            self.params,
            self.dtrain,
            num_boost_round=2500,
            early_stopping_rounds=10,
            evals=watchlist,
        )
        assert bst.best_score > 0.98

    def test_cv(self):
        """
        Test cross-validation with a group specified
        """
        cv = xgboost.cv(
            self.params,
            self.dtrain,
            num_boost_round=2500,
            early_stopping_rounds=10,
            nfold=10,
            as_pandas=False,
        )
        assert isinstance(cv, dict)
        assert set(cv.keys()) == {
            "test-ndcg-mean",
            "train-ndcg-mean",
            "test-ndcg-std",
            "train-ndcg-std",
        }, "CV results dict key mismatch."

    def test_cv_no_shuffle(self):
        """
        Test cross-validation with a group specified
        """
        cv = xgboost.cv(
            self.params,
            self.dtrain,
            num_boost_round=2500,
            early_stopping_rounds=10,
            shuffle=False,
            nfold=10,
            as_pandas=False,
        )
        assert isinstance(cv, dict)
        assert len(cv) == 4

    def test_get_group(self):
        """
        Retrieve the group number from the dmatrix
        """
        # test the new getter
        self.dtrain.get_uint_info("group_ptr")

        for d, qid in [
            (self.dtrain, self.qid_train),
            (self.dvalid, self.qid_valid),
            (self.dtest, self.qid_test),
        ]:
            # size of each group
            group_sizes = np.array(
                [len(list(items)) for _key, items in itertools.groupby(qid)]
            )
            # indexes of group boundaries
            group_limits = d.get_uint_info("group_ptr")
            assert len(group_limits) == len(group_sizes) + 1
            assert np.array_equal(np.diff(group_limits), group_sizes)
            assert np.array_equal(group_sizes, np.diff(d.get_uint_info("group_ptr")))
            assert np.array_equal(group_sizes, np.diff(d.get_uint_info("group_ptr")))
            assert np.array_equal(group_limits, d.get_uint_info("group_ptr"))
