import math

import lightgbm
import pandas as pd
import pytest
from sklearn.datasets import load_breast_cancer
from sklearn.model_selection import train_test_split

import ray
import ray.data
from ray.train import ScalingConfig
from ray.train.constants import TRAIN_DATASET_KEY
from ray.train.lightgbm import (
    LightGBMTrainer,
    RayTrainReportCallback,
    normalize_pandas_for_lightgbm,
)
from ray.train.v2._internal.constants import is_v2_enabled

assert is_v2_enabled()


@pytest.fixture
def ray_start_6_cpus():
    address_info = ray.init(num_cpus=6)
    yield address_info
    # The code after the yield will run as teardown code.
    ray.shutdown()


scale_config = ScalingConfig(num_workers=2)

data_raw = load_breast_cancer()
dataset_df = pd.DataFrame(data_raw["data"], columns=data_raw["feature_names"])
dataset_df["target"] = data_raw["target"]
train_df, test_df = train_test_split(dataset_df, test_size=0.3)

params = {
    "objective": "binary",
    "metric": ["binary_logloss", "binary_error"],
}


def test_fit_with_categoricals(ray_start_6_cpus):
    @ray.remote
    class ValidationCollector:
        def __init__(self):
            self.validation_scores = {}

        def report(self, rank, binary_logloss, binary_error):
            self.validation_scores[rank] = {
                "binary_logloss": binary_logloss,
                "binary_error": binary_error,
            }

        def get_validation_scores(self):
            return self.validation_scores

    # Ensure all workers have the same model in data parallel training
    # by comparing their validation scores.
    # Comparing lightgbm models directly seems less reliable.
    collector = ValidationCollector.remote()

    def lightgbm_train_fn_per_worker(
        config: dict,
        label_column: str,
        valid_dataset: ray.data.Dataset,
        num_boost_round: int = 10,
    ):
        remaining_iters = num_boost_round
        train_ds_iter = ray.train.get_dataset_shard(TRAIN_DATASET_KEY)
        train_df = normalize_pandas_for_lightgbm(
            train_ds_iter.materialize().to_pandas()
        )

        eval_df = normalize_pandas_for_lightgbm(valid_dataset.materialize().to_pandas())
        eval_X, eval_y = eval_df.drop(label_column, axis=1), eval_df[label_column]
        valid_set = lightgbm.Dataset(eval_X, label=eval_y)

        train_X, train_y = train_df.drop(label_column, axis=1), train_df[label_column]
        train_set = lightgbm.Dataset(train_X, label=train_y)

        # Add network params of the worker group to enable distributed training.
        config.update(ray.train.lightgbm.get_network_params())

        # Add lightgbm-specific distributed training params.
        config.update(
            {
                "tree_learner": "data_parallel",
                "pre_partition": True,
            }
        )

        booster = lightgbm.train(
            params=config,
            train_set=train_set,
            num_boost_round=remaining_iters,
            # NOTE: Include the training dataset in the evaluation datasets.
            # This allows `train-*` metrics to be calculated and reported.
            valid_sets=[valid_set, train_set],
            valid_names=["valid", TRAIN_DATASET_KEY],
            init_model=None,
            callbacks=[RayTrainReportCallback()],
        )
        collector.report.remote(
            ray.train.get_context().get_world_rank(),
            booster.best_score["valid"]["binary_logloss"],
            booster.best_score["valid"]["binary_error"],
        )

    train_df_with_cat = train_df.copy()
    test_df_with_cat = test_df.copy()
    train_df_with_cat["categorical_column"] = pd.Series(
        (["A", "B"] * math.ceil(len(train_df_with_cat) / 2))[: len(train_df_with_cat)]
    ).astype("category")
    test_df_with_cat["categorical_column"] = pd.Series(
        (["A", "B"] * math.ceil(len(test_df_with_cat) / 2))[: len(test_df_with_cat)]
    ).astype("category")

    train_dataset = ray.data.from_pandas(train_df_with_cat)
    valid_dataset = ray.data.from_pandas(test_df_with_cat)
    trainer = LightGBMTrainer(
        train_loop_per_worker=lambda: lightgbm_train_fn_per_worker(
            config=params,
            label_column="target",
            # Do not shard the validation dataset across workers to ensure all workers compute
            # the same validation score. See https://github.com/microsoft/LightGBM/issues/4392.
            valid_dataset=valid_dataset,
        ),
        scaling_config=scale_config,
        datasets={TRAIN_DATASET_KEY: train_dataset},
    )
    result = trainer.fit()
    checkpoint = result.checkpoint
    model = RayTrainReportCallback.get_model(checkpoint)
    assert model.pandas_categorical == [["A", "B"]]
    validation_scores = ray.get(collector.get_validation_scores.remote())
    assert validation_scores[0]["binary_logloss"] == pytest.approx(
        validation_scores[1]["binary_logloss"], abs=1e-6
    )
    assert validation_scores[0]["binary_error"] == pytest.approx(
        validation_scores[1]["binary_error"], abs=1e-6
    )


if __name__ == "__main__":
    import sys

    import pytest

    sys.exit(pytest.main(["-v", "-x", __file__]))