"""Market share flows related util functions."""

from collections import defaultdict


def cast_data(data_set, type_mapping):
    """Cast data to the type described by type mapping.

    Args:
        data_set (iterable): Source data set.
        type_mapping (dict): Fields to types mapping.
    Returns:
        list: Data cast to proper types.
    """
    rows = []
    for row in data_set:
        new_row = {}

        for k, v in row.items():
            new_row[k] = type_mapping[k](v)

        rows.append(new_row)

    return rows


def groupby(data_set, key, aggregation_dict):
    """Group and aggregate data by column.

    Args:
        data_set (iterable): Source data set.
        key (str): Grouping column name.
        aggregation_dict (dict): Aggregation dict. Contains column names and
            aggregation operations. E.g.:
                {
                    'col1': sum,
                    'col2': sum
                }
    Returns:
        list: Grouped data set.
    """
    rows = []
    res = defaultdict(dict)

    for row in data_set:
        for col in aggregation_dict.keys():
            if col not in res[row[key]]:
                res[row[key]][col] = [row[col]]
            else:
                res[row[key]][col].append(row[col])

    for k, v in res.items():
        new_row = {key: k}

        for field, agg_lst in v.items():
            new_row[field] = aggregation_dict[field](agg_lst)

        rows.append(new_row)

    return rows


def join(key, ds, other):
    """Join two data sets.

    Args:
        key (str): Join key.
        ds (iterable): First data set.
        other (iterable): Second data set.
    Returns:
        list: Joined data set.
    """
    def find(search_clause, ds):
        """Find row in data set.

        Args:
            search_clause (tuple): Tuple containing column name and value.
            ds (iterable): Data Set.
        Returns:
            dict: First occurrence or None.
        """
        for row in ds:
            if row[search_clause[0]] == search_clause[1]:
                return row
        return None

    rows = []

    for row in ds:
        row_to_join = find((key, row[key]), other)
        new_row = row.copy()
        new_row.update(row_to_join)
        rows.append(new_row)

    return rows