import itertools
import sys
from collections import OrderedDict, defaultdict

import numpy as np


def get_user_item_sparse_indices(
    data, user_unique_vals, item_unique_vals, mode, ordered
):
    user_indices = column_sparse_indices(
        data.user.to_numpy(), user_unique_vals, mode, ordered
    )
    item_indices = column_sparse_indices(
        data.item.to_numpy(), item_unique_vals, mode, ordered
    )
    return user_indices, item_indices


def merge_sparse_col(sparse_col, multi_sparse_col):
    flatten_cols = list(itertools.chain.from_iterable(multi_sparse_col))
    return flatten_cols if not sparse_col else sparse_col + flatten_cols


def merge_sparse_indices(data_class, data, sparse_col, multi_sparse_col, mode, ordered):
    if sparse_col and multi_sparse_col:
        sparse_indices = get_sparse_indices_matrix(
            data_class, data, sparse_col, mode, ordered
        )
        sparse_offset = get_sparse_offset(data_class, sparse_col)
        sparse_indices = sparse_indices + sparse_offset[:-1]

        multi_sparse_indices = get_multi_sparse_indices_matrix(
            data_class, data, multi_sparse_col, mode, ordered
        )
        multi_sparse_offset = get_multi_sparse_offset(data_class, multi_sparse_col)
        multi_sparse_indices = (
            multi_sparse_indices + sparse_offset[-1] + multi_sparse_offset
        )
        return np.concatenate([sparse_indices, multi_sparse_indices], axis=1)

    elif sparse_col:
        sparse_indices = get_sparse_indices_matrix(
            data_class, data, sparse_col, mode, ordered
        )
        sparse_offset = get_sparse_offset(data_class, sparse_col)
        return sparse_indices + sparse_offset[:-1]

    elif multi_sparse_col:
        multi_sparse_indices = get_multi_sparse_indices_matrix(
            data_class, data, multi_sparse_col, mode, ordered
        )
        multi_sparse_offset = get_multi_sparse_offset(data_class, multi_sparse_col)
        return multi_sparse_indices + multi_sparse_offset


def get_sparse_indices_matrix(data_class, data, sparse_col, mode, ordered):
    n_samples, n_features = len(data), len(sparse_col)
    sparse_indices = np.zeros((n_samples, n_features), dtype=np.int32)
    for i, col in enumerate(sparse_col):
        col_values = data[col].to_numpy()
        unique_values = data_class.sparse_unique_vals[col]
        sparse_indices[:, i] = column_sparse_indices(
            col_values, unique_values, mode, ordered
        )
    return sparse_indices


def get_multi_sparse_indices_matrix(data_class, data, multi_sparse_col, mode, ordered):
    n_samples = len(data)
    # n_fields = len(multi_sparse_col)
    n_features = len(list(itertools.chain.from_iterable(multi_sparse_col)))
    multi_sparse_indices = np.zeros((n_samples, n_features), dtype=np.int32)
    i = 0
    while i < n_features:
        for field in multi_sparse_col:
            unique_values = data_class.multi_sparse_unique_vals[field[0]]
            for col in field:
                col_values = data[col].to_numpy()
                multi_sparse_indices[:, i] = column_sparse_indices(
                    col_values, unique_values, mode, ordered, multi_sparse=True
                )
                i += 1
    return multi_sparse_indices


# def get_dense_indices_matrix(data, dense_col):
#    n_samples, n_features = len(data), len(dense_col)
#    dense_indices = np.tile(np.arange(n_features), [n_samples, 1])
#    return dense_indices


def get_sparse_offset(data_class, sparse_col):
    # plus one for oov value
    unique_values = [len(data_class.sparse_unique_vals[col]) + 1 for col in sparse_col]
    return np.cumsum(np.array([0] + unique_values))


def get_multi_sparse_offset(data_class, multi_sparse_col):
    unique_values = [
        len(data_class.multi_sparse_unique_vals[field[0]]) + 1
        for field in multi_sparse_col
    ]
    field_offset = np.cumsum(np.array([0] + unique_values)).tolist()[:-1]
    offset = []
    for i, field in enumerate(multi_sparse_col):
        offset.extend([field_offset[i]] * len(field))
    return np.array(offset)


def merge_offset(data_class, sparse_col, multi_sparse_col):
    if sparse_col and multi_sparse_col:
        sparse_offset = get_sparse_offset(data_class, sparse_col)
        multi_sparse_offset = (
            get_multi_sparse_offset(data_class, multi_sparse_col) + sparse_offset[-1]
        )
        return np.concatenate([sparse_offset[:-1], multi_sparse_offset])
    elif sparse_col:
        sparse_offset = get_sparse_offset(data_class, sparse_col)
        return sparse_offset[:-1]
    elif multi_sparse_col:
        multi_sparse_offset = get_multi_sparse_offset(data_class, multi_sparse_col)
        return multi_sparse_offset


def sparse_oov(data_class, sparse_col):
    unique_values = [len(data_class.sparse_unique_vals[col]) + 1 for col in sparse_col]
    return np.cumsum(unique_values) - 1


