# Copyright 2017 The TensorFlow Authors. 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.
# ==============================================================================
"""Package for histogram compression."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import collections

import numpy as np

# Normal CDF for std_devs: (-Inf, -1.5, -1, -0.5, 0, 0.5, 1, 1.5, Inf)
# naturally gives bands around median of width 1 std dev, 2 std dev, 3 std dev,
# and then the long tail.
NORMAL_HISTOGRAM_BPS = (0, 668, 1587, 3085, 5000, 6915, 8413, 9332, 10000)


CompressedHistogramValue = collections.namedtuple('CompressedHistogramValue',
                                                  ['basis_point', 'value'])


# TODO(@jart): Unfork these methods.
def compress_histogram_proto(histo, bps=NORMAL_HISTOGRAM_BPS):
  """Creates fixed size histogram by adding compression to accumulated state.

  This routine transforms a histogram at a particular step by interpolating its
  variable number of buckets to represent their cumulative weight at a constant
  number of compression points. This significantly reduces the size of the
  histogram and makes it suitable for a two-dimensional area plot where the
  output of this routine constitutes the ranges for a single x coordinate.

  Args:
    histo: A HistogramProto object.
    bps: Compression points represented in basis points, 1/100ths of a percent.
        Defaults to normal distribution.

  Returns:
    List of values for each basis point.
  """
  # See also: Histogram::Percentile() in core/lib/histogram/histogram.cc
  if not histo.num:
    return [CompressedHistogramValue(b, 0.0) for b in bps]
  bucket = np.array(histo.bucket)
  bucket_limit = list(histo.bucket_limit)
  weights = (bucket * bps[-1] / (bucket.sum() or 1.0)).cumsum()
  values = []
  j = 0
  while j < len(bps):
    i = np.searchsorted(weights, bps[j], side='right')
    while i < len(weights):
      cumsum = weights[i]
      cumsum_prev = weights[i - 1] if i > 0 else 0.0
      if cumsum == cumsum_prev:  # prevent lerp divide by zero
        i += 1
        continue
      if not i or not cumsum_prev:
        lhs = histo.min
      else:
        lhs = max(bucket_limit[i - 1], histo.min)
      rhs = min(bucket_limit[i], histo.max)
      weight = _lerp(bps[j], cumsum_prev, cumsum, lhs, rhs)
      values.append(CompressedHistogramValue(bps[j], weight))
      j += 1
      break
    else:
      break
  while j < len(bps):
    values.append(CompressedHistogramValue(bps[j], histo.max))
    j += 1
  return values


def compress_histogram(buckets, bps=NORMAL_HISTOGRAM_BPS):
  """Creates fixed size histogram by adding compression to accumulated state.

  This routine transforms a histogram at a particular step by linearly
  interpolating its variable number of buckets to represent their cumulative
  weight at a constant number of compression points. This significantly reduces
  the size of the histogram and makes it suitable for a two-dimensional area
  plot where the output of this routine constitutes the ranges for a single x
  coordinate.

  Args:
    buckets: A list of buckets, each of which is a 3-tuple of the form
      `(min, max, count)`.
    bps: Compression points represented in basis points, 1/100ths of a percent.
        Defaults to normal distribution.

  Returns:
    List of values for each basis point.
  """
  # See also: Histogram::Percentile() in core/lib/histogram/histogram.cc
  buckets = np.array(buckets)
  if not buckets.size:
    return [CompressedHistogramValue(b, 0.0) for b in bps]
  (minmin, maxmax) = (buckets[0][0], buckets[-1][1])
  counts = buckets[:, 2]
  right_edges = list(buckets[:, 1])
  weights = (counts * bps[-1] / (counts.sum() or 1.0)).cumsum()

  result = []
  bp_index = 0
  while bp_index < len(bps):
    i = np.searchsorted(weights, bps[bp_index], side='right')
    while i < len(weights):
      cumsum = weights[i]
      cumsum_prev = weights[i - 1] if i > 0 else 0.0
      if cumsum == cumsum_prev:  # prevent division-by-zero in `_lerp`
        i += 1
        continue
      if not i or not cumsum_prev:
        lhs = minmin
      else:
        lhs = max(right_edges[i - 1], minmin)
      rhs = min(right_edges[i], maxmax)
      weight = _lerp(bps[bp_index], cumsum_prev, cumsum, lhs, rhs)
      result.append(CompressedHistogramValue(bps[bp_index], weight))
      bp_index += 1
      break
    else:
      break
  while bp_index < len(bps):
    result.append(CompressedHistogramValue(bps[bp_index], maxmax))
    bp_index += 1
  return result


def _lerp(x, x0, x1, y0, y1):
  """Affinely map from [x0, x1] onto [y0, y1]."""
  return y0 + (x - x0) * float(y1 - y0) / (x1 - x0)