# 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.

from collections import namedtuple
from functools import partial

import numpy as np
import pandas as pd

from .... import opcodes
from ....serialization.serializables import BoolField
from ..aggregation import BaseDataFrameExpandingAgg

_stage_info = namedtuple(
    "_stage_info",
    (
        "map_groups",
        "map_sources",
        "combine_sources",
        "combine_columns",
        "combine_funcs",
        "key_to_funcs",
        "valid_columns",
        "min_periods_func_name",
    ),
)

_cum_alpha_coeff_func = "_cum_alpha_coeff"
_cum_square_alpha_coeff_func = "_cum_square_alpha_coeff"


def _add_pred_results(pred_results, local_results, axis=0):
    if pred_results[0].ndim == 1:
        df_filler = 0
    else:
        df_filler = pred_results[0].iloc[-1, :].dropna()
        df_filler[:] = 0

    new_locals = []
    combine_axis = pred_results[0].ndim - axis - 1
    for pred_result, local_result in zip(pred_results, local_results):
        local_result = local_result.fillna(df_filler, axis=axis)
        new_locals.append(
            local_result.add(pred_result.sum(axis=axis), axis=combine_axis)
        )
    return new_locals


def _combine_arithmetic(pred_results, local_results, axis=0):
    if pred_results is None:
        return local_results[0]
    return _add_pred_results(pred_results, local_results, axis=axis)[0]


def _combine_minmax(pred_results, local_results, axis=0, fun_name=None):
    if pred_results is None:
        return local_results[0]

    pred_size = len(pred_results[0])
    con = pd.concat([pred_results[0], local_results[0]], axis=axis)
    result = con.expanding(axis=axis).agg(fun_name)
    if result.ndim == 2:
        return result.iloc[pred_size:, :] if axis == 0 else result.iloc[:, pred_size:]
    else:
        return result.iloc[pred_size:]


def _combine_mean(pred_results, local_results, axis=0):
    local_sum_data, local_count_data = local_results

    if pred_results is not None:
        local_sum_data, local_count_data = _add_pred_results(
            pred_results, local_results, axis=axis
        )
    return local_sum_data / local_count_data


def _combine_var(pred_results, local_results, axis=0):
    local_sum_data, local_count_data, local_var_data = local_results
    if pred_results is None:
        return local_var_data * local_count_data / (local_count_data - 1)

    pred_sum_data, pred_count_data, pred_var_data = pred_results

    local_sum_square = (
        local_count_data * local_var_data + local_sum_data**2 / local_count_data
    )
    pred_sum_square = (
        pred_count_data * pred_var_data + pred_sum_data**2 / pred_count_data
    )

    local_sum_square, local_sum_data, local_count_data = _add_pred_results(
        [pred_sum_square, pred_sum_data, pred_count_data],
        [local_sum_square, local_sum_data, local_count_data],
        axis=axis,
    )

    return (local_sum_square - local_sum_data**2 / local_count_data) / (
        local_count_data - 1
    )


def _combine_std(pred_results, local_results, axis=0):
    return np.sqrt(_combine_var(pred_results, local_results, axis=axis))


class DataFrameExpandingAgg(BaseDataFrameExpandingAgg):
    _op_type_ = opcodes.EXPANDING_AGG

    _center = BoolField("center")

    def __init__(self, center=None, **kw):
        super().__init__(_center=center, **kw)

    @property
    def center(self):
        return self._center

    @classmethod
    def _get_stage_functions(cls, op: "DataFrameExpandingAgg", func):
        if func == "_data_count":
            return ["count"], _combine_arithmetic
        elif func in ("sum", "prod", "count"):
            return [func], _combine_arithmetic
        elif func in ("min", "max"):
            return [func], partial(_combine_minmax, fun_name=func)
        elif func == "mean":
            return ["sum", "count"], _combine_mean
        elif func in {"var", "std"}:
            return (
                ["sum", "count", "var"],
                _combine_var if func == "var" else _combine_std,
            )
        else:  # pragma: no cover
            raise NotImplementedError

    @classmethod
    def _execute_map_function(cls, op: "DataFrameExpandingAgg", func, in_data):
        min_periods = 1 if op.min_periods > 0 else 0

        expanding = in_data.expanding(
            min_periods=min_periods, center=op.center, axis=op.axis
        )
        if func == "var":
            result = expanding.var(ddof=0)
        else:
            result = expanding.agg(func)

        if op.output_agg:
            summary = result.iloc[len(result) - 1 : len(result)]
        else:
            summary = None
        return result, summary

    @classmethod
    def _execute_combine_function(
        cls, op: "DataFrameExpandingAgg", func, pred_inputs, local_inputs, func_cols
    ):
        return func(pred_inputs, local_inputs, axis=op.axis)

    @classmethod
    def _execute_raw_function(cls, op: "DataFrameExpandingAgg", in_data):
        expanding = in_data.expanding(
            min_periods=op.min_periods, center=op.center, axis=op.axis
        )
        return expanding.agg(op.func)
