text: """
  Neuraxle's AutoML Scope Manager Classes
  ====================================
  
  
  Scope manager objects used by AutoML algorithm classes.
  
  Although the name "manager" for a class is disliked by the Neuraxle community,
  it is used as these are Pythonic context managers, as described in
  `Python's contextlib documentation <https://docs.python.org/3/library/contextlib.html>`__.
  
  Classes are splitted like this for the AutoML:
  - Projects
  - Clients
  - Rounds (runs)
  - Trials
  - TrialSplits
  - MetricResults
  
  
  ..
      Copyright 2022, Neuraxio Inc.
  
      Licensed under the Apache License, Version 2.0 (the "License");
      you may not use this file except in compliance with the License.
      You may obtain a copy of the License at
  
          http://www.apache.org/licenses/LICENSE-2.0
  
      Unless required by applicable law or agreed to in writing, software
      distributed under the License is distributed on an "AS IS" BASIS,
      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
      See the License for the specific language governing permissions and
      limitations under the License.
  
  
  """
  
  import copy
  import gc
  import hashlib
  import os
  import typing
  from abc import abstractmethod
  from collections import OrderedDict
  from types import TracebackType
  from typing import (Callable, ContextManager, Dict, Generic, Iterable, List,
                      Optional, Tuple, Type, TypeVar)
  
  import numpy as np
  from neuraxle.base import BaseService
  from neuraxle.base import ExecutionContext as CX
  from neuraxle.base import Flow, TrialStatus, _CouldHaveContext
  from neuraxle.hyperparams.space import (FlatDict, HyperparameterSamples,
                                          HyperparameterSpace, RecursiveDict)
  from neuraxle.metaopt.data.vanilla import (DEFAULT_CLIENT, DEFAULT_PROJECT,
                                             RETRAIN_TRIAL_SPLIT_ID,
                                             AutoMLContext, BaseDataclass,
                                             BaseHyperparameterOptimizer,
                                             ClientDataclass,
                                             HyperparamsRepository,
                                             MetricResultsDataclass,
                                             ProjectDataclass, RootDataclass,
                                             RoundDataclass, ScopedLocation,
                                             ScopedLocationAttr,
                                             ScopedLocationAttrInt,
                                             SubDataclassT, TrialDataclass,
                                             TrialSplitDataclass,
                                             VanillaHyperparamsRepository,
                                             dataclass_2_id_attr)
  
  SubAggregateT = TypeVar('SubAggregateT', bound=Optional['BaseAggregate'])
  ParentAggregateT = TypeVar('ParentAggregateT', bound=Optional['BaseAggregate'])
  
  
  def _with_method_as_context_manager(
      func: Callable[['BaseAggregate'], 'BaseAggregate']
  ) -> Callable[['BaseAggregate'], ContextManager['SubAggregateT']]:
      """
      This is a method to be used as a context manager.
      This will sync items with the repos.
      Example:
  
      .. code-block:: python
          with obj.func() as managed_context:
              # obj.__enter__() is called
              managed_context.do_something()
          # obj.__exit__() is called
  
      .. seealso::
          :class:`neuraxle.metaopt.data.aggregates.BaseAggregate`
          :func:`neuraxle.metaopt.data.aggregates.BaseAggregate.__enter__`
          :func:`neuraxle.metaopt.data.aggregates.BaseAggregate.__exit__`
          :class:`neuraxle.metaopt.auto_ml.Trainer`
          :func:`neuraxle.metaopt.auto_ml.Trainer.train`
  
      """
      # Also add the docstring of `_with_method_as_context_manager' to the docstring of the `func`:
      func.__doc__ = (func.__doc__ or "") + _with_method_as_context_manager.__doc__
      return func
  
  
  class BaseAggregate(_CouldHaveContext, BaseService, ContextManager[SubAggregateT], Generic[ParentAggregateT, SubAggregateT, SubDataclassT]):
      """
      Base class for aggregated objects using the repo and the dataclasses to manipulate them.
  
      For more information, read this `article by Martin Fowler on DDD Aggregates <https://martinfowler.com/bliki/DDD_Aggregate.html>`_.
      """
  
      def __init__(self, _dataclass: SubDataclassT, context: AutoMLContext, is_deep=False, parent: ParentAggregateT = None):
          BaseService.__init__(self, name=f"{self.__class__.__name__}_{_dataclass.get_id()}")
          _CouldHaveContext.__init__(self)
          self._dataclass: SubDataclassT = _dataclass
          self._spare: SubDataclassT = copy.copy(_dataclass).shallow()
          # TODO: pre-push context to allow for dc auto-loading and easier parent auto-loading?
          self.context: AutoMLContext = context.push_attr(_dataclass)
          self.loc: ScopedLocation = self.context.loc.copy()
          self.is_deep = is_deep
          self._parent: ParentAggregateT = parent
  
          self.service_assertions = [Flow, HyperparamsRepository]
          self._invariant()
  
      @staticmethod
      def from_context(context: AutoMLContext, is_deep=True) -> 'BaseAggregate':
          """
          Return an aggregate from a context.
          Requirement: Have already pre-push the attr of dataclass into
          the context before calling this.
          """
          _dataclass = context.load_dc(deep=True)
          aggregate_class = dataclass_2_aggregate[_dataclass.__class__]
          return aggregate_class(_dataclass, context.pop_attr(), is_deep=is_deep)
  
      @classmethod
      def dummy(cls: Type['BaseAggregate'], context: CX = None) -> 'BaseAggregate':
          """
          Create a dummy object of the desired type for testing purposes.
          The parent subtree will be created in a temporary repository,
          or in a real repository if the context is provided.
          """
          context = context or AutoMLContext.from_context(context)
  
          # create dataclass tree:
          dataclasses: List[BaseDataclass] = dataclass_2_id_attr.keys()
          root_dc: RootDataclass = RootDataclass()
          target_dc_class: Type[BaseDataclass] = aggregate_2_dataclass[cls]
          prev_dc: BaseDataclass = root_dc
          for _dataclass in dataclasses:
              # loop for every type of dataclass and create them until we find the one we want:
              dc: BaseDataclass = _dataclass()
              prev_dc.store(dc)
  
              prev_dc = dc
              context = context.push_attr(dc)
              if isinstance(dc, target_dc_class):
                  break
  
          # create the aggregate:
          context.repo.save(root_dc, ScopedLocation(), deep=True)
          return cls.from_context(context, is_deep=False)
  
      def _invariant(self):
          _type: Type[SubDataclassT] = self.dataclass
          self._assert(isinstance(self._dataclass, _type),
                       f"self._dataclass should be of type {_type.__name__} but is of type "
                       f"{self._dataclass.__class__.__name__}", self.context)
          self._assert(isinstance(self.context, AutoMLContext),
                       f"self.context should be of type AutoMLContext but is of type "
                       f"{self.context.__class__.__name__}", self.context)
          self._assert(isinstance(self.loc, ScopedLocation),
                       f"self.loc should be of type ScopedLocation but is of type "
                       f"{self.loc.__class__.__name__}", self.context)
          self._assert(self.loc.as_list() == self.context.loc.as_list(),
                       f"{self.loc} should be equal to {self.context.loc}", self.context)
          self._assert(self.loc[-1] == self._dataclass.get_id(),
                       f"{self.loc}'s last attr should be equal to self._dataclass.get_id() "
                       f"and the id is {self._dataclass.get_id()}", self.context)
          self._assert_at_lifecycle(self.context)
  
      def without_context(self) -> 'BaseAggregate':
          """
          Return a copy of this aggregate without the context.
          Useful for initializing a temporary aggregate,
          such as a filtered or reduced aggregate without all its subaggregates
          to disallow saving the reduced aggregate.
          """
          self_copy = copy.copy(self)
          self_copy.context = None
          return self_copy
  
      @property
      def flow(self) -> Flow:
          return self.context.flow
  
      @property
      def repo(self) -> HyperparamsRepository:
          return self.context.repo
  
      def subaggregate(self, _dataclass: SubDataclassT, context: AutoMLContext, is_deep=False, parent: ParentAggregateT = None) -> SubAggregateT:
          return aggregate_2_subaggregate[self.__class__](_dataclass, context, is_deep, parent)
  
      @property
      def parent(self) -> ParentAggregateT:
          if self._parent is None and not isinstance(self, Root):
              self._parent = self.from_context(self.context.pop_attr(), is_deep=False)
              self._assert(isinstance(self._parent, self.parentaggregate),
                           f"parent should be of type {self.parentaggregate.__name__} but is of type {self._parent.__class__.__name__}", self.context)
          return self._parent
  
      @property
      def parentaggregate(self) -> Type[ParentAggregateT]:
          return {v: k for k, v in aggregate_2_subaggregate.items()}[self.__class__]
  
      @property
      def dataclass(self) -> Type[BaseDataclass]:
          return aggregate_2_dataclass[self.__class__]
  
      @property
      def subdataclass(self) -> Type[SubDataclassT]:
          return aggregate_2_dataclass[aggregate_2_subaggregate[self.__class__]]
  
      def sanitize_metric_name(self, metric_name: str = None):
          """
          If the argument metric is None, the optimizer's metric is taken.
          If the optimizer's metric is None, the parent metric is taken.
          """
          if metric_name is not None:
              return metric_name
          elif hasattr(self._dataclass, "main_metric_name") and self._dataclass.main_metric_name is not None:
              return self._dataclass.main_metric_name
          return metric_name or (self.parent.sanitize_metric_name(metric_name) if self.parent is not None else None)
  
      def refresh(self, deep: bool = True):
          with self.context.lock:
              _new_dc: SubDataclassT = self.repo.load(self.loc, deep=deep)
  
              def _fail_on_unsaved_changes(_self_now: SubDataclassT, _self_before: SubDataclassT, _self_new_loaded: SubDataclassT):
                  _has_self_changed = _self_now != _self_before
                  if _has_self_changed:
                      _is_new_dc_different_from_now = _self_now != _self_new_loaded
                      if _is_new_dc_different_from_now:
  
                          raise RuntimeError(f"{_self_now.__class__.__name__} with id `{_self_now.get_id()}` "
                                             f"has unsaved changes while it tried to refresh itself. "
                                             f"Before changes: {_self_before}. "
                                             f"After changes, that is not saved: {_self_now}. "
                                             f"DC that we tried to load: {_self_new_loaded}. ")
              _fail_on_unsaved_changes(
                  self._dataclass.shallow(),
                  self._spare.shallow(),
                  _new_dc.shallow()
              )
  
              self._dataclass = _new_dc
  
          self.is_deep = deep
          self._invariant()
  
      def save(self, deep: bool = True) -> 'BaseAggregate':
          if deep and deep != self.is_deep:
              raise ValueError(
                  f"Cannot save {str(self)} with deep=True when self "
                  f"is not already deep. You might want to use self.refresh(deep=True) at "
                  f"some point to refresh self before saving deeply then.")
  
          self._invariant()
          self._spare = copy.copy(self._dataclass).shallow()
  
          with self.context.lock:
              self.repo.save(self._dataclass, self.loc, deep=deep)
          return self
  
      def save_subaggregate(self, subagg: SubAggregateT, deep=False) -> 'BaseAggregate':
          self._dataclass.store(subagg._dataclass)
  
          with self.context.lock:
              self.save(deep=False)
              subagg.save(deep=deep)
          return self
  
      def __enter__(self) -> SubAggregateT:
          # self.context.free_scoped_logger_handler_file()
          self._invariant()
          self._managed_resource._invariant()
  
          self._managed_resource.context.add_scoped_logger_file_handler()
  
          return self._managed_resource
  
      def __exit__(
          self,
          exc_type: Optional[Type[BaseException]],
          exc_val: Optional[BaseException],
          exc_tb: Optional[TracebackType]
      ) -> Optional[bool]:
          self._managed_resource.context.free_scoped_logger_file_handler()
          # self.context.add_scoped_logger_file_handler()
  
          handled_err: bool = self._release_managed_subresource(self._managed_resource, exc_val)
          return handled_err
  
      @_with_method_as_context_manager
      def managed_subresource(self, *args, **kwds) -> SubAggregateT:
          self._managed_subresource(*args, **kwds)
          return self
  
      def _managed_subresource(self, *args, **kwds) -> ContextManager[SubAggregateT]:
          self._invariant()
          self._managed_resource: SubAggregateT = self._acquire_managed_subresource(
              *args, **kwds)
          return self
  
      @abstractmethod
      def _acquire_managed_subresource(self, *args, **kwds) -> SubAggregateT:
          """
          Acquire a new subaggregate that is managed such that it is deep saved at the
          beginning.
          """
          # Create subaggregate:
          with self.context.lock:
              self.refresh(self.is_deep)
              subdataclass: SubDataclassT = self.repo.load(*args, **kwds)
              subagg: SubAggregateT = self.subaggregate(subdataclass, self.context)
              self.save_subaggregate(subagg, deep=False)
              return subagg
  
      def _release_managed_subresource(self, resource: SubAggregateT, e: Exception = None) -> bool:
          """
          Release a subaggregate that was acquired with managed_subresource. The subaggregate
          normally has already saved itself. We may update things again here if needed.
  
          Exceptions may be handled here. If handled, return True, if not, then return False.
          """
          with self.context.lock:
              self.refresh(self.is_deep)
              self.save(False)
          handled_error = e is None
          return handled_error
  
      def __len__(self) -> int:
          return len(self._dataclass.get_sublocation())
  
      def __iter__(self) -> Iterable[SubAggregateT]:
          if not self.is_deep:
              self.refresh(True)
          for subdataclass in self._dataclass.get_sublocation_values():
              if subdataclass is not None:
                  yield self.subaggregate(subdataclass, self.context, is_deep=True, parent=self)
  
      def __getitem__(self, item: int) -> 'Trial':
          """
          Get trial at the given index.
  
          :param item: trial index
          :return:
          """
          return self.subaggregate(self._dataclass.get_sublocation()[item], self.context, self.is_deep, self)
  
      def __str__(self) -> str:
          return self.__repr__()
  
      def __repr__(self) -> str:
          if self.is_deep:
              prefix = "<deep>"
              subobjs = str([str(t) for t in self.__iter__()])
          else:
              prefix = "<shallow>"
              subobjs = f"{self._dataclass.get_sublocation()}>"
          return (
              f"{self.__class__.__name__}{prefix}("
              f"id={self._dataclass.get_id()}, {subobjs}"
              f")"
          )
  
      def __eq__(self, other: object) -> bool:
          return self._dataclass == other._dataclass
  
  
  class Root(BaseAggregate[None, 'Project', RootDataclass]):
  
      def save(self, deep: bool = False):
          if deep:
              for p in self.projects:
                  p.save(deep=p.is_deep)
  
      @property
      def projects(self) -> List['Project']:
          return list(self)
  
      @staticmethod
      def from_repo(context: CX, repo: HyperparamsRepository) -> 'Root':
          _dataclass: RootDataclass = repo.load(ScopedLocation())
          automl_context: AutoMLContext = AutoMLContext.from_context(context, repo)
          return Root(_dataclass, automl_context)
  
      @staticmethod
      def vanilla(context: CX = None) -> 'Root':
          if context is None:
              context = CX()
          vanilla_repo: HyperparamsRepository = VanillaHyperparamsRepository(
              os.path.join(context.get_path(), "hyperparams"))
          return Root.from_repo(context, vanilla_repo)
  
      @_with_method_as_context_manager
      def get_project(self, name: str) -> 'Root':
          self._managed_subresource(project_name=name)
          return self
  
      @_with_method_as_context_manager
      def default_project(self) -> 'Root':
          self._managed_subresource(project_name=DEFAULT_PROJECT)
          return self
  
      def _acquire_managed_subresource(self, project_name: str) -> 'Project':
          project_loc: ScopedLocation = self.loc.with_id(project_name)
          return super()._acquire_managed_subresource(project_loc)
  
  
  class Project(BaseAggregate[None, 'Client', ProjectDataclass]):
  
      @_with_method_as_context_manager
      def get_client(self, name: str) -> 'Project':
          self._managed_subresource(client_name=name)
          return self
  
      @_with_method_as_context_manager
      def default_client(self) -> 'Project':
          self._managed_subresource(client_name=DEFAULT_CLIENT)
          return self
  
      def _acquire_managed_subresource(self, client_name: str) -> 'Client':
          client_loc: ScopedLocation = self.loc.with_id(client_name)
          return super()._acquire_managed_subresource(client_loc)
  
  
  class Client(BaseAggregate[Project, 'Round', ClientDataclass]):
  
      @_with_method_as_context_manager
      def new_round(self, main_metric_name: str) -> 'Client':
          self._managed_subresource(new_round=True, main_metric_name=main_metric_name)
          return self
  
      @_with_method_as_context_manager
      def resume_last_round(self) -> 'Client':
          self._managed_subresource(new_round=False)
          return self
  
      @_with_method_as_context_manager
      def optim_round(self, new_round: bool, main_metric_name: str) -> 'Client':
          self._managed_subresource(new_round=new_round, main_metric_name=main_metric_name)
          return self
  
      def _acquire_managed_subresource(self, new_round: bool = True, main_metric_name: str = None) -> 'Round':
          with self.context.lock:
              self.refresh(self.is_deep)
  
              # Get new round loc:
              round_id: int = self._dataclass.get_next_i()
              if not new_round:
                  round_id = max(0, round_id - 1)
              if round_id == 0:
                  new_round = True
              round_loc: ScopedLocation = self.loc.with_id(round_id)
  
              # Get round to return:
              _round_dataclass: RoundDataclass = self.repo.load(round_loc, deep=True)
              if main_metric_name is not None:
                  _round_dataclass.main_metric_name = main_metric_name
  
              subagg: Round = Round(_round_dataclass, self.context, is_deep=True)
              if new_round:
                  self.save_subaggregate(subagg, deep=True)
              return subagg
  
  
  class Round(BaseAggregate[Client, 'Trial', RoundDataclass]):
  
      def with_optimizer(
          self,
          hp_optimizer: BaseHyperparameterOptimizer,
          hp_space: HyperparameterSpace
      ) -> 'Round':
          self.hp_optimizer: BaseHyperparameterOptimizer = hp_optimizer
          self.hp_space: HyperparameterSpace = hp_space
          return self
  
      def with_metric(self, metric_name: str) -> 'Round':
          self._dataclass.main_metric_name = metric_name
          return self.save(False)
  
      @property
      def main_metric_name(self) -> str:
          return self._dataclass.main_metric_name
  
      @_with_method_as_context_manager
      def new_rvs_trial(self, continue_on_error: bool = False) -> 'Round':
          self._managed_subresource(new_trial=True, continue_on_error=continue_on_error)
          return self
  
      @_with_method_as_context_manager
      def last_trial(self, continue_on_error: bool = False) -> 'Round':
          self._managed_subresource(new_trial=False, continue_on_error=continue_on_error)
          return self
  
      @_with_method_as_context_manager
      def refitting_best_trial(self) -> 'Round':
          self.refresh(deep=True)
          self._managed_subresource(new_trial=None, continue_on_error=False)
          return self
  
      @property
      def _trials(self) -> List['Trial']:
          if not self.is_deep:
              self.refresh(True)
          return [Trial(t, self.context) for t in self._dataclass.get_sublocation()]
  
      @property
      def get_id(self) -> int:
          return self._dataclass.get_id()
  
      def _acquire_managed_subresource(self, new_trial=True, continue_on_error: bool = False) -> 'Trial':
          """
          Get a trial.
  
          :param new_trial: If True, will create a new trial. If false, will load the last trial. If None, will load the best trial.
          :param continue_on_error: If True, will continue to the next trial if the current trial fails.
                                    Otherwise, will let the exception be raised for the failure (won't catch).
          """
          with self.context.lock:
              self.refresh(self.is_deep)
              # self.context.add_scoped_logger_file_handler()
  
              # Get new trial loc:
              trial_id: int = self._dataclass.get_next_i()
              if new_trial is None:
                  trial_id: int = self.get_best_trial()._dataclass.get_id()
              elif not new_trial:
                  # Try get last trial
                  trial_id = max(0, trial_id - 1)
              trial_loc = self.loc.with_id(trial_id)
  
              # Get trial to return:
              _trial_dataclass: TrialDataclass = self.repo.load(trial_loc, deep=True)
              new_hps: HyperparameterSamples = self.hp_optimizer.find_next_best_hyperparams(self)
              _trial_dataclass.hyperparams = new_hps
              self.flow.log_planned(new_hps)
              subagg: Trial = Trial(_trial_dataclass, self.context, is_deep=True)
  
              if continue_on_error:
                  subagg.continue_loop_on_error()
  
              if new_trial or trial_id == 0:
                  self.save_subaggregate(subagg, deep=True)
              return subagg
  
      def _release_managed_subresource(self, resource: 'Trial', e: Exception = None) -> bool:
          """
          Release a subaggregate that was acquired with managed_subresource. The subaggregate
          normally has already saved itself. We may update things again here if needed.
  
          Exceptions may be handled here. If handled, return True, if not, then return False.
          """
          resource: Trial = resource  # typing
          handled_exception = False
          with self.context.lock:
              self.refresh(True)
  
              is_all_success: bool = resource.are_all_splits_successful()
              is_all_failure: bool = resource.are_all_splits_failures()
  
              if e is None:
                  handled_exception = True
                  if is_all_success:
                      resource._set_success()
                  elif is_all_failure:
                      e = RuntimeError("All trial splits failed for this trial.")
                      resource._set_failed(e)
              else:
                  resource._set_failed(e)
  
              self.save_subaggregate(resource, deep=resource.is_deep)
  
              if not is_all_failure and len(self) > 0:
                  main_metric_name = self.main_metric_name
                  self.flow.log('Finished round hp search!')
                  self.flow.log_best_hps(
                      main_metric_name,
                      self.get_best_hyperparams(main_metric_name)
                  )
              else:
                  self.flow.log_failure(e)
  
              self.save(False)
          return handled_exception
  
      def get_best_hyperparams(self, metric_name: str = None) -> HyperparameterSamples:
          """
          Get best hyperparams from all trials.
  
          :return:
          """
          if not self.is_deep:
              self.refresh(True)
          metric_name = self.sanitize_metric_name(metric_name)
  
          if len(self) == 0:
              return HyperparameterSamples()
          elif len(self) == 1:
              return self._trials[0].get_hyperparams()
          else:
              return self.get_best_trial(metric_name).get_hyperparams()
  
      def get_best_trial(self, metric_name: str = None) -> Optional['Trial']:
          """
          :return: trial with best score from all trials
          """
          best_score, best_trial = None, None
          metric_name = self.sanitize_metric_name(metric_name)
  
          if len(self) == 0:
              raise Exception('Could not get best trial because there were no successful trial.')
  
          for trial in self._trials:
              trial_score = trial.get_avg_validation_score(metric_name)
  
              _has_better_score = best_score is None or (
                  trial_score is not None and trial.is_success() and (
                      self.is_higher_score_better(metric_name) == (trial_score > best_score)
                  )
              )
  
              if _has_better_score:
                  best_score = trial_score
                  best_trial = trial
  
          return best_trial
  
      def is_higher_score_better(self, metric_name: str = None) -> bool:
          """
          Return true if higher score is better. If metric_name is None, the optimizer's
          metric is taken.
  
          :return
          """
          if len(self) == 0:
              return False
  
          metric_name = self.sanitize_metric_name(metric_name)
  
          is_higher_score_better_attr = "is_higher_score_better_attr_" + metric_name
          if not hasattr(self, is_higher_score_better_attr):
              setattr(
                  self,
                  is_higher_score_better_attr,
                  self._trials[-1].is_higher_score_better(metric_name)
              )
          return getattr(self, is_higher_score_better_attr)
  
      def append(self, trial: 'Trial'):
          """
          Add a new trial. Will also save the trial shallowly.
  
          :param trial: new trial
          :return:
          """
          self.save_subaggregate(trial, deep=False)
  
      def copy(self) -> 'Round':
          return Round(copy.deepcopy(self._dataclass), self.context.copy().with_loc(self.loc.popped()))
  
      def get_number_of_split(self):
          if len(self) > 0:
              return len(self[0]._validation_splits)
          return 0
  
      def get_metric_names(self) -> List[str]:
          if len(self) > 0:
              return list(self[-1]._dataclass.validation_splits[-1].metric_results.keys())
          return []
  
      def best_result_summary(self, metric_name: str = None) -> Tuple[float, ScopedLocationAttrInt, FlatDict]:
          return self.summary(metric_name)[0]
  
      def summary(
          self, metric_name: str = None
      ) -> List[Tuple[float, ScopedLocationAttrInt, FlatDict]]:
          """
          Get a summary of the round. Best score is first.
          Values in the returned triplet tuples are: (score, trial_number, hyperparams),
          sorted by score such that the best score is first.
          """
          if not self.is_deep:
              self.refresh(True)
          metric_name = self.sanitize_metric_name(metric_name)
  
          results: List[float, ScopedLocationAttrInt, FlatDict] = list()
  
          for trial in self._trials:
              # TODO: what happens to failed or ongoing trials?
              score = trial.get_avg_validation_score(metric_name)
              trial_number = trial._dataclass.trial_number
              hp = trial.get_hyperparams().to_flat_dict()
  
              results.append((score, trial_number, hp))
  
          is_reverse: bool = self.is_higher_score_better(metric_name)
          results = list(sorted(results, reverse=is_reverse))
          return results
  
      def get_all_hyperparams(self) -> List[FlatDict]:
          """
          Get all hyperparams from all trials.
  
          :return:
          """
          self.refresh(True)
  
          hyperparams: List[FlatDict] = list()
          for trial in self._trials:
              hyperparams.append(trial.get_hyperparams().to_flat_dict())
  
          return hyperparams
  
  
  class Trial(BaseAggregate[Round, 'TrialSplit', TrialDataclass]):
      """
      This class is a sub-contextualization of the :class:`HyperparameterRepository`
      class for holding a trial and manipulating it within its context.
      A Trial contains the results for each validation split.
      Each trial split contains both the training set results, and the validation set results.
  
      .. seealso::
          :class:`Trial`,
          :class:`TrialSplit`,
          :class:`HyperparamsRepository`,
          :class:`AutoML`,
          :class:`ExecutionContext`
      """
  
      def continue_loop_on_error(self) -> 'Trial':
          self._continue_loop_on_error: bool = True
          return self
  
      @property
      def _validation_splits(self) -> List['TrialSplit']:
          if not self.is_deep:
              self.refresh(True)
          return [TrialSplit(s, self.context, is_deep=self.is_deep)
                  for s in self._dataclass.validation_splits]
  
      @property
      def main_metric_name(self) -> str:
          return self.parent.main_metric_name
  
      @_with_method_as_context_manager
      def new_validation_split(self, continue_loop_on_error: bool = False) -> 'Trial':
          continue_loop_on_error = continue_loop_on_error or (
              hasattr(self, "_continue_loop_on_error") and self._continue_loop_on_error)
  
          self._managed_subresource(continue_loop_on_error=continue_loop_on_error)
          return self
  
      def retrain_split(self) -> 'Trial':
          self._managed_subresource(continue_loop_on_error=False, retrain_split=True)
          return self
  
      def _acquire_managed_subresource(self, continue_loop_on_error: bool, retrain_split=False) -> 'TrialSplit':
  
          self.error_types_to_raise = (
              SystemError, SystemExit, EOFError, KeyboardInterrupt
          ) if continue_loop_on_error else (Exception,)
  
          with self.context.lock:
              self.refresh(self.is_deep)
              # self.context.add_scoped_logger_file_handler()
  
              # Get new split loc:
              if retrain_split:
                  split_id = RETRAIN_TRIAL_SPLIT_ID
              else:
                  split_id: int = self._dataclass.get_next_i()
                  split_id = max(0, split_id)
                  if split_id == 0:
                      self.flow.log_start()
              split_loc = self.loc.with_id(split_id)
  
              # Get split to return:
              _split_dataclass: TrialSplitDataclass = self.repo.load(split_loc)
              _split_dataclass.hyperparams = self.get_hyperparams()
              subagg: TrialSplit = TrialSplit(_split_dataclass, self.context, is_deep=True)
  
              self.save_subaggregate(subagg, deep=True)
              return subagg
  
      def _release_managed_subresource(self, resource: 'TrialSplit', e: Exception = None) -> bool:
          gc.collect()
          handled_error = False
  
          with self.context.lock:
              self.refresh(self.is_deep)
              # self.context.free_scoped_logger_handler_file()
  
              if e is None:
                  resource._set_success()
                  handled_error = True
              else:
                  resource._set_failed(e)
  
                  if any((isinstance(e, c) for c in self.error_types_to_raise)):
                      self._set_failed(e)
                      handled_error = False
                  else:
                      handled_error = True
  
              if resource._dataclass.split_number == RETRAIN_TRIAL_SPLIT_ID:
                  self._dataclass.retrained_split = resource._dataclass
                  resource.save(deep=True)
                  self.save(deep=False)
              else:
                  self.save_subaggregate(resource, deep=True)
  
          return handled_error
  
      def is_success(self):
          """
          Checks if the trial is successful from its dataclass record.
          """
          return self._dataclass.status == TrialStatus.SUCCESS
  
      def get_status(self) -> TrialStatus:
          """
          Get the status of the trial.
          """
          return self._dataclass.status
  
      def are_all_splits_successful(self) -> bool:
          """
          Return true if all splits are successful.
          """
          return all(i.is_success() for i in self._validation_splits)
  
      def are_all_splits_failures(self) -> bool:
          """
          Return true if all splits are failed.
          """
          return all(not i.is_success() for i in self._validation_splits)
  
      def get_avg_validation_score(self, metric_name: str = None) -> Optional[float]:
          """
          Returns the average score for all validation splits's
          best validation score for the specified scoring metric.
  
          :return: validation score
          """
          if self.is_success():
              metric_name = self.sanitize_metric_name(metric_name)
  
              scores = [
                  val_split[metric_name].get_best_validation_score()
                  for val_split in self._validation_splits
                  if val_split.is_success() and metric_name in val_split.get_metric_names()
              ]
              scores = [s for s in scores if s is not None]
              return sum(scores) / len(scores) if len(scores) > 0 else None
  
      def get_avg_n_epoch_to_best_validation_score(self, metric_name: str = None) -> float:
          metric_name = self.sanitize_metric_name(metric_name)
  
          n_epochs = [
              val_split.get_n_epochs_to_best_validation_score(metric_name)
              for val_split in self._validation_splits if val_split.is_success()
          ]
  
          n_epochs = sum(n_epochs) / len(n_epochs) if len(n_epochs) > 0 else None
          return n_epochs
  
      def get_hyperparams(self) -> HyperparameterSamples:
          """
          Return hyperparams dict.
  
          :return: hyperparams dict
          """
          return self._dataclass.hyperparams
  
      def _set_success(self) -> 'Trial':
          """
          Set trial status to success. Must save after to ensure consistency.
  
          :return: self
          """
          self._dataclass.end(TrialStatus.SUCCESS)
          self.flow.log_success()
          return self
  
      def _set_failed(self, error: Exception) -> 'Trial':
          """
          Set failed trial with exception. Must save after to ensure consistency.
  
          :param error: catched exception
          :return: self
          """
          self._dataclass.end(TrialStatus.FAILED)
          self.flow.log_failure(exception=error)
          return self
  
      def get_trial_id(self, hp_dict: Dict):
          """
          Hash hyperparams with blake2s to create a trial hash.
  
          :param hp_dict: hyperparams dict
          :return:
          """
          current_hyperparameters_hash = hashlib.blake2s(str.encode(str(hp_dict))).hexdigest()
          return f"{self._dataclass.trial_number}_{current_hyperparameters_hash}"
  
      def is_higher_score_better(self, metric_name: str = None) -> bool:
          """
          Return if the metric is higher score is better.
  
          :param metric_name: metric name
          :return: bool
          """
          metric_name = self.sanitize_metric_name(metric_name)
          return self._dataclass.validation_splits[-1].metric_results[metric_name].higher_score_is_better
  
  
  class TrialSplit(BaseAggregate[Trial, 'MetricResults', TrialSplitDataclass]):
      """
      One split of a trial. This is where a model is trained and evaluated on a specific dataset split.
      """
  
      def __init__(self, _dataclass: SubDataclassT, context: AutoMLContext, is_deep=False, parent: Trial = None):
          if parent is not None:
              _dataclass.hyperparams = parent.get_hyperparams()
          super().__init__(_dataclass, context, is_deep=is_deep, parent=parent)
          self.epoch: int = 0
          self.n_epochs: int = None
  
      def _invariant(self):
          self._assert(self.is_deep,
                       f"self.is_deep should always be set to True for "
                       f"{self.__class__.__name__}", self.context)
          super()._invariant()
  
      @property
      def main_metric_name(self) -> str:
          return self.parent.main_metric_name
  
      def get_hyperparams(self) -> RecursiveDict:
          return self._dataclass.hyperparams
  
      def with_n_epochs(self, n_epochs: int) -> 'TrialSplit':
          self.n_epochs: int = n_epochs
          return self
  
      def next_epoch(self) -> int:
          """
          Increment epoch. Returns the new epoch id.
          Epochs are 1-indexed like lengths: first epoch is 1, second is 2, etc.
          """
          if self.n_epochs is None:
              raise ValueError(
                  "self.n_epochs is not set. Please call self.with_n_epochs(n_epochs) first on your TrialSplit.")
          if self.epoch == 0:
              self.start()
          self.epoch: int = self.epoch + 1
          self.flow.log_epoch(self.epoch, self.n_epochs)
          return self.epoch
  
      def start(self) -> 'TrialSplit':
          """
          Start the trial split.
          """
          self.flow.log_start()
          with self.context.lock:
              self._dataclass.start()
              self.save(False)
          return self
  
      @property
      def metric_results(self) -> Dict[str, 'MetricResults']:
          return {
              metric_name: MetricResults(mr, self.context)  # .at_epoch(self.epoch, self.n_epochs)
              for metric_name, mr in self._dataclass.metric_results.items()
          }
  
      def metric_result(self, metric_name: str = None) -> 'MetricResults':
          """
          Get a metric result that is not managed with a "with" statement. Access it read-only.
          """
          metric_name = self.sanitize_metric_name(metric_name)
          mr: MetricResultsDataclass = self._dataclass.get_sublocation()[metric_name]
          return MetricResults(mr, self.context, is_deep=True)
  
      @_with_method_as_context_manager
      def managed_metric(self, metric_name: str, higher_score_is_better: bool) -> 'TrialSplit':
          """
          To be used as a with statement to get the managed metric.
          """
          self._managed_subresource(metric_name=metric_name, higher_score_is_better=higher_score_is_better)
          return self
  
      def _acquire_managed_subresource(self, metric_name: str, higher_score_is_better: bool) -> 'MetricResults':
          with self.context.lock:
              self.refresh(True)
  
              subdataclass: MetricResultsDataclass = self._create_or_get_metric_results(
                  metric_name, higher_score_is_better)
  
              subagg: MetricResults = self.subaggregate(subdataclass, self.context, is_deep=True, parent=self)
              return subagg
  
      def _create_or_get_metric_results(self, name, higher_score_is_better):
          name = self.sanitize_metric_name(name)
          if name not in self._dataclass.metric_results:
              self._dataclass.metric_results[name] = MetricResultsDataclass(
                  metric_name=name,
                  validation_values=[],
                  train_values=[],
                  higher_score_is_better=higher_score_is_better,
              )
          return self._dataclass.metric_results[name]
  
      def _release_managed_subresource(self, resource: 'MetricResults', e: Exception = None) -> bool:
          handled_error = False
          with self.context.lock:
              if e is not None:
                  handled_error = False
                  self.context.flow.log_error(e)
              else:
                  handled_error = True
              self.save_subaggregate(resource, deep=True)
          return handled_error
  
      def train_context(self) -> 'AutoMLContext':
          return self.context.train()
  
      def validation_context(self) -> 'AutoMLContext':
          return self.context.validation()
  
      def get_metric_names(self) -> List[str]:
          return list(self._dataclass.metric_results.keys())
  
      def _set_success(self) -> 'TrialSplit':
          """
          Set trial status to success. Must save after to ensure consistency.
  
          :return: self
          """
          self._dataclass.end(status=TrialStatus.SUCCESS)
          self.flow.log_success()
          return self
  
      def is_success(self):
          """
          Set trial status to success.
          """
          return self._dataclass.status == TrialStatus.SUCCESS
  
      def _set_failed(self, error: Exception) -> 'TrialSplit':
          """
          Set failed trial with exception. Must save after to ensure consistency.
  
          :param error: catched exception
          :return: self
          """
          if isinstance(error, SystemExit) or isinstance(error, KeyboardInterrupt):
              self._dataclass.end(TrialStatus.ABORTED)
              self.flow.log_aborted(error)
          else:
              self._dataclass.end(TrialStatus.FAILED)
              self.flow.log_failure(error)
          return self
  
  
  class MetricResults(BaseAggregate[TrialSplit, None, MetricResultsDataclass]):
  
      def _invariant(self):
          self._assert(self.is_deep,
                       f"self.is_deep should always be set to True for "
                       f"{self.__class__.__name__}", self.context)
          super()._invariant()
  
      @property
      def metric_name(self) -> str:
          return self._dataclass.metric_name
  
      def _acquire_managed_subresource(self, *args, **kwds) -> SubAggregateT:
          raise NotImplementedError("MetricResults has no subresource to manage as a terminal resource.")
  
      def add_train_result(self, score: float):
          """
          Add a train metric result.
  
          :param name: name of the metric. If None, use the main metric name.
          :param score: the value to be logged
          :param higher_score_is_better: wheter or not a higher score is better for this metric
          """
          with self.context.lock:
              self._dataclass.train_values.append(score)
              self.flow.log_train_metric(self.metric_name, score)
              self.save(True)
  
      def add_valid_result(self, score: float):
          """
          Add a validation metric result.
  
          :param name: name of the metric. If None, use the main metric name.
          :param score: the value to be logged
          :param higher_score_is_better: wheter or not a higher score is better for this metric
          """
          with self.context.lock:
              self._dataclass.validation_values.append(score)
              self.flow.log_valid_metric(self.metric_name, score)
              self.save(True)
  
      def get_train_scores(self) -> List[float]:
          return self._dataclass.train_values
  
      def get_valid_scores(self) -> List[float]:
          """
          Return the validation scores for the given scoring metric.
          """
          return self._dataclass.validation_values
  
      def get_final_validation_score(self) -> float:
          """
          Return the latest validation score for the given scoring metric.
          """
          return self.get_valid_scores()[-1]
  
      def get_best_validation_score(self) -> float:
          """
          Return the best validation score for the given scoring metric.
          """
          scores = self.get_valid_scores()
          if len(scores) == 0:
              return None
  
          if self.is_higher_score_better():
              f = np.max
          else:
              f = np.min
  
          return f(scores)
  
      def get_n_epochs_to_best_validation_score(self) -> int:
          """
          Return the number of epochs
          """
          if self.is_higher_score_better():
              f = np.argmax
          else:
              f = np.argmin
  
          return f(self.get_valid_scores())
  
      def is_higher_score_better(self) -> bool:
          """
          Return True if higher scores are better for the main metric.
  
          :return:
          """
          return self._dataclass.higher_score_is_better
  
      def is_new_best_score(self) -> bool:
          """
          Return True if the latest validation score is the new best score.
  
          :return:
          """
          if self.get_best_validation_score() in self.get_valid_scores()[:-1]:
              return False
          return True
  
      def __iter__(self) -> Iterable[SubAggregateT]:
          """
          Loop over validation values.
          """
          return self._dataclass.validation_values
  
  
  aggregate_2_subaggregate: typing.OrderedDict[Type[BaseAggregate], Type[BaseAggregate]] = OrderedDict([
      (type(None), Root),
      (Root, Project),
      (Project, Client),
      (Client, Round),
      (Round, Trial),
      (Trial, TrialSplit),
      (TrialSplit, MetricResults),
      (MetricResults, type(None)),
  ])
  
  aggregate_2_dataclass: typing.OrderedDict[BaseAggregate, BaseDataclass] = OrderedDict([
      (Root, RootDataclass),
      (Client, ClientDataclass),
      (Round, RoundDataclass),
      (Project, ProjectDataclass),
      (Trial, TrialDataclass),
      (TrialSplit, TrialSplitDataclass),
      (MetricResults, MetricResultsDataclass),
  ])
  dataclass_2_aggregate: typing.OrderedDict[BaseDataclass, BaseAggregate] = {
      dc: agg
      for agg, dc in aggregate_2_dataclass.items()
  }
  