#!/usr/bin/python2
#
# Copyright 2018 Google LLC
#
# Use of this source code is governed by a BSD-style
# license that can be found in the LICENSE file or at
# https://developers.google.com/open-source/licenses/bsd

"""Functions for data source to DSPL model conversion."""
from __future__ import print_function


__author__ = 'Benjamin Yolken <yolken@google.com>'

import itertools

import data_source
from dspllib.model import dspl_model


def _CalculateSlices(column_bundle):
  """Calculate all the possible slices to be produced from a column bundle.

  Args:
    column_bundle: A DataSourceColumnBundle object produced by a data source

  Returns:
    A sequence of DataSourceColumn sequences; each sequence contains the columns
    for one slice.
  """
  all_slices = []

  binary_elements = []
  non_binary_elements = []

  child_parent_dict = {}

  for column in column_bundle.GetColumnIterator():
    if column.rollup:
      binary_elements.append(column)
    else:
      non_binary_elements.append(column)

    if column.parent_ref:
      child_parent_dict[column] = column_bundle.GetColumnByID(column.parent_ref)

  # Expand out slices using powerset operator
  for selection in _Powerset(binary_elements):
    transformed_slice = non_binary_elements + list(selection)

    all_slices.append([s for s in transformed_slice])

  # Prune slices that contain both a concept and an ancestor
  slices_to_evaluate = []

  for data_slice in all_slices:
    keep_slice = True

    for key in child_parent_dict.keys():
      if key in data_slice:
        curr_val = child_parent_dict[key]

        # Loop through concept ancestors
        while True:
          if curr_val in data_slice:
            keep_slice = False
            break
          elif curr_val in child_parent_dict:
            curr_val = child_parent_dict[curr_val]
          else:
            break

        if not keep_slice:
          break

    if keep_slice:
      slices_to_evaluate.append(data_slice)

  return slices_to_evaluate


def _Powerset(input_list):
  """Create powerset iterator from the elements in a list.

  Note: Based on example in official Python itertools documentation.

  Example:
    [p for p in Powerset(['a', 'b')] == [(), ('a',), ('b'), ('a', 'b')]

  Args:
    input_list: A sequence of Python objects

  Returns:
    An iterator which loops through all (tuple) elements in the powerset of
    the input_list.
  """
  return (
      itertools.chain.from_iterable(
          itertools.combinations(input_list, r) for r
          in range(len(input_list) + 1)))


def _CreateConceptTable(
        column, instance_data, parent_column=None, verbose=True):
  """Create a DSPL table object that enumerates the instances of a concept.

  If the concept extends 'entity:entity' or 'geo:location', extra columns
  are added to the resulting table for the required inherited properties.
  Otherwise, the table has just a single column.

  By convention, the table id is given as [column_id]_table and the
  corresponding CSV is named [column_id]_table.csv.

  Args:
    column: A DataSourceColumn object corresponding to a dataset dimension
    instance_data: A TableRows object containing a list of concept instances
    parent_column: A DataSourceColumn object corresponding to the parent
                   of this concept
    verbose: Print out status messages to stdout

  Returns:
    A DSPL Table object
  """
  dspl_columns = [dspl_model.TableColumn(column_id=column.column_id,
                                         data_type=column.data_type)]

  if parent_column:
    dspl_columns += [dspl_model.TableColumn(column_id=parent_column.column_id,
                                            data_type=parent_column.data_type)]

  if column.concept_extension == 'entity:entity':
    # Add a 'name' column and populate it with the instance IDs
    dspl_columns += [
        dspl_model.TableColumn(column_id='name', data_type='string')]

    dspl_table_data = instance_data.MergeValues(instance_data).rows
  elif column.concept_extension == 'geo:location':
    # Add 'name', 'latitude', and 'longitude' columns; populate the first with
    # the instance IDs, the others with blank values
    dspl_columns += [
        dspl_model.TableColumn(column_id='name', data_type='string'),
        dspl_model.TableColumn(column_id='latitude', data_type='float'),
        dspl_model.TableColumn(column_id='longitude', data_type='float')]

    dspl_table_data = (instance_data.MergeValues(instance_data)
                       .MergeConstant('').MergeConstant('').rows)
  else:
    dspl_table_data = instance_data.rows

  # Create table, including header row in table data
  concept_table = dspl_model.Table(
      table_id='%s_table' % (column.column_id),
      columns=dspl_columns,
      file_name='%s_table.csv' % (column.column_id),
      table_data=[[c.column_id for c in dspl_columns]] + dspl_table_data,
      verbose=verbose)

  return concept_table


def _CreateSliceTable(slice_columns, table_id, file_name,
                      slice_data, verbose=True):
  """Create a DSPL Table object for a dataset slice.

  Args:
    slice_columns: Sequence of DataSourceColumn objects representing concepts
                   in this slice
    table_id: ID for the table
    file_name: Name of the CSV file containing the table data
    slice_data: A TableRows object containing the data for the slice table
    verbose: Print out status messages to stdout

  Returns:
    A DSPL Table object
  """
  dspl_columns = [
      dspl_model.TableColumn(c.column_id, c.data_type, c.data_format)
      for c in slice_columns]

  # Create table, including header row in table data
  slice_table = dspl_model.Table(
      table_id=table_id,
      columns=dspl_columns,
      file_name=file_name,
      table_data=[[c.column_id for c in dspl_columns]] + slice_data.rows,
      verbose=verbose)

  return slice_table


def PopulateDataset(data_source_obj, verbose):
  """Create a DSPL dataset from a data source.

  Loops through the set of possible slices (provided by the _CalculateSlices
  function), creating the necessary DSPL concept, slice, and table objects as
  needed.

  The following naming convention is used:

    DSPL concept ID                  := DataSource column ID
    DSPL table ID for concept tables := DSPL concept ID + "_table"
    DSPL slice ID                    := "slice_" + n, where n=0,1,2,...
    DSPL table ID for slice tables   := DSPL slice ID + "_table"
    DSPL table file name             := DSPL table ID + ".csv"

  Args:
    data_source_obj: An object that implements the DataSource interface