import os
import json
import random
import unittest
import numpy as np
import sys
import tempfile
import shutil
from unittest.mock import MagicMock

import ray
from ray.tune.result import TRAINING_ITERATION
from ray.tune.schedulers import (HyperBandScheduler, AsyncHyperBandScheduler,
                                 Populat
ionBasedTraining, MedianStoppingRule,
                                 TrialScheduler, HyperBandForBOHB)

from ray.tune.schedulers.pbt import explore
from ray.tune.trial import Trial, Checkpoint
from ray.tune.trial_executor import TrialExecutor
from ray.tune.resources import Reso
urces

from ray.rllib import _register_all
_register_all()


def result(t, rew):
    return dict(
        time_total_s=t, episode_reward_mean=rew, training_iteration=int(t))


def mock_trial_runner(trials=None):
    trial_runner = MagicMock()
    trial_runner.get_trials.return_value = trials or []
    return trial_runner


class EarlyStoppingSuite(unittest.TestCase):
    def setUp(self):
        ray.init()

    def tearDown(self):
        ray.shutdown()
        _register_all()  # re-register the evicted objects

    def basicSetup(self, rule):
        t1 = Trial("PPO")  # mean is 450, max 900, t_max=10
        t2 = Trial("PPO")  # mean is 450, max 450, t_max=5
        runner = mock_trial_runner()
        for i in range(10):
            r1 = result(i, i * 100)
            print("basicSetup:", i)
            self.assertEqual(
                rule.on_trial_result(runner, t1, r1), TrialScheduler.CONTINUE)
        for i in range(5):
            r2 = result(i, 450)
            self.assertEqual(
                rule.on_trial_result(runner, t2, r2), TrialScheduler.CONTINUE)
        return t1, t2

    def testMedianStoppingConstantPerf(self):
        rule = MedianStoppingRule(grace_period=0, min_samples_required=1)
        t1, t2 = self.basicSetup(rule)
        runner = mock_trial_runner()
        rule.on_trial_complete(runner, t1, result(10, 1000))
        self.assertEqual(
            rule.on_trial_result(runner, t2, result(5, 450)),
            TrialScheduler.CONTINUE)
        self.assertEqual(
            rule.on_trial_result(runner, t2, result(6, 0)),
            TrialScheduler.CONTINUE)
        self.assertEqual(
            rule.on_trial_result(runner, t2, result(10, 450)),
            TrialScheduler.STOP)

    def testMedianStoppingOnCompleteOnly(self):
        rule = MedianStoppingRule(grace_period=0, min_samples_required=1)
        t1, t2 = self.basicSetup(rule)
        runner = mock_trial_runner()
        self.assertEqual(
            rule.on_trial_result(runner, t2, result(100, 0)),
            TrialScheduler.CONTINUE)
        rule.on_trial_complete(runner, t1, result(101, 1000))
        self.assertEqual(
            rule.on_trial_result(runner, t2, result(101, 0)),
            TrialScheduler.STOP)

    def testMedianStoppingGracePeriod(self):
        rule = MedianStoppingRule(grace_period=2.5, min_samples_required=1)
        t1, t2 = self.basicSetup(rule)
        runner = mock_trial_runner()
        rule.on_trial_complete(runner, t1, result(10, 1000))
        rule.on_trial_complete(runner, t2, result(10, 1000))
        t3 = Trial("PPO")
        self.assertEqual(
            rule.on_trial_result(runner, t3, result(1, 10)),
            TrialScheduler.CONTINUE)
        self.assertEqual(
            rule.on_trial_result(runner, t3, result(2, 10)),
            TrialScheduler.CONTINUE)
        self.assertEqual(
            rule.on_trial_result(runner, t3, result(3, 10)),
            TrialScheduler.STOP)

    def testMedianStoppingMinSamples(self):
        rule = MedianStoppingRule(grace_period=0, min_samples_required=2)
        t1, t2 = self.basicSetup(rule)
        runner = mock_trial_runner()
        rule.on_trial_complete(runner, t1, result(10, 1000))
        t3 = Trial("PPO")
        # Insufficient samples to evaluate t3
        self.assertEqual(
            rule.on_trial_result(runner, t3, result(5, 10)),
            TrialScheduler.CONTINUE)
        rule.on_trial_complete(runner, t2, result(5, 1000))
        # Sufficient samples to evaluate t3
        self.assertEqual(
            rule.on_trial_result(runner, t3, result(5, 10)),
            TrialScheduler.STOP)

    def testMedianStoppingUsesMedian(self):
        rule = MedianStoppingRule(grace_period=0, min_samples_required=1)
        t1, t2 = self.basicSetup(rule)
        runner = mock_trial_runner()
        rule.on_trial_complete(runner, t1, result(10, 1000))
        rule.on_trial_complete(runner, t2, result(10, 1000))
        t3 = Trial("PPO")
        self.assertEqual(
            rule.on_trial_result(runner, t3, result(1, 260)),
            TrialScheduler.CONTINUE)
        self.assertEqual(
            rule.on_trial_result(runner, t3, result(2, 260)),
            TrialScheduler.STOP)

    def testMedianStoppingSoftStop(self):
        rule = MedianStoppingRule(
            grace_period=0, min_samples_required=1, hard_stop=False)
        t1, t2 = self.basicSetup(rule)
        runner = mock_trial_runner()
        rule.on_trial_complete(runner, t1, result(10, 1000))
        rule.on_trial_complete(runner, t2, result(10, 1000))
        t3 = Trial("PPO")
        self.assertEqual(
            rule.on_trial_result(runner, t3, result(1, 260)),
            TrialScheduler.CONTINUE)
        self.assertEqual(
            rule.on_trial_result(runner, t3, result(2, 260)),
            TrialScheduler.PAUSE)

    def _test_metrics(self, result_func, metric, mode):
        rule = MedianStoppingRule(
            grace_period=0,
            min_samples_required=1,
            time_attr="training_iteration",
            metric=metric,
            mode=mode)
        t1 = Trial("PPO")  # mean is 450, max 900, t_max=10
        t2 = Trial("PPO")  # mean is 450, max 450, t_max=5
        runner = mock_trial_runner()
        for i in range(10):
            self.assertEqual(
                rule.on_trial_result(runner, t1, result_func(i, i * 100)),
                TrialScheduler.CONTINUE)
        for i in range(5):
            self.assertEqual(
                rule.on_trial_result(runner, t2, result_func(i, 450)),
                TrialScheduler.CONTINUE)
        rule.on_trial_complete(runner, t1, result_func(10, 1000))
        self.assertEqual(
            rule.on_trial_result(runner, t2, result_func(5, 450)),
            TrialScheduler.CONTINUE)
        self.assertEqual(
            rule.on_trial_result(runner, t2, result_func(6, 0)),
            TrialScheduler.CONTINUE)

    def testAlternateMetrics(self):
        def result2(t, rew):
            return dict(training_iteration=t, neg_mean_loss=rew)

        self._test_metrics(result2, "neg_mean_loss", "max")

    def testAlternateMetricsMin(self):
        def result2(t, rew):
            return dict(training_iteration=t, mean_loss=-rew)

        self._test_metrics(result2, "mean_loss", "min")


class _MockTrialExecutor(TrialExecutor):
    def start_trial(self, trial, checkpoint_obj=None):
        trial.logger_running = True
        trial.restored_checkpoint = checkpoint_obj.value
        trial.status = Trial.RUNNING

    def stop_trial(self, trial, error=False, error_msg=None, stop_logger=True):
        trial.status = Trial.ERROR if error else Trial.TERMINATED
        if stop_logger:
            trial.logger_running = False

    def restore(self, trial, checkpoint=None):
        pass

    def save(self, trial, type=Checkpoint.PERSISTENT, result=None):
        return Checkpoint(Checkpoint.PERSISTENT, trial.trainable_name, result)

    def reset_trial(self, trial, new_config, new_experiment_tag):
        return False


class _MockTrialRunner():
    def __init__(self, scheduler):
        self._scheduler_alg = scheduler
        self.trials = []
        self.trial_executor = _MockTrialExecutor()

    def process_action(self, trial, action):
        if action == TrialScheduler.CONTINUE:
            pass
        elif action == TrialScheduler.PAUSE:
            self._pause_trial(trial)
        elif action == TrialScheduler.STOP:
            self.trial_executor.stop_trial(trial)

    def stop_trial(self, trial):
        if trial.status in [Trial.ERROR, Trial.TER