import easydict
import pytest
from ding.worker.learner import register_learner_hook, build_learner_hook_by_cfg, LearnerHook
from ding.worker.learner.learner_hook import SaveCkptHook, LoadCkptHook, LogShowHook, LogReduceHook
from ding.worker.learner.learner_hook import show_hooks, add_learner_hook, merge_hooks
from easydict import EasyDict


@pytest.fixture(scope='function')
def setup_simplified_hook_cfg():
    return dict(
        save_ckpt_after_iter=20,
        save_ckpt_after_run=True,
    )


@pytest.fixture(scope='function')
def fake_setup_simplified_hook_cfg():
    return dict(
        log_show_after_iter=20,
        log_reduce_after_iter=True,
    )


@pytest.mark.unittest
class TestLearnerHook:

    def test_register(self):

        class FakeHook(LearnerHook):
            pass

        register_learner_hook('fake', FakeHook)
        with pytest.raises(AssertionError):
            register_learner_hook('placeholder', type)

    def test_build_learner_hook_by_cfg(self, setup_simplified_hook_cfg):
        hooks = build_learner_hook_by_cfg(setup_simplified_hook_cfg)
        show_hooks(hooks)
        assert len(hooks['before_run']) == 0
        assert len(hooks['before_iter']) == 0
        assert len(hooks['after_iter']) == 1
        assert isinstance(hooks['after_iter'][0], SaveCkptHook)
        assert len(hooks['after_run']) == 1
        assert isinstance(hooks['after_run'][0], SaveCkptHook)

    def test_add_learner_hook(self, setup_simplified_hook_cfg):
        hooks = build_learner_hook_by_cfg(setup_simplified_hook_cfg)
        hook_1 = LogShowHook('log_show', 20, position='after_iter', ext_args=EasyDict({'freq': 100}))
        add_learner_hook(hooks, hook_1)
        hook_2 = LoadCkptHook('load_ckpt', 20, position='before_run', ext_args=EasyDict({'load_path': './model.pth'}))
        add_learner_hook(hooks, hook_2)
        hook_3 = LogReduceHook('log_reduce', 10, position='after_iter')
        add_learner_hook(hooks, hook_3)

        show_hooks(hooks)
        assert len(hooks['after_iter']) == 3
        assert len(hooks['after_run']) == 1
        assert len(hooks['before_run']) == 1
        assert len(hooks['before_iter']) == 0
        assert isinstance(hooks['after_run'][0], SaveCkptHook)
        assert isinstance(hooks['before_run'][0], LoadCkptHook)

    def test_merge_hooks(self, setup_simplified_hook_cfg, fake_setup_simplified_hook_cfg):
        hooks = build_learner_hook_by_cfg(setup_simplified_hook_cfg)
        show_hooks(hooks)
        fake_hooks = build_learner_hook_by_cfg(fake_setup_simplified_hook_cfg)
        show_hooks(fake_hooks)
        hooks_ = merge_hooks(hooks, fake_hooks)
        show_hooks(hooks_)
        assert len(hooks_['after_iter']) == 3
        assert len(hooks_['after_run']) == 1
        assert len(hooks_['before_run']) == 0
        assert len(hooks_['before_iter']) == 0
        assert isinstance(hooks['after_run'][0], SaveCkptHook)