# Copyright 1999-2021 Alibaba Group Holding Ltd.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import itertools
import logging
import operator
import pickle
from collections import defaultdict
from functools import reduce

import numpy as np

from .... import opcodes
from ....core import ExecutableTuple, OutputType, recursive_tile
from ....core.context import get_context
from ....core.operand import MergeDictOperand
from ....serialization.serializables import (
    FieldTypes,
    DictField,
    Int32Field,
    KeyField,
    ListField,
    StringField,
)
from ...utils import concat_chunks, collect_ports
from ._align import align_data_set
from .core import LGBMModelType, get_model_cls_from_type

logger = logging.getLogger(__name__)


class LGBMTrain(MergeDictOperand):
    _op_type_ = opcodes.LGBM_TRAIN

    _model_type = Int32Field(
        "model_type", on_serialize=lambda x: x.value, on_deserialize=LGBMModelType
    )
    _params = DictField("params", key_type=FieldTypes.string)
    _data = KeyField("data")
    _label = KeyField("label")
    _sample_weight = KeyField("sample_weight")
    _init_score = KeyField("init_score")
    _kwds = DictField("kwds", key_type=FieldTypes.string)

    _eval_datas = ListField("eval_datas", FieldTypes.key)
    _eval_labels = ListField("eval_labels", FieldTypes.key)
    _eval_sample_weights = ListField("eval_sample_weights", FieldTypes.key)
    _eval_init_scores = ListField("eval_init_scores", FieldTypes.key)

    _workers = ListField("workers", FieldTypes.string)
    _worker_id = Int32Field("worker_id")
    _worker_ports = KeyField("worker_ports")

    _tree_learner = StringField("tree_learner")
    _timeout = Int32Field("timeout")

    def __init__(
        self,
        model_type=None,
        data=None,
        label=None,
        sample_weight=None,
        init_score=None,
        eval_datas=None,
        eval_labels=None,
        eval_sample_weights=None,
        eval_init_scores=None,
        params=None,
        kwds=None,
        workers=None,
        worker_id=None,
        worker_ports=None,
        tree_learner=None,
        timeout=None,
        **kw,
    ):
        super().__init__(
            _model_type=model_type,
            _params=params,
            _data=data,
            _label=label,
            _sample_weight=sample_weight,
            _init_score=init_score,
            _eval_datas=eval_datas,
            _eval_labels=eval_labels,
            _eval_sample_weights=eval_sample_weights,
            _eval_init_scores=eval_init_scores,
            _kwds=kwds,
            _workers=workers,
            _worker_id=worker_id,
            _worker_ports=worker_ports,
            _tree_learner=tree_learner,
            _timeout=timeout,
            **kw,
        )
        if self.output_types is None:
            self.output_types = [OutputType.object]

    @property
    def model_type(self) -> LGBMModelType:
        return self._model_type

    @property
    def data(self):
        return self._data

    @property
    def label(self):
        return self._label

    @property
    def sample_weight(self):
        return self._sample_weight

    @property
    def init_score(self):
        return self._init_score

    @property
    def eval_datas(self) -> list:
        return self._eval_datas or []

    @property
    def eval_labels(self) -> list:
        return self._eval_labels or []

    @property
    def eval_sample_weights(self) -> list:
        return self._eval_sample_weights or []

    @property
    def eval_init_scores(self) -> list:
        return self._eval_init_scores or []

    @property
    def params(self) -> dict:
        return self._params or dict()

    @property
    def kwds(self) -> dict:
        return self._kwds or dict()

    @property
    def workers(self) -> list:
        return self._workers

    @property
    def worker_id(self) -> int:
        return self._worker_id

    @property
    def worker_ports(self):
        return self._worker_ports

    @property
    def timeout(self) -> int:
        return self._timeout

    @property
    def tree_learner(self) -> str:
        return self._tree_learner

    def _set_inputs(self, inputs):
        super()._set_inputs(inputs)
        it = iter(inputs)
        for attr in ["_data", "_label", "_sample_weight", "_init_score"]:
            if getattr(self, attr) is not None:
                setattr(self, attr, next(it))
        for attr in [
            "_eval_datas",
            "_eval_labels",
            "_eval_sample_weights",
            "_eval_init_scores",
        ]:
            new_list = []
            for c in getattr(self, attr, None) or []:
                if c is not None:
                    new_list.append(next(it))
            setattr(self, attr, new_list or None)

        if self._worker_ports is not None:
            self._worker_ports = next(it)

    def __call__(self):
        inputs = []
        for attr in ["_data", "_label", "_sample_weight", "_init_score"]:
            if getattr(self, attr) is not None:
                inputs.append(getattr(self, attr))
        for attr in [
            "_eval_datas",
            "_eval_labels",
            "_eval_sample_weights",
            "_eval_init_scores",
        ]:
            for c in getattr(self, attr, None) or []:
                if c is not None:
                    inputs.append(c)
        return self.new_tileable(inputs)

    @staticmethod
    def _get_data_chunks_workers(ctx, data):
        # data_chunk.inputs is concat, and concat's input is the co-allocated chunks
        metas = ctx.get_chunks_meta(
            [c.key for c in data.chunks], fields=["ip", "bands"]
        )

        ips = []
        ip_to_worker = {}
        for m in metas:
            ip = m["ip"]
            assert ip, "There is meta {meta} who doesn't contain ip."
            ips.append(ip)
            bands = m["bands"]
            ip_to_worker[ip] = bands[0][0] if bands else None
        return ips, ip_to_worker

    @staticmethod
    def _concat_chunks_by_worker(chunks, chunk_workers):
        worker_to_chunks = defaultdict(list)
        for chunk, worker in zip(chunks, chunk_workers):
            worker_to_chunks[worker].append(chunk)
        worker_to_concat = dict()
        for worker, chunks in worker_to_chunks.items():
            worker_to_concat[worker] = concat_chunks(chunks)
        return worker_to_concat

    @classmethod
    def tile(cls, op: "LGBMTrain"):
        ctx = get_context()
        data = op.data
        worker_to_args = defaultdict(dict)

        # Note: Mars worker is band address, and LGBMTrain worker is machine ip.
        ips, ip_to_worker = cls._get_data_chunks_workers(ctx, data)

        for arg in ["_data", "_label", "_sample_weight", "_init_score"]:
            if getattr(op, arg) is not None:
                for worker, chunk in cls._concat_chunks_by_worker(
                    getattr(op, arg).chunks, ips
                ).items():
                    worker_to_args[worker][arg] = chunk

        if op.eval_datas:
            eval_workers_list = [
                cls._get_data_chunks_workers(ctx, d)[0] for d in op.eval_datas
            ]
            extra_workers = reduce(
                operator.or_, (set(w) for w in eval_workers_list)
            ) - set(ips)
            worker_remap = dict(zip(extra_workers, itertools.cycle(ips)))
            if worker_remap:
                eval_workers_list = [
                    [worker_remap.get(w, w) for w in wl] for wl in eval_workers_list
                ]

            for arg in [
                "_eval_datas",
                "_eval_labels",
                "_eval_sample_weights",
                "_eval_init_scores",
            ]:
                if getattr(op, arg):
                    for tileable, eval_workers in zip(
                        getattr(op, arg), eval_workers_list
                    ):
                        for worker, chunk in cls._concat_chunks_by_worker(
                            tileable.chunks, eval_workers
                        ).items():
                            if arg not in worker_to_args[worker]:
                                worker_to_args[worker][arg] = []
                            worker_to_args[worker][arg].append(chunk)

        out_chunks = []
        ips = list(set(ips))
        workers = list(ip_to_worker.values())
        for worker_id, worker in enumerate(ips):
            chunk_op = op.copy().reset_key()
            chunk_op.expect_worker = ip_to_worker[worker]

            input_chunks = []
            concat_args = worker_to_args.get(worker, {})
            for arg in [
                "_data",
                "_label",
                "_sample_weight",
                "_init_score",
                "_eval_datas",
                "_eval_labels",
                "_eval_sample_weights",
                "_eval_init_scores",
            ]:
                arg_val = getattr(op, arg)
                if arg_val:
                    arg_chunk = concat_args.get(arg)
                    setattr(chunk_op, arg, arg_chunk)
                    if isinstance(arg_chunk, list):
                        input_chunks.extend(arg_chunk)
                    else:
                        input_chunks.append(arg_chunk)

            worker_ports_chunk = (
                yield from recursive_tile(collect_ports(workers, op.data))
            ).chunks[0]
            input_chunks.append(worker_ports_chunk)

            chunk_op._workers = ips
            chunk_op._worker_ports = worker_ports_chunk
            chunk_op._worker_id = worker_id

            data_chunk = concat_args["_data"]
            out_chunk = chunk_op.new_chunk(
                input_chunks, shape=(np.nan,), index=data_chunk.index[:1]
            )
            out_chunks.append(out_chunk)

        new_op = op.copy()
        return new_op.new_tileables(
            op.inputs, chunks=out_chunks, nsplits=((np.nan for _ in out_chunks),)
        )

    @classmethod
    def execute(cls, ctx, op: "LGBMTrain"):
        if op.merge:
            return super().execute(ctx, op)

        from lightgbm.basic import _safe_call, _LIB

        data_val = ctx[op.data.key]
        data_val = data_val.spmatrix if hasattr(data_val, "spmatrix") else data_val

        label_val = ctx[op.label.key]
        sample_weight_val = (
            ctx[op.sample_weight.key] if op.sample_weight is not None else None
        )
        init_score_val = ctx[op.init_score.key] if op.init_score is not None else None

        if op.eval_datas is None:
            eval_set, eval_sample_weight, eval_init_score = None, None, None
        else:
            eval_set, eval_sample_weight, eval_init_score = [], [], []
            for data, label in zip(op.eval_datas, op.eval_labels):
                data_eval = ctx[data.key]
                data_eval = (
                    data_eval.spmatrix if hasattr(data_eval, "spmatrix") else data_eval
                )
                eval_set.append((data_eval, ctx[label.key]))
            for weight in op.eval_sample_weights:
                eval_sample_weight.append(
                    ctx[weight.key] if weight is not None else None
                )
            for score in op.eval_init_scores:
                eval_init_score.append(ctx[score.key] if score is not None else None)

            eval_set = eval_set or None
            eval_sample_weight = eval_sample_weight or None
            eval_init_score = eval_init_score or None

        params = op.params.copy()
        # if model is trained, remove unsupported parameters
        params.pop("out_dtype_", None)
        worker_ports = ctx[op.worker_ports.key]
        worker_endpoints = [
            f"{worker}:{port}" for worker, port in zip(op.workers, worker_ports)
        ]

        params["machines"] = ",".join(worker_endpoints)
        params["time_out"] = op.timeout
        params["num_machines"] = len(worker_endpoints)
        params["local_listen_port"] = worker_ports[op.worker_id]

        if (op.tree_learner or "").lower() not in {"data", "feature", "voting"}:
            logger.warning(
                "Parameter tree_learner not set or set to incorrect value "
                f'{op.tree_learner}, using "data" as default'
            )
            params["tree_learner"] = "data"
        else:
            params["tree_learner"] = op.tree_learner

        try:
            model_cls = get_model_cls_from_type(op.model_type)
            model = model_cls(**params)
            model.fit(
                data_val,
                label_val,
                sample_weight=sample_weight_val,
                init_score=init_score_val,
                eval_set=eval_set,
                eval_sample_weight=eval_sample_weight,
                eval_init_score=eval_init_score,
                **op.kwds,
            )

            if (
                op.model_type == LGBMModelType.RANKER
                or op.model_type == LGBMModelType.REGRESSOR
            ):
                model.set_params(out_dtype_=np.dtype("float"))
            elif hasattr(label_val, "dtype"):
                model.set_params(out_dtype_=label_val.dtype)
            else:
                model.set_params(out_dtype_=label_val.dtypes[0])

            ctx[op.outputs[0].key] = pickle.dumps(model)
        finally:
            _safe_call(_LIB.LGBM_NetworkFree())


def train(params, train_set, eval_sets=None, **kwargs):
    eval_sets = eval_sets or []
    model_type = kwargs.pop("model_type", LGBMModelType.CLASSIFIER)

    evals_result = kwargs.pop("evals_result", dict())
    session = kwargs.pop("session", None)
    run_kwargs = kwargs.pop("run_kwargs", None)
    if run_kwargs is None:
        run_kwargs = dict()
    timeout = kwargs.pop("timeout", 120)
    base_port = kwargs.pop("base_port", None)

    aligns = align_data_set(train_set)
    for eval_set in eval_sets:
        aligns += align_data_set(eval_set)

    aligned_iter = iter(ExecutableTuple(aligns).execute(session))
    datas, labels, sample_weights, init_scores = [], [], [], []
    for dataset in [train_set] + eval_sets:
        train_kw = dict()
        for arg in ["data", "label", "sample_weight", "init_score"]:
            if getattr(dataset, arg) is not None:
                train_kw[arg] = next(aligned_iter)
            else:
                train_kw[arg] = None

        datas.append(train_kw["data"])
        labels.append(train_kw["label"])
        sample_weights.append(train_kw["sample_weight"])
        init_scores.append(train_kw["init_score"])

    train_cls = kwargs.pop("train_cls", LGBMTrain)
    op = train_cls(
        params=params,
        data=datas[0],
        label=labels[0],
        sample_weight=sample_weights[0],
        init_score=init_scores[0],
        eval_datas=datas[1:],
        eval_labels=labels[1:],
        eval_weights=sample_weights[1:],
        eval_init_score=init_scores[1:],
        model_type=model_type,
        timeout=timeout,
        lgbm_port=base_port,
        kwds=kwargs,
    )
    ret = op().execute(session=session, **run_kwargs).fetch(session=session)
    # Note: There may be multiple results if running on clusters, so we should
    # combine the results if so.
    if isinstance(ret, list):
        ret = b"".join(ret)

    bst = pickle.loads(ret)
    evals_result.update(bst.evals_result_ or {})
    return bst
