#!/usr/bin/env python3
# (c) Facebook, Inc. and its affiliates. Confidential and proprietary.

from typing import Type

import numpy as np
from flsim.utils.async_trainer.training_event_generator import (
    AsyncTrainingEventGenerator,
    AsyncTrainingEventGeneratorConfig,
    AsyncTrainingEventGeneratorFromList,
    AsyncTrainingEventGeneratorFromListConfig,
    AsyncTrainingStartTimeDistrConfig,
    ConstantAsyncTrainingStartTimeDistr,
    ConstantAsyncTrainingStartTimeDistrConfig,
    EventTimingInfo,
    IEventGenerator,
    PoissonAsyncTrainingStartTimeDistr,
    PoissonAsyncTrainingStartTimeDistrConfig,
)
from flsim.utils.timing.training_duration_distribution import (
    PerExampleGaussianDurationDistribution,
    PerExampleGaussianDurationDistributionConfig,
    PerUserGaussianDurationDistribution,
    PerUserGaussianDurationDistributionConfig,
)
from libfb.py import testutil
from omegaconf import OmegaConf


class EventDistributionsUtilTest(testutil.BaseFacebookTestCase):
    def setUp(self) -> None:
        super().setUp()

    def test_simulated_training_training_event_generator(self) -> None:
        """Check that EventDistributionFromList works correctly by inputing
        a sample distribution, and confirming that the output is correct
        """
        timing_info1 = EventTimingInfo(prev_event_start_to_current_start=1, duration=2)
        timing_info2 = EventTimingInfo(prev_event_start_to_current_start=2, duration=1)
        timing_info3 = EventTimingInfo(prev_event_start_to_current_start=2, duration=5)
        random_list = [timing_info1, timing_info2, timing_info3]
        distr = AsyncTrainingEventGeneratorFromList(
            **OmegaConf.structured(
                AsyncTrainingEventGeneratorFromListConfig(training_events=random_list)
            )
        )
        self.assertTrue(
            distr.time_to_next_event_start()
            == timing_info1.prev_event_start_to_current_start
        )
        self.assertTrue(
            distr.training_duration(num_training_examples=1) == timing_info1.duration
        )
        self.assertTrue(
            distr.time_to_next_event_start()
            == timing_info2.prev_event_start_to_current_start
        )
        self.assertTrue(
            distr.training_duration(num_training_examples=1) == timing_info2.duration
        )
        self.assertTrue(
            distr.time_to_next_event_start()
            == timing_info3.prev_event_start_to_current_start
        )
        self.assertTrue(
            distr.training_duration(num_training_examples=1) == timing_info3.duration
        )

    def _duration_normality_check(
        self,
        event_generator: IEventGenerator,
        sample_count: int,
        expected_mean: float,
        expected_sd: float,
        epsilon: float,
    ) -> None:
        durations = []
        for _ in range(sample_count):
            durations.append(event_generator.training_duration(num_training_examples=1))
        # normality check doesn't verify mean and variance
        self.assertAlmostEqual(np.mean(durations), expected_mean, delta=epsilon)
        self.assertAlmostEqual(np.std(durations), expected_sd, delta=epsilon)

    def test_poisson_training_event_generator(self) -> None:
        """Check that TrainingEventDistritubion makes sense by checking that
        generated event durations follow normal distribution.
        """
        np.random.seed(1)
        # TODO: (jesikmin) T54593732 Add test to ensure that time_to_next_event_start
        # follows Poisson
        event_rate_per_sec, duration_mean, duration_sd = 10, 1, 5
        # set training_duration_min to a very negative number, to not bound the distribution
        duration_distr = PerExampleGaussianDurationDistributionConfig(
            training_duration_mean=duration_mean, training_duration_sd=duration_sd
        )
        training_start_time_distr = PoissonAsyncTrainingStartTimeDistrConfig(
            training_rate=event_rate_per_sec
        )
        distr = AsyncTrainingEventGenerator(
            **OmegaConf.structured(
                AsyncTrainingEventGeneratorConfig(
                    training_start_time_distribution=training_start_time_distr,
                    duration_distribution_generator=duration_distr,
                )
            )
        )
        self._duration_normality_check(
            distr,
            1000,  # sample_count
            duration_mean,
            duration_sd,
            epsilon=1,
        )

    def test_constant_training_event_distribution(self) -> None:
        """Check that ConstantAsyncTrainingStartTimeDistr generates the right
        next_event_time
        """
        np.random.seed(1)
        min_mean = 0.0001
        max_mean = 10
        max_sd = 1
        event_rate_per_sec = np.random.uniform(min_mean, max_mean)
        duration_mean = np.random.uniform(0, max_mean)
        duration_sd = np.random.uniform(0, max_sd)
        training_start_time_distr = ConstantAsyncTrainingStartTimeDistrConfig(
            training_rate=event_rate_per_sec
        )
        duration_distr = PerExampleGaussianDurationDistributionConfig(
            training_duration_mean=duration_mean, training_duration_sd=duration_sd
        )
        distr = AsyncTrainingEventGenerator(
            **OmegaConf.structured(
                AsyncTrainingEventGeneratorConfig(
                    training_start_time_distribution=training_start_time_distr,
                    duration_distribution_generator=duration_distr,
                )
            )
        )
        self.assertEqual(distr.time_to_next_event_start(), 1 / event_rate_per_sec)
        self._duration_normality_check(
            distr,
            1000,  # sample_count
            duration_mean,
            duration_sd,
            epsilon=1,
        )

    def test_constant_training_event_distribution_zero_sd(self) -> None:
        """Check that ConstantAsyncTrainingStartTimeDistr generates the right
        constant training_duration and next_event_time when SD is zero
        """
        np.random.seed(1)
        min_mean = 0.0001
        max_mean = 10
        for _num_rand_vals in range(10):
            event_rate_per_sec = np.random.uniform(min_mean, max_mean)
            duration_mean = np.random.uniform(0, max_mean)
            training_start_time_distr = ConstantAsyncTrainingStartTimeDistrConfig(
                training_rate=event_rate_per_sec
            )
            duration_distr = PerExampleGaussianDurationDistributionConfig(
                training_duration_mean=duration_mean, training_duration_sd=0
            )
            distr = AsyncTrainingEventGenerator(
                **OmegaConf.structured(
                    AsyncTrainingEventGeneratorConfig(
                        training_start_time_distribution=training_start_time_distr,
                        duration_distribution_generator=duration_distr,
                    )
                )
            )
            # generate duration and time_to_next_event_start two times
            for _ in range(2):
                self.assertAlmostEqual(
                    distr.training_duration(num_training_examples=1),
                    duration_mean,
                    delta=1e-4,
                )
                self.assertEqual(
                    distr.time_to_next_event_start(), 1 / event_rate_per_sec
                )

    def test_training_duration_min_bound(self) -> None:
        """Check that training_duration_min bound is followed"""
        np.random.seed(1)
        max_mean = 10
        max_sd = 1
        # generate 10 random event generators
        for _num_rand_vals in range(10):
            duration_mean = np.random.uniform(0, max_mean)
            duration_sd = np.random.uniform(0, max_sd)
            # choose a duration_min that is likely to be hit often
            duration_min = duration_mean
            per_user_duration_distr = PerUserGaussianDurationDistribution(
                **OmegaConf.structured(
                    PerUserGaussianDurationDistributionConfig(
                        training_duration_mean=duration_mean,
                        training_duration_sd=duration_sd,
                        training_duration_min=duration_min,
                    )
                )
            )
            per_example_duration_distr = PerExampleGaussianDurationDistribution(
                **OmegaConf.structured(
                    PerExampleGaussianDurationDistributionConfig(
                        training_duration_mean=duration_mean,
                        training_duration_sd=duration_sd,
                        training_duration_min=duration_min,
                    )
                )
            )
            # generate 100 random num_examples
            for _ in range(100):
                num_examples = np.random.randint(low=1, high=1e5)
                per_user_gaussian_duration = per_user_duration_distr.training_duration(
                    num_training_examples=num_examples
                )
                self.assertGreaterEqual(per_user_gaussian_duration, duration_min)
                # for per-example training duration, duration-min_bound applies to each example
                # while training_duration() returns time for each user
                # so actual bound is
                per_example_gaussian_duration = (
                    per_example_duration_distr.training_duration(
                        num_training_examples=num_examples
                    )
                )
                self.assertGreaterEqual(
                    per_example_gaussian_duration, duration_min * num_examples
                )

    @testutil.data_provider(
        lambda: (
            {
                "start_time_distr_config_class": PoissonAsyncTrainingStartTimeDistrConfig,
                "start_time_distr_class": PoissonAsyncTrainingStartTimeDistr,
            },
            {
                "start_time_distr_config_class": ConstantAsyncTrainingStartTimeDistrConfig,
                "start_time_distr_class": ConstantAsyncTrainingStartTimeDistr,
            },
        )
    )
    def test_string_conversion(
        self,
        start_time_distr_config_class: AsyncTrainingStartTimeDistrConfig,
        start_time_distr_class: Type,
    ) -> None:
        """Check that strings are correctly converted to TrainingEventGenerator"""
        training_rate = 1
        duration_mean_sec = 1
        training_duration_sd = 1
        training_start_time_distr = OmegaConf.structured(start_time_distr_config_class)
        training_start_time_distr.training_rate = training_rate
        duration_distr = PerExampleGaussianDurationDistributionConfig(
            training_duration_mean=duration_mean_sec,
            training_duration_sd=training_duration_sd,
        )
        self.assertEqual(
            AsyncTrainingEventGenerator(
                **OmegaConf.structured(
                    AsyncTrainingEventGeneratorConfig(
                        training_start_time_distribution=training_start_time_distr,
                        duration_distribution_generator=duration_distr,
                    )
                )
            )._training_start_time_distr.__class__,
            start_time_distr_class,
        )
