import json

from django.test import TestCase

from suggestion.models import Study
from suggestion.models import Trial
from suggestion.models import TrialMetric
from suggestion.early_stop_algorithm.early_stop_first_trial import EarlyStopFirstTrialAlgorithm


class EarlyStopFirstTrialAlgorithmTest(TestCase):
  def setUp(self):
    study_configuration_json = {
        "goal":
        "MAXIMIZE",
        "maxTrials":
        5,
        "maxParallelTrials":
        1,
        "params": [{
            "parameterName": "hidden1",
            "type": "INTEGER",
            "minValue": 40,
            "maxValue": 400,
            "scalingType": "LINEAR"
        }]
    }
    study_configuration = json.dumps(study_configuration_json)
    self.study = Study.create("RandomSearchStudy", study_configuration)
    trial1 = Trial.create(self.study.id, "RandomSearchTrial1")
    trial2 = Trial.create(self.study.id, "RandomSearchTrial2")
    self.trials = [trial1, trial2]
    TrialMetric.create(trial1.id, 10, 0.5)
    TrialMetric.create(trial1.id, 20, 0.6)
    TrialMetric.create(trial2.id, 10, 0.6)
    TrialMetric.create(trial2.id, 20, 0.5)

  def tearDown(self):
    pass

  # Test EarlyStopFirstTrialAlgorithm
  def test_get_early_stop_trials1(self):
    algorithm = EarlyStopFirstTrialAlgorithm()
    early_stop_trials = algorithm.get_early_stop_trials(self.trials)
    self.assertEqual(len(early_stop_trials), 1)
    self.assertEqual(early_stop_trials[0].name, "RandomSearchTrial1")
