from copy import deepcopy

import apache_beam as beam


class _PairWithIndexNumberDoFn(beam.DoFn):
    def process(self, row):
        _, value_iter = row
        for i, value in enumerate(sorted(value_iter)):
            yield (value, i)


class PairWithIndexNumber(beam.PTransform):
    def expand(self, pcoll):
        return (
            pcoll
            | beam.Distinct()
            | beam.Map(lambda r: (None, r))
            | beam.GroupByKey()
            | beam.ParDo(_PairWithIndexNumberDoFn())
        )


@beam.ptransform_fn
def DigestCategoricalColumns(pcoll, categorical_columns):
    def to_dict(src_pcoll, column):
        dict_pcoll = (
            src_pcoll
            | f"Map[{column}]" >> beam.Map(lambda r: r[column])
            | f"Pair[{column}]" >> PairWithIndexNumber()
        )
        return beam.pvalue.AsDict(dict_pcoll)

    dicts = {cat_col: to_dict(pcoll, cat_col) for cat_col in categorical_columns}

    def set_cat_cols(row, cat_cols, **dicts):
        result = {**row}
        for cat_col in cat_cols:
            result[cat_col] = dicts[cat_col][result[cat_col]]
        return result

    return pcoll | beam.Map(set_cat_cols, categorical_columns, **dicts)


@beam.ptransform_fn
def _CreateCategoricalDict(pcoll, existing_dict_pairs):
    """
    For a specific column, creates a new "categorical dict"
    mapping values to unique ints.
    """
    existing_max_value = (
        existing_dict_pairs
        | "just values" >> beam.Map(lambda r: r[1])
        | "get max" >> beam.combiners.Top.Of(1)
        | "extract" >> beam.FlatMap(lambda r: r)
    )

    new_pairs = (
        pcoll
        | "filter for unseen"
        >> beam.Filter(
            lambda row, existing: row not in existing,
            existing=beam.pvalue.AsDict(existing_dict_pairs),
        )
        | beam.Distinct()
        | "group into single list" >> beam.combiners.ToList()
        | "append unique values"
        >> beam.FlatMap(
            lambda row, max_v: [(key, max_v + 1 + i) for i, key in enumerate(row)],
            max_v=beam.pvalue.AsSingleton(existing_max_value, default_value=0),
        )
    )

    return [new_pairs, existing_dict_pairs] | "combine new/existing" >> beam.Flatten()


@beam.ptransform_fn
def CreateCategoricalDicts(pcoll, cat_cols, existing_dict_rows):
    """
    For a set of columnular data inputs this function takes:

      - cat_cols:

          Type: [str]

          An array of "categorical" columns

      - existing_dicts:

          Type: PCollection[(string, string, int)]

          Rows of tuples of type:
            (column, previously_seen_value, mapped_unique_int)

          Mapping a set of "previously seen" keys to unique int values for each
          column.
          Not optional.
          If none exist, pass an empty PCollection.

    It then creates a transform which takes a pcollection and
      - looks at the input pcoll for unseen values in each categorical column
      - creates new unique integers for each distinct unseen value, starting at
        max(previous value for column)+1
      - ammends the existing mappings with (col, unseen_value, new_unique_int)

    Output is:

        - Type: PCollection[(string, string, int)]

    This is useful for preparing data to be trained by LightGBM
    """

    acc = []

    existing_dicts = (
        existing_dict_rows
        | beam.Map(lambda r: beam.pvalue.TaggedOutput(r[0], (r[1], r[2]))).with_outputs()
    )

    for column in cat_cols:
        existing_dict_pairs = existing_dicts[column]

        acc.append(
            pcoll
            | f"extract column value [{column}]"
            # Note: column=column used to capture value at loop time.
            #       prevents column being final value of loop due to
            #       value by reference.
            >> beam.Map(lambda r, column=column: r[column])
            | _CreateCategoricalDict(existing_dict_pairs)
            | f"re-append column [{column}]" >> beam.Map(lambda r, column=column: (column, *r))
        )

    return acc | beam.Flatten()


@beam.ptransform_fn
def ReplaceCategoricalColumns(
    inputs,
    cat_cols,
    categorical_dict_rows,
    default_unseen=None,
    features_key=None,
):
    """
    Utilizes the "categorical dictionary rows" generated by
    CreateCategoricalDicts which is a list of pairs of type of (column, value,
    unique_int), and replaces each column with the appropriate value found in
    the mapping.

    In there case where no mapping exists, `default_unseen` will be returned.
    This defaults to None, but can be set as a kwarg.
    """

    categorical_mapping_dicts = (
        categorical_dict_rows
        | "index by column" >> beam.Map(lambda r: (r[0], (r[1], r[2])))
        | "group by column" >> beam.GroupByKey()
        | "create dict per column" >> beam.MapTuple(lambda k, v: (k, dict(v)))
    )

    if features_key:
        # Nesting -- process separately from row
        def _process_nested(row, d):
            digested = {}

            features_items = row[features_key]

            for col, val in features_items.items():
                if col in cat_cols:
                    digested[col] = d.get(col, {}).get(val, default_unseen)
                else:
                    digested[col] = val

            ret = deepcopy(row)
            ret[features_key] = digested

            return ret

        return inputs | beam.Map(
            _process_nested,
            d=beam.pvalue.AsDict(categorical_mapping_dicts),
        )

    else:
        # No nesting - use row in place
        def _process(row, d):
            digested = {}

            for col, val in row.items():
                if col in cat_cols:
                    digested[col] = d.get(col, {}).get(val, default_unseen)
                else:
                    digested[col] = val

            return digested

        return inputs | beam.Map(
            _process,
            d=beam.pvalue.AsDict(categorical_mapping_dicts),
        )
