import importlib.util
import itertools
import json
from fractions import Fraction
from math import expm1, log1p
from pathlib import Path
import unittest

ROOT = Path(__file__).resolve().parent
SPEC = importlib.util.spec_from_file_location("tail_model", ROOT / "model.py")
model = importlib.util.module_from_spec(SPEC)
SPEC.loader.exec_module(model)


class TailSamplingTests(unittest.TestCase):
    def test_independent_enumeration(self):
        for n in range(1, 9):
            for p in (Fraction(0), Fraction(1, 4), Fraction(1, 2), Fraction(1)):
                mass = Fraction(0)
                count_mass = Fraction(0)
                for state in itertools.product((0, 1), repeat=n):
                    k = sum(state)
                    weight = p ** k * (1 - p) ** (n - k)
                    if k:
                        mass += weight
                        count_mass += k * weight
                result = model.union_model(p, n)
                self.assertAlmostEqual(result["independent_any_slow"], float(mass))
                self.assertEqual(result["independent_conditional_count"],
                                 float(count_mass / mass) if mass else None)
                self.assertLessEqual(result["equal_marginal_lower_bound"], float(mass))
                self.assertGreaterEqual(result["equal_marginal_upper_bound"], float(mass))

    def test_stable_float_oracle(self):
        for p, n in ((Fraction(1, 100), 100), (Fraction(1, 10000), 2000)):
            actual = model.union_model(p, n)["independent_any_slow"]
            self.assertAlmostEqual(actual, -expm1(n * log1p(-float(p))), places=14)
        self.assertGreater(model.union_model(Fraction(1, 100), 100)["independent_conditional_count"], 1)

    def test_invalid_parameters(self):
        for p, n in ((Fraction(-1), 2), (Fraction(2), 2), (0.5, 2), (Fraction(1, 2), 0),
                     (Fraction(1, 2), True), (Fraction(1, 2), 1.5)):
            with self.assertRaises(ValueError):
                model.union_model(p, n)
        for values in ([], [-1], [True], [float("nan")]):
            with self.assertRaises(ValueError):
                model.sample_summary(values)

    def test_same_mean_different_tails(self):
        a, b = model.results()["same_mean"]
        self.assertEqual(a["mean_ms"], b["mean_ms"])
        self.assertEqual((a["p99_ms"], a["p999_ms"]), (10, 10))
        self.assertEqual((b["p99_ms"], b["p999_ms"]), (1, 901))

    def test_printed_counts(self):
        hdd, ssd = model.results()["observed_counts"]
        self.assertEqual(round(hdd["slow_drive_hour_percent"], 2), 0.22)
        self.assertEqual(round(hdd["tail_group_hour_percent"], 2), 1.54)
        self.assertEqual(round(ssd["slow_drive_hour_percent"], 2), 0.57)
        # The paper prints 0.58%, while its integer counts round to 0.57%.
        self.assertEqual(round(ssd["tail_group_hour_percent"], 2), 2.23)
        self.assertEqual(round(hdd["slow_drives_per_tail_group_hour"], 2), 1.70)
        self.assertEqual(round(ssd["slow_drives_per_tail_group_hour"], 2), 1.80)

    def test_generated_bytes(self):
        self.assertEqual((ROOT / "results.json").read_text(),
                         json.dumps(model.results(), indent=2) + "\n")


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