# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you 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.

import unittest
from datasketches import kll_ints_sketch, kll_floats_sketch, kll_doubles_sketch
from datasketches import kll_items_sketch, ks_test, PyStringsSerDe
import copy
import numpy as np

class KllTest(unittest.TestCase):
    def test_kll_floats_example(self):
      k = 160
      n = 2 ** 20

      # create a sketch and inject ~1 million N(0,1) points as an array and as a single item
      kll = kll_floats_sketch(k)
      kll.update(np.random.normal(size=n-1))
      kll.update(0.0)

      # 0 should be near the median
      self.assertAlmostEqual(0.5, kll.get_rank(0.0), delta=0.035)
      
      # the median should be near 0
      self.assertAlmostEqual(0.0, kll.get_quantile(0.5), delta=0.035)

      # we also track the min/max independently from the rest of the data
      # which lets us know the full observed data range
      self.assertLessEqual(kll.get_min_value(), kll.get_quantile(0.01))
      self.assertLessEqual(0.0, kll.get_rank(kll.get_min_value()))
      self.assertGreaterEqual(kll.get_max_value(), kll.get_quantile(0.99))
      self.assertGreaterEqual(1.0, kll.get_rank(kll.get_max_value()))

      # we can also extract a list of values at a time,
      # here the values should give us something close to [-2, -1, 0, 1, 2].
      # then get the CDF, which will return something close to
      # the original values used in get_quantiles()
      # finally, can check the normalized rank error bound
      pts = kll.get_quantiles([0.0228, 0.1587, 0.5, 0.8413, 0.9772])
      cdf = kll.get_cdf(pts)  # include 1.0 at end to account for all probability mass
      self.assertEqual(len(cdf), len(pts)+1)
      err = kll.normalized_rank_error(False)
      self.assertEqual(err, kll_floats_sketch.get_normalized_rank_error(k, False))

      # and a few basic queries about the sketch
      self.assertFalse(kll.is_empty())
      self.assertTrue(kll.is_estimation_mode())
      self.assertEqual(kll.n, n)
      self.assertEqual(kll.k, k)
      self.assertLess(kll.num_retained, n)

      # merging itself will double the number of items the sketch has seen
      # but need to do that with a copy
      kll_copy = copy.copy(kll)
      kll.merge(kll_copy)
      self.assertEqual(kll.n, 2*n)

      # we can then serialize and reconstruct the sketch
      kll_bytes = kll.serialize()
      new_kll = kll_floats_sketch.deserialize(kll_bytes)
      self.assertEqual(kll.num_retained, new_kll.num_retained)
      self.assertEqual(kll.get_min_value(), new_kll.get_min_value())
      self.assertEqual(kll.get_max_value(), new_kll.get_max_value())
      self.assertEqual(kll.get_quantile(0.7), new_kll.get_quantile(0.7))
      self.assertEqual(kll.get_rank(0.0), new_kll.get_rank(0.0))

      # A Kolmogorov-Smirnov Test of kll and new_kll should match, even for
      # a fairly small p-value -- cannot reject the null hypothesis that
      # they come from the same distribution (since they do)
      self.assertFalse(ks_test(kll, new_kll, 0.001))

      total_weight = 0
      for tuple in kll:
        item = tuple[0]
        weight = tuple[1]
        total_weight = total_weight + weight
      self.assertEqual(total_weight, kll.n)

    def test_kll_ints_sketch(self):
        k = 100
        n = 10
        kll = kll_ints_sketch(k)
        for i in range(0, n):
          kll.update(i)

        self.assertEqual(kll.get_min_value(), 0)
        self.assertEqual(kll.get_max_value(), n-1)
        self.assertEqual(kll.n, n)
        self.assertFalse(kll.is_empty())
        self.assertFalse(kll.is_estimation_mode()) # n < k
        self.assertEqual(kll.k, k)

        pmf = kll.get_pmf([round(n/2)])
        self.assertIsNotNone(pmf)
        self.assertEqual(len(pmf), 2)

        cdf = kll.get_cdf([round(n/2)])
        self.assertIsNotNone(cdf)
        self.assertEqual(len(cdf), 2)

        self.assertEqual(kll.get_quantile(0.5), round(n/2))
        quants = kll.get_quantiles([0.25, 0.5, 0.75])
        self.assertIsNotNone(quants)
        self.assertEqual(len(quants), 3)

        self.assertEqual(kll.get_rank(round(n/2)), 0.5)

        # merge copy of self
        kll_copy = copy.copy(kll)
        kll.merge(kll_copy)
        self.assertEqual(kll.n, 2 * n)

        sk_bytes = kll.serialize()
        self.assertTrue(isinstance(kll_ints_sketch.deserialize(sk_bytes), kll_ints_sketch))

    def test_kll_doubles_sketch(self):
      # already tested float and ints and it's templatized, so just make sure it instantiates properly
      k = 75
      kll = kll_doubles_sketch(k)
      self.assertTrue(kll.is_empty())

    def test_kll_items_sketch(self):
      # most functionality has been tested, but we need to ensure objects and sorting work
      # as well as serialization
      k = 100
      n = 2 ** 16

      # create a sketch and inject enough points to force compaction
      kll = kll_items_sketch(k)
      for i in range(0, n):
        kll.update(str(i))
      
      # check iterator for items
      total_weight = 0
      for tuple in kll:
        item = tuple[0]
        weight = tuple[1]
        total_weight = total_weight + weight
      self.assertEqual(total_weight, kll.n)

      kll_copy = copy.copy(kll)
      kll.merge(kll_copy)
      self.assertEqual(kll.n, 2 * n)
      
      kll_bytes = kll.serialize(PyStringsSerDe())
      new_kll = kll_items_sketch.deserialize(kll_bytes, PyStringsSerDe())
      self.assertEqual(kll.num_retained, new_kll.num_retained)
      self.assertEqual(kll.get_min_value(), new_kll.get_min_value())
      self.assertEqual(kll.get_max_value(), new_kll.get_max_value())
      self.assertEqual(kll.get_quantile(0.7), new_kll.get_quantile(0.7))
      self.assertEqual(kll.get_rank(str(n/4)), new_kll.get_rank(str(n/4)))

      # The sketches are identical so we cannot reject the null hypothesis that they
      # reflect the same underlying distribtion
      self.assertFalse(ks_test(kll, new_kll, 0.001))

      # check that printing works as expected
      self.assertGreater(len(kll.to_string(True, True)), 0)
      self.assertEqual(len(kll.__str__()), len(kll.to_string()))

if __name__ == '__main__':
    unittest.main()
