import inspect
import detectron2.utils.comm as comm
from detectron2.engine import EvalHook as _EvalHook
from detectron2.evaluation.testing import flatten_results_dict


class EvalHook(_EvalHook):
    def __init__(self, eval_period, eval_function):
        super().__init__(eval_period, eval_function)
        func_args = inspect.getfullargspec(eval_function).args
        assert {"final_iter", "next_iter"}.issubset(set(func_args)), (
            f"Eval function must have either 'final_iter' or 'next_iter' as an argument."
            f"Got {func_args} instead."
        )

    def _do_eval(self, final_iter=False, next_iter=0):
        results = self._func(final_iter=final_iter, next_iter=next_iter)

        if results:
            assert isinstance(
                results, dict
            ), "Eval function must return a dict. Got {} instead.".format(results)

            flattened_results = flatten_results_dict(results)
            for k, v in flattened_results.items():
                try:
                    v = float(v)
                except Exception as e:
                    raise ValueError(
                        "[EvalHook] eval_function should return a nested dict of float. "
                        "Got '{}: {}' instead.".format(k, v)
                    ) from e
            self.trainer.storage.put_scalars(**flattened_results, smoothing_hint=False)

        # Evaluation may take different time among workers.
        # A barrier make them start the next iteration together.
        comm.synchronize()

    def after_step(self):
        next_iter = self.trainer.iter + 1
        if self._period > 0 and next_iter % self._period == 0:
            # do the last eval in after_train
            if next_iter != self.trainer.max_iter:
                self._do_eval(next_iter=next_iter)

    def after_train(self):
        # This condition is to prevent the eval from running after a failed training
        if self.trainer.iter + 1 >= self.trainer.max_iter:
            self._do_eval(final_iter=True)
        # func is likely a closure that holds reference to the trainer
        # therefore we clean it to avoid circular reference in the end
        del self._func