#
# Copyright 2014 Google Inc. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Utility class used to collate data and fill gaps in multiple reports."""

import datetime
import itertools

DATE_FORMAT = u'%Y-%m-%d'
MONTH_FORMAT = u'%Y-%m'


class DataCollator(object):
  """Used to collate reports and fill the chronological gaps.

  Combines multiple reports in one and makes sure that every dimension value
  is present, filling with zeros when appropiate. Also, fills the chronological
  gaps (for DATE, MONTH and WEEK).

  Example:

  REPORT 1
  ==============================
  MONTH     COUNTRY_CODE  CLICKS
  FEBRUARY  US            100
  FEBRUARY  UK            50
  APRIL     ES            10

  REPORT 2
  ==============================
  MONTH     COUNTRY_CODE  CLICKS
  APRIL     US            80

  RESULT
  ==============================
  MONTH     COUNTRY_CODE  CLICKS
  FEBRUARY  US            100
  FEBRUARY  UK            50
  FEBRUARY  ES            0
  MARCH     US            0
  MARCH     UK            0
  MARCH     ES            0
  APRIL     US            80
  APRIL     UK            0
  APRIL     ES            10

  Attributes:
    reports: A list of reports generated by the Google Client Library
  """
  # Stores all possible values for a given dimension.
  dimension_values = {}
  date_index = -1
  week_index = -1
  month_index = -1
  number_of_columns = 0
  start_date = None
  end_date = None
  ratios_and_currencies_indices = []

  def __init__(self, reports):
    self.reports = reports
    self._check_input_headers()
    self.number_of_columns = len(reports[0]['headers'])
    self._collect_all_dimension_values()
    self._check_input_dates()
    self._set_global_start_end_dates()
    self._get_all_ratios_and_currencies()

  def collate_data(self):
    """Collects and combines data from multiple reports, removing gaps.

    Returns:
      A single report containing collated data from every initial report and
      filled with missing data.

    Raises:
      ValueError: The input reports have problems.
    """

    # The result will have the same structure/order as the first report.
    result = {}
    # Set start and end dates (already calculated).
    result['startDate'] = date_to_date_st(self.start_date)
    result['endDate'] = date_to_date_st(self.end_date)
    result['rows'] = []
    # Headers are the same for every report, copy headers from the first one.
    result['headers'] = self.reports[0]['headers']

    cursor_date = self.start_date
    while cursor_date <= self.end_date:
      if self.dimension_values:
        for combination in self._generate_every_dimension_combination():
          # Combination is something like ('dim1_value1', 'dim2_value3').
          row = self._get_dimension_combination_row(combination, cursor_date)
          if row is None:
            row = self._create_new_row(combination, cursor_date)
          result['rows'].append(row)

      # Increment the date cursor.
      # Increment one day.
      if self.date_index != -1:
        cursor_date += datetime.timedelta(days=1)
      # Increment one week.
      elif self.week_index != -1:
        prev_month = cursor_date.month
        cursor_date += datetime.timedelta(days=7)
        # Weeks can run between months, so the same week with different month
        # must be added as well.
        if cursor_date.month != prev_month:
          cursor_date = datetime.datetime(
              cursor_date.year, cursor_date.month, 1)
      # Increment one month.
      elif self.month_index != -1:
        cursor_date = increase_month(cursor_date)
        cursor_date = datetime.datetime(cursor_date.year, cursor_date.month, 1)
      else:
        # No time dimensions.
        break

    return result

  def _check_input_headers(self):
    reports = self.reports
    if not reports:
      raise ValueError('No reports given.')
    for report_index, report in enumerate(reports):
      if not report['headers']:
        raise ValueError('Report %d has no headers.' % (report_index))
      if len(report['headers']) != len(reports[0]['headers']):
        raise ValueError('All reports must have the same number of columns.')
      has_dimensions = False
      for h, header in enumerate(report['headers']):
        if header['type'] == 'DIMENSION':
          has_dimensions = True
        if header['name'] != reports[0]['headers'][h]['name']:
          raise ValueError(
              'All reports must have the same columns in the same order.')
      if not has_dimensions:
        raise ValueError('Report %d has no dimensions' % (report_index))

  def _check_input_dates(self):
    reports = self.reports
    for report_index, report in enumerate(reports):
      if report_index == 0:
        continue
      prev_end_date = date_st_to_date(reports[report_index - 1]['endDate'])
      start_date = date_st_to_date(report['startDate'])
      if start_date > prev_end_date + datetime.timedelta(days=1):
        raise ValueError('There are date gaps between reports. '
                         'Only gaps inside reports are filled.')

  def _set_global_start_end_dates(self):
    for report in self.reports:
      report_start_date = date_st_to_date(report['startDate'])
      report_end_date = date_st_to_date(report['endDate'])
      if self.start_date is None or report_start_date < self.start_date:
        self.start_date = report_start_date
      if self.end_date is None or report_end_date > self.end_date:
        self.end_date = report_end_date

  def _generate_every_dimension_combination(self):
    """Generates every combination of dimension values.

    Example: for two dimensions, with two possible values each:
    [(UK, MALE),
     (UK, FEMALE),
     (US, MALE),
     (US, FEMALE)]

    Returns:
      A generator for every combination of dimension values.
    """
    # Edge case: dimensions were defined but the report is empty. The dimension
    # values need to be flattened:
    number_of_values = len(
        list(itertools.chain(*self.dimension_values.values())))

    if self.dimension_values.keys() and not number_of_values:
      comb = ([] for _ in xrange(len(self.dimension_values)))
      return comb

    # Normal case:
    return itertools.product(*self.dimension_values.values())

  def _create_new_row(self, combination, date):
    """Creates a new row in the resulting report.

    Args:
      combination: A tuple of dimension values, generated by
        self._generate_every_dimension_combination()
      date: a datetime object.

    Returns:
      The newly created row.
    """

    # Fill everything with zeros
    row = [0 for _ in range(self.number_of_columns)]

    # Fill dates.
    if self.date_index != -1:
      row[self.date_index] = date_to_date_st(date)
    if self.week_index != -1:
      row[self.week_index] = date_to_week_st(date)
    if self.month_index != -1:
      row[self.month_index] = date_to_month_st(date)

    # Fill dimensions.
    if combination:
      for combination_index, dimension_index in enumerate(
          self.dimension_values):
        row[dimension_index] = combination[combination_index]

    # Instead of a zero, add None to ratios and currencies.
    for ratio_or_currency_index in self.ratios_and_currencies_indices:
      row[ratio_or_currency_index] = None

    return row

  def _get_header_index(self, report, header):
    try:
      return [x['name'] for x in report['headers']].index(header)
    except ValueError:
      return -1

  def _collect_all_dimension_values(self):
    """Gets all the existing values for every dimension."""

    self.dimension_values = {}

    for i, report in enumerate(self.reports):

      # Get the date, week and month indices
      date_index = self._get_header_index(report, 'DATE')
      week_index = self._get_header_index(report, 'WEEK')
      month_index = self._get_header_index(report, 'MONTH')

      # The result will have the same indices as the first report.
      if i == 0:
        self.date_index = date_index
        self.week_index = week_index
        self.month_index = month_index

      for dimension_index, unused_dim in self._get_all_dimensions_from_report(
          report):
        if dimension_index in (date_index, week_index, month_index):
          continue
        if dimension_index not in self.dimension_values:
          self.dimension_values[dimension_index] = []
        # Add all the values for this dimension.
        try:
          for row in report['rows']:
            if row[dimension_index] not in self.dimension_values[dimension_index]:
              self.dimension_values[dimension_index].append(row[dimension_index])
        except KeyError:
          self.dimension_values = None

  def _get_all_dimensions_from_report(self, report):
    # Returns [(index0, dimension0), (index1, dimension1)...].
    return [(i, x['name']) for i, x in enumerate(report['headers'])
            if x['type'] == 'DIMENSION']

  def _get_all_ratios_and_currencies(self):
    self.ratios_and_currencies_indices = [
        i for i, x in enumerate(self.reports[0]['headers'])
        if x['type'] == 'METRIC_RATIO' or x['type'] == 'METRIC_CURRENCY']

  def _get_dimension_combination_row(self, combination, date):
    """Checks if a particular combination of dimensions exists.

    Args:
      combination: a tuple with a combination of values.
      date: a datetime object representing the row's date.

    Returns:
      The index of the row that matches the combinations.
    """
    for report in self.reports:
      for row in report['rows']:
        # Filter date.
        if self.date_index != -1:
          if row[self.date_index] != date_to_date_st(date):
            continue
        elif self.week_index != -1:
          if row[self.week_index] != date_to_week_st(date):
            continue
        elif self.month_index != -1:
          if row[self.month_index] != date_to_month_st(date):
            continue

        # Filter dimension combination.
        row_found = True
        # Check if all the dimensions of the row fit the combination:
        # For every dimension:
        for dimension_values_index, combination_index in enumerate(
            self.dimension_values):
          # Check if the value is the same
          row_value = row[combination_index]
          if row_value != combination[dimension_values_index]:
            # next row
            row_found = False
            break
        # Nobody protested, it's our row.
        if row_found:
          return row
    return None


def date_to_date_st(date):
  """Returns a string representing the date with DATE_FORMAT."""
  return date.strftime(DATE_FORMAT)


def date_to_week_st(date):
  """Returns a string representing the week with DATE_FORMAT."""
  return (
      date - datetime.timedelta(days=date.weekday())).strftime(DATE_FORMAT)


def date_to_month_st(date):
  """Returns a string representing the month with MONTH_FORMAT."""
  return date.strftime(MONTH_FORMAT)


def date_st_to_date(date_st):
  """Returns a datetime object from a string in DATE_FORMAT."""
  return datetime.datetime.strptime(date_st, DATE_FORMAT)


def increase_month(date):
  """Increase a date by one month.

  Jan 31 will return Feb 28 (or 29)

  Args:
    date: A datetime object

  Returns:
    A datetime object one month from the initial date.
  """
  try:
    if date.month == 12:
      # If the month is Dec, add a year and set the month to Jan.
      return datetime.datetime(date.year + 1, 1, date.day)
    return datetime.datetime(date.year, date.month + 1, date.day)
  except ValueError:
    # Trying to create invalid date like Feb 31st.
    # Take the first day of the month after the next and subtract one day.
    # Note that date.month will never be 11 nor 12 as there's no way to trigger
    # a ValueError (both January and December have 31 days).
    return (datetime.datetime(date.year, date.month + 2, 1)
            - datetime.timedelta(days=1))

