#!/usr/bin/env python3 -u
# -*- coding: utf-8 -*-
# copyright: sktime developers, BSD-3-Clause License (see LICENSE file)
"""Hierarchical Data Generators."""

__author__ = ["ltsaprounis"]

from itertools import product
from typing import Tuple, Union

import numpy as np
import pandas as pd
from sklearn.utils import check_random_state

from sktime.datasets import load_airline
from sktime.utils._testing.series import _make_index


def _make_hierarchical(
    hierarchy_levels: Tuple = (2, 4),
    max_timepoints: int = 12,
    min_timepoints: int = 12,
    same_cutoff: bool = True,
    n_columns: int = 1,
    all_positive: bool = True,
    index_type: str = None,
    random_state: Union[int, np.random.RandomState] = None,
    add_nan: bool = False,
) -> pd.DataFrame:
    """Generate hierarchical multiindex mtype for testing.

    Parameters
    ----------
    hierarchy_levels : Tuple, optional
        the number of groups at each hierarchy level, by default (2, 4)
    max_timepoints : int, optional
        maximum time points a series can have, by default 12
    min_timepoints : int, optional
        minimum time points a seires can have, by default 12
    same_cutoff : bool, optional
        If it's True all series will end at the same date, by default True
    n_columns : int, optional
        number of columns in the output dataframe, by default 1
    all_positive : bool, optional
        If True the time series will be , by default True
    index_type : str, optional
        type of index, by default None
        Supported types are "period", "datetime", "range" or "int".
        If it's not provided, "datetime" is selected.
    random_state : int, np.random.RandomState or None
        Controls the randomness of the estimator, by default None
    add_nan : bool, optional
        If it's true the series will contain NaNs, by default False

    Returns
    -------
    pd.DataFrame
        hierarchical mtype dataframe
    """
    levels = [
        [f"h{i}_{j}" for j in range(hierarchy_levels[i])]
        for i in range(len(hierarchy_levels))
    ]
    level_names = [f"h{i}" for i in range(len(hierarchy_levels))]
    rng = check_random_state(random_state)
    if min_timepoints == max_timepoints:
        time_index = _make_index(max_timepoints, index_type)
        index = pd.MultiIndex.from_product(
            levels + [time_index], names=level_names + ["time"]
        )
    else:
        df_list = []
        for levels_tuple in product(*levels):
            n_timepoints = rng.randint(low=min_timepoints, high=max_timepoints)
            if same_cutoff:
                time_index = _make_index(max_timepoints, index_type)[-n_timepoints:]
            else:
                time_index = _make_index(n_timepoints, index_type)
            d = dict(zip(level_names, levels_tuple))
            d["time"] = time_index
            df_list.append(pd.DataFrame(d))
        index = pd.MultiIndex.from_frame(
            pd.concat(df_list), names=level_names + ["time"]
        )

    total_time_points = len(index)
    data = rng.normal(size=(total_time_points, n_columns))
    if add_nan:
        # add some nan values
        data[int(len(data) / 2)] = np.nan
        data[0] = np.nan
        data[-1] = np.nan
    if all_positive:
        data -= np.min(data, axis=0) - 1
    df = pd.DataFrame(
        data=data, index=index, columns=[f"c{i}" for i in range(n_columns)]
    )

    return df


# todo adds some docs
def _bottom_hier_datagen(
    no_levels=3,
    no_bottom_nodes=6,
    intercept_max=20,
    coef_1_max=20,
    coef_2_max=0.1,
):
    """Hierarchical data generator using the flights dataset."""
    if no_levels > no_bottom_nodes:
        raise ValueError("no_levels should be less than no_bottom_nodes")

    base_ts = load_airline()
    df = pd.DataFrame(base_ts, index=base_ts.index)
    df.index.rename(None, inplace=True)

    if no_levels == 0:
        df.columns = ["passengers"]
        df.index.rename("timepoints", inplace=True)
        return df
    else:

        df.columns = ["l1_node01"]

        intercept = np.arange(0, intercept_max, 0.01)
        coef_1 = np.arange(0, coef_1_max, 0.01)
        coef_2 = np.arange(0, coef_2_max, 0.01)
        power_2 = ([0.5, 1, 1.5, 2],)

        node_lookup = pd.DataFrame(
            ["l1_node" + f"{x:02d}" for x in range(1, no_bottom_nodes + 1)]
        )
        node_lookup.columns = ["l1_agg"]

        if no_levels >= 2:

            for i in range(2, no_levels + 1):
                node_lookup["l" + str(i) + "_agg"] = node_lookup.groupby(
                    ["l" + str(i - 1) + "_agg"]
                )["l1_agg"].transform(
                    lambda x: "l"
                    + str(i)
                    + "_node"
                    + f"{int(np.sort(np.random.choice(np.arange(1,np.floor(len(node_lookup.index)/i)+1, 1), size=1))):02d}"  # noqa from flake8 E501
                )

        node_lookup = node_lookup.set_index("l1_agg", drop=True)

        for i in range(2, no_bottom_nodes + 1):
            df["l1_node" + f"{i:02d}"] = (
                np.random.choice(intercept, size=1)
                + np.random.choice(coef_1, size=1) * df["l1_node01"]
                + (
                    np.random.choice(coef_2, size=1)
                    * (df["l1_node01"] ** np.random.choice(power_2, size=1))
                )
            )

        df = (
            df.melt(ignore_index=False)
            .reset_index(drop=False)
            .rename(
                columns={
                    "variable": "l1_agg",
                    "index": "timepoints",
                    "value": "passengers",
                }
            )
        )

        df = pd.merge(left=df, right=node_lookup.reset_index(), on="l1_agg")
        df = df[df.columns.sort_values(ascending=True)]

        df_newindex = ["l" + str(x) + "_agg" for x in range(1, no_levels + 1)][::-1]
        df_newindex.append("timepoints")

        df = df.set_index(df_newindex)
        df.sort_index(inplace=True)

        return df