def multi_sparse_oov(data_class, multi_sparse_col, extend=True):
    unique_values = [
        len(data_class.multi_sparse_unique_vals[field[0]]) + 1
        for field in multi_sparse_col
    ]
    # return np.cumsum(unique_values) - 1
    field_oov = np.cumsum(unique_values) - 1
    if extend:
        oov = []
        for i, field in enumerate(multi_sparse_col):
            oov.extend([field_oov[i]] * len(field))
        return np.array(oov)
    else:
        return field_oov


def get_oov_pos(data_class, sparse_col, multi_sparse_col):
    if sparse_col and multi_sparse_col:
        sparse = sparse_oov(data_class, sparse_col)
        multi_sparse = multi_sparse_oov(data_class, multi_sparse_col) + sparse[-1] + 1
        return np.concatenate([sparse, multi_sparse])
    elif sparse_col:
        return sparse_oov(data_class, sparse_col)
    elif multi_sparse_col:
        return multi_sparse_oov(data_class, multi_sparse_col)


def multi_sparse_combine_info(
    data_class, all_sparse_cols, sparse_col, multi_sparse_col
):
    from ..data import MultiSparseInfo

    # sparse_last_offset = get_sparse_offset(data_class, sparse_col)[-1]
    # unique_values = [
    #    len(data_class.multi_sparse_unique_vals[field[0]]) + 1
    #    for field
    #    in multi_sparse_col
    # ]
    # offset = np.cumsum(np.array([0] + unique_values))
    # offset = offset.tolist()[:-1] + sparse_last_offset
    # all_sparse_cols = data_class.sparse_col.name
    field_offset = [all_sparse_cols.index(field[0]) for field in multi_sparse_col]
    field_length = [len(col) for col in multi_sparse_col]
    sparse_offset = sparse_oov(data_class, sparse_col)[-1] + 1 if sparse_col else 0
    feat_oov = multi_sparse_oov(data_class, multi_sparse_col, extend=False)
    return MultiSparseInfo(field_offset, field_length, feat_oov + sparse_offset)


# def multi_sparse_true_size(data, multi_sparse_col, pad_val):
#    sizes = [
#        np.sum(data[field].to_numpy() != pad_val, axis=1) for field in multi_sparse_col
#    ]
#    return np.vstack(sizes)


def check_unknown(values, uniques):
    # diff = list(np.setdiff1d(values, uniques, assume_unique=True))
    mask = np.in1d(values, uniques, invert=True)
    return mask


def column_sparse_indices(
    values, unique, mode="train", ordered=True, multi_sparse=False
):
    if mode not in ("train", "test"):
        raise ValueError('mode must either be "train" or "test" ')
    if ordered:
        if mode == "test" or multi_sparse:
            not_in_mask = check_unknown(values, unique)
            col_indices = np.searchsorted(unique, values)
            col_indices[not_in_mask] = len(unique)
        else:
            col_indices = np.searchsorted(unique, values)
    else:
        map_vals = dict(zip(unique, range(len(unique))))
        oov_val = len(unique)
        if mode == "test" or multi_sparse:
            col_indices = np.array(
                [map_vals[v] if v in map_vals else oov_val for v in values]
            )
        else:
            col_indices = np.array([map_vals[v] for v in values])
    return col_indices


def interaction_consumed(user_indices, item_indices):
    user_consumed = defaultdict(list)
    item_consumed = defaultdict(list)
    for u, i in zip(user_indices, item_indices):
        if isinstance(u, np.integer):
            u = u.item()
        if isinstance(i, np.integer):
            i = i.item()
        user_consumed[u].append(i)
        item_consumed[i].append(u)
    return remove_duplicates(user_consumed, item_consumed)


def remove_duplicates(user_consumed, item_consumed):
    # keys will preserve order in dict since Python3.7
    if sys.version_info[:2] >= (3, 7):
        dict_func = dict.fromkeys
    else:
        dict_func = OrderedDict.fromkeys
    user_dedup = {u: list(dict_func(items)) for u, items in user_consumed.items()}
    item_dedup = {i: list(dict_func(users)) for i, users in item_consumed.items()}
    return user_dedup, item_dedup


def multi_sparse_col_map(multi_sparse_col):
    multi_sparse_map = dict()
    for field in multi_sparse_col:
        if len(field) > 1:
            for col in field[1:]:
                multi_sparse_map[col] = field[0]
    return multi_sparse_map


def recover_sparse_cols(data_info):
    total_sparse_cols = data_info.sparse_col.name
    if data_info.sparse_unique_vals:
        sparse_cols = [
            col for col in total_sparse_cols if col in data_info.sparse_unique_vals
        ]
    else:
        sparse_cols = None

    if data_info.multi_sparse_unique_vals:
        multi_sparse_cols = []
        i, field = 0, 0
        while i < len(total_sparse_cols):
            col = total_sparse_cols[i]
            if col in data_info.multi_sparse_unique_vals:
                field_len = data_info.multi_sparse_combine_info.field_len[field]
                multi_sparse_cols.append(
                    [total_sparse_cols[k] for k in range(i, i + field_len)]
                )
                i += field_len
                field += 1
            else:
                i += 1
    else:
        multi_sparse_cols = None

    return sparse_cols, multi_sparse_cols
