import json

from suggestion.models import Study
from suggestion.models import TrialMetric
from suggestion.early_stop_algorithm.abstract_early_stop import (
    AbstractEarlyStopAlgorithm,
)


class EarlyStopDescendingAlgorithm(AbstractEarlyStopAlgorithm):
    def get_early_stop_trials(self, trials):
        result = []

        for trial in trials:
            study = Study.objects.get(name=trial.study_name)
            study_configuration_json = json.loads(study.study_configuration)
            study_goal = study_configuration_json["goal"]

            metrics = TrialMetric.objects.filter(trial_id=trial.id).order_by(
                "-training_step"
            )
            metrics = [metric for metric in metrics]

            if len(metrics) >= 2:
                if study_goal == "MAXIMIZE":
                    if metrics[0].objective_value < metrics[1].objective_value:
                        result.append(trial)
                elif study_goal == "MINIMIZE":
                    if metrics[0].objective_value > metrics[1].objective_value:
                        result.append(trial)

        return result
