# Copyright (c) OpenMMLab. All rights reserved.
import copy
import logging
import os.path as osp
import tempfile
from unittest import TestCase
from unittest.mock import Mock, patch

import torch
from torch import nn
from torch.nn.init import constant_

from mmengine.logging.logger import MMLogger
from mmengine.model import BaseModule, ModuleDict, ModuleList, Sequential
from mmengine.registry import Registry, build_from_cfg

COMPONENTS = Registry('component')
FOOMODELS = Registry('model')

Logger = MMLogger.get_current_instance()


@COMPONENTS.register_module()
class FooConv1d(BaseModule):

    def __init__(self, init_cfg=None):
        super().__init__(init_cfg)
        self.conv1d = nn.Conv1d(4, 1, 4)

    def forward(self, x):
        return self.conv1d(x)


@COMPONENTS.register_module()
class FooConv2d(BaseModule):

    def __init__(self, init_cfg=None):
        super().__init__(init_cfg)
        self.conv2d = nn.Conv2d(3, 1, 3)

    def forward(self, x):
        return self.conv2d(x)


@COMPONENTS.register_module()
class FooLinear(BaseModule):

    def __init__(self, init_cfg=None):
        super().__init__(init_cfg)
        self.linear = nn.Linear(3, 4)

    def forward(self, x):
        return self.linear(x)


@COMPONENTS.register_module()
class FooLinearConv1d(BaseModule):

    def __init__(self, linear=None, conv1d=None, init_cfg=None):
        super().__init__(init_cfg)
        if linear is not None:
            self.linear = build_from_cfg(linear, COMPONENTS)
        if conv1d is not None:
            self.conv1d = build_from_cfg(conv1d, COMPONENTS)

    def forward(self, x):
        x = self.linear(x)
        return self.conv1d(x)


@FOOMODELS.register_module()
class FooModel(BaseModule):

    def __init__(self,
                 component1=None,
                 component2=None,
                 component3=None,
                 component4=None,
                 init_cfg=None) -> None:
        super().__init__(init_cfg)
        if component1 is not None:
            self.component1 = build_from_cfg(component1, COMPONENTS)
        if component2 is not None:
            self.component2 = build_from_cfg(component2, COMPONENTS)
        if component3 is not None:
            self.component3 = build_from_cfg(component3, COMPONENTS)
        if component4 is not None:
            self.component4 = build_from_cfg(component4, COMPONENTS)

        # its type is not BaseModule, it can be initialized
        # with "override" key.
        self.reg = nn.Linear(3, 4)


class TestBaseModule(TestCase):

    def setUp(self) -> None:
        self.temp_dir = tempfile.TemporaryDirectory()
        self.BaseModule = BaseModule()
        self.model_cfg = dict(
            type='FooModel',
            init_cfg=[
                dict(type='Constant', val=1, bias=2, layer='Linear'),
                dict(type='Constant', val=3, bias=4, layer='Conv1d'),
                dict(type='Constant', val=5, bias=6, layer='Conv2d')
            ],
            component1=dict(type='FooConv1d'),
            component2=dict(type='FooConv2d'),
            component3=dict(type='FooLinear'),
            component4=dict(
                type='FooLinearConv1d',
                linear=dict(type='FooLinear'),
                conv1d=dict(type='FooConv1d')))

        self.model = build_from_cfg(self.model_cfg, FOOMODELS)
        self.logger = MMLogger.get_instance(self._testMethodName)

    def tearDown(self) -> None:
        self.temp_dir.cleanup()
        logging.shutdown()
        MMLogger._instance_dict.clear()
        return super().tearDown()

    def test_is_init(self):
        assert self.BaseModule.is_init is False

    def test_init_weights(self):
        """
        Config
        model (FooModel, Linear: weight=1, bias=2, Conv1d: weight=3, bias=4,
                        Conv2d: weight=5, bias=6)
        ├──component1 (FooConv1d)
        ├──component2 (FooConv2d)
        ├──component3 (FooLinear)
        ├──component4 (FooLinearConv1d)
            ├──linear (FooLinear)
            ├──conv1d (FooConv1d)
        ├──reg (nn.Linear)
        Parameters after initialization
        model (FooModel)
        ├──component1 (FooConv1d, weight=3, bias=4)
        ├──component2 (FooConv2d, weight=5, bias=6)
        ├──component3 (FooLinear, weight=1, bias=2)
        ├──component4 (FooLinearConv1d)
            ├──linear (FooLinear, weight=1, bias=2)
            ├──conv1d (FooConv1d, weight=3, bias=4)
        ├──reg (nn.Linear, weight=1, bias=2)
        """
        self.model.init_weights()

        assert torch.equal(
            self.model.component1.conv1d.weight,
            torch.full(self.model.component1.conv1d.weight.shape, 3.0))
        assert torch.equal(
            self.model.component1.conv1d.bias,
            torch.full(self.model.component1.conv1d.bias.shape, 4.0))
        assert torch.equal(
            self.model.component2.conv2d.weight,
            torch.full(self.model.component2.conv2d.weight.shape, 5.0))
        assert torch.equal(
            self.model.component2.conv2d.bias,
            torch.full(self.model.component2.conv2d.bias.shape, 6.0))
        assert torch.equal(
            self.model.component3.linear.weight,
            torch.full(self.model.component3.linear.weight.shape, 1.0))
        assert torch.equal(
            self.model.component3.linear.bias,
            torch.full(self.model.component3.linear.bias.shape, 2.0))
        assert torch.equal(
            self.model.component4.linear.linear.weight,
            torch.full(self.model.component4.linear.linear.weight.shape, 1.0))
        assert torch.equal(
            self.model.component4.linear.linear.bias,
            torch.full(self.model.component4.linear.linear.bias.shape, 2.0))
        assert torch.equal(
            self.model.component4.conv1d.conv1d.weight,
            torch.full(self.model.component4.conv1d.conv1d.weight.shape, 3.0))
        assert torch.equal(
            self.model.component4.conv1d.conv1d.bias,
            torch.full(self.model.component4.conv1d.conv1d.bias.shape, 4.0))
        assert torch.equal(self.model.reg.weight,
                           torch.full(self.model.reg.weight.shape, 1.0))
        assert torch.equal(self.model.reg.bias,
                           torch.full(self.model.reg.bias.shape, 2.0))

        # Test build model from Pretrained weights

        class CustomLinear(BaseModule):

            def __init__(self, init_cfg=None):
                super().__init__(init_cfg)
                self.linear = nn.Linear(1, 1)

            def init_weights(self):
                constant_(self.linear.weight, 1)
                constant_(self.linear.bias, 2)

        @FOOMODELS.register_module()
        class PratrainedModel(FooModel):

            def __init__(self,
                         component1=None,
                         component2=None,
                         component3=None,
                         component4=None,
                         init_cfg=None) -> None:
                super().__init__(component1, component2, component3,
                                 component4, init_cfg)
                self.linear = CustomLinear()

        checkpoint_path = osp.join(self.temp_dir.name, 'test.pth')
        torch.save(self.model.state_dict(), checkpoint_path)
        model_cfg = copy.deepcopy(self.model_cfg)
        model_cfg['type'] = 'PratrainedModel'
        model_cfg['init_cfg'] = dict(
            type='Pretrained', checkpoint=checkpoint_path)
        model = FOOMODELS.build(model_cfg)
        ori_layer_weight = model.linear.linear.weight.clone()
        ori_layer_bias = model.linear.linear.bias.clone()
        model.init_weights()

        self.assertTrue((ori_layer_weight != model.linear.linear.weight).any())
        self.assertTrue((ori_layer_bias != model.linear.linear.bias).any())

        class FakeDDP(nn.Module):

            def __init__(self, module) -> None:
                super().__init__()
                self.module = module

        # Test initialization of nested modules in DDPModule which define
        # `init_weights`.
        with patch('mmengine.model.base_module.is_model_wrapper',
                   lambda x: isinstance(x, FakeDDP)):
            model = FOOMODELS.build(model_cfg)
            model.ddp = FakeDDP(CustomLinear())
            model.init_weights()
            self.assertTrue((model.ddp.module.linear.weight == 1).all())
            self.assertTrue((model.ddp.module.linear.bias == 2).all())

        # Test submodule.init_weights will be skipped if `is_init` is set
        # to True in root model
        model: FooModel = FOOMODELS.build(copy.deepcopy(self.model_cfg))
        for child in model.children():
            child.init_weights = Mock()
        model.is_init = True
        model.init_weights()
        for child in model.children():
            child.init_weights.assert_not_called()

        # Test submodule.init_weights will be skipped if submodule's `is_init`
        # is set to True
        model: FooModel = FOOMODELS.build(copy.deepcopy(self.model_cfg))
        for child in model.children():
            child.init_weights = Mock()
        model.component1.is_init = True
        model.reg.is_init = True
        model.init_weights()
        model.component1.init_weights.assert_not_called()
        model.component2.init_weights.assert_called_once()
        model.component3.init_weights.assert_called_once()
        model.component4.init_weights.assert_called_once()
        model.reg.init_weights.assert_not_called()

    def test_dump_init_info(self):
        import os
        import shutil
        dump_dir = 'tests/test_model/test_dump_info'
        if not (os.path.exists(dump_dir) and os.path.isdir(dump_dir)):
            os.makedirs(dump_dir)
        for filename in os.listdir(dump_dir):
            file_path = os.path.join(dump_dir, filename)
            if os.path.isfile(file_path) or os.path.islink(file_path):
                os.unlink(file_path)
            elif os.path.isdir(file_path):
                shutil.rmtree(file_path)

        MMLogger.get_instance('logger1')  # add logger without FileHandler
        model1 = build_from_cfg(self.model_cfg, FOOMODELS)
        model1.init_weights()
        assert len(os.listdir(dump_dir)) == 0
        log_path = os.path.join(dump_dir, 'out.log')
        MMLogger.get_instance(
            'logger2', log_file=log_path)  # add logger with FileHandler
        model2 = build_from_cfg(self.model_cfg, FOOMODELS)
        model2.init_weights()
        assert len(os.listdir(dump_dir)) == 1
        assert os.stat(log_path).st_size != 0
        # `FileHandler` should be closed in Windows, otherwise we cannot
        # delete the temporary directory
        logging.shutdown()
        MMLogger._instance_dict.clear()
        shutil.rmtree(dump_dir)


class TestModuleList(TestCase):

    def test_modulelist_weight_init(self):
        models_cfg = [
            dict(
                type='FooConv1d',
                init_cfg=dict(
                    type='Constant', layer='Conv1d', val=0., bias=1.)),
            dict(
                type='FooConv2d',
                init_cfg=dict(
                    type='Constant', layer='Conv2d', val=2., bias=3.)),
        ]
        layers = [build_from_cfg(cfg, COMPONENTS) for cfg in models_cfg]
        modellist = ModuleList(layers)
        modellist.init_weights()
        self.assertTrue(
            torch.equal(modellist[0].conv1d.weight,
                        torch.full(modellist[0].conv1d.weight.shape, 0.)))
        self.assertTrue(
            torch.equal(modellist[0].conv1d.bias,
                        torch.full(modellist[0].conv1d.bias.shape, 1.)))
        self.assertTrue(
            torch.equal(modellist[1].conv2d.weight,
                        torch.full(modellist[1].conv2d.weight.shape, 2.)))
        self.assertTrue(
            torch.equal(modellist[1].conv2d.bias,
                        torch.full(modellist[1].conv2d.bias.shape, 3.)))
        # inner init_cfg has higher priority
        layers = [build_from_cfg(cfg, COMPONENTS) for cfg in models_cfg]
        modellist = ModuleList(
            layers,
            init_cfg=dict(
                type='Constant', layer=['Conv1d', 'Conv2d'], val=4., bias=5.))
        modellist.init_weights()
        self.assertTrue(
            torch.equal(modellist[0].conv1d.weight,
                        torch.full(modellist[0].conv1d.weight.shape, 0.)))
        self.assertTrue(
            torch.equal(modellist[0].conv1d.bias,
                        torch.full(modellist[0].conv1d.bias.shape, 1.)))
        self.assertTrue(
            torch.equal(modellist[1].conv2d.weight,
                        torch.full(modellist[1].conv2d.weight.shape, 2.)))
        self.assertTrue(
            torch.equal(modellist[1].conv2d.bias,
                        torch.full(modellist[1].conv2d.bias.shape, 3.)))


class TestModuleDict(TestCase):

    def test_moduledict_weight_init(self):
        models_cfg = dict(
            foo_conv_1d=dict(
                type='FooConv1d',
                init_cfg=dict(
                    type='Constant', layer='Conv1d', val=0., bias=1.)),
            foo_conv_2d=dict(
                type='FooConv2d',
                init_cfg=dict(
                    type='Constant', layer='Conv2d', val=2., bias=3.)),
        )
        layers = {
            name: build_from_cfg(cfg, COMPONENTS)
            for name, cfg in models_cfg.items()
        }
        modeldict = ModuleDict(layers)
        modeldict.init_weights()
        self.assertTrue(
            torch.equal(
                modeldict['foo_conv_1d'].conv1d.weight,
                torch.full(modeldict['foo_conv_1d'].conv1d.weight.shape, 0.)))
        self.assertTrue(
            torch.equal(
                modeldict['foo_conv_1d'].conv1d.bias,
                torch.full(modeldict['foo_conv_1d'].conv1d.bias.shape, 1.)))
        self.assertTrue(
            torch.equal(
                modeldict['foo_conv_2d'].conv2d.weight,
                torch.full(modeldict['foo_conv_2d'].conv2d.weight.shape, 2.)))
        self.assertTrue(
            torch.equal(
                modeldict['foo_conv_2d'].conv2d.bias,
                torch.full(modeldict['foo_conv_2d'].conv2d.bias.shape, 3.)))
        # inner init_cfg has higher priority
        layers = {
            name: build_from_cfg(cfg, COMPONENTS)
            for name, cfg in models_cfg.items()
        }
        modeldict = ModuleDict(
            layers,
            init_cfg=dict(
                type='Constant', layer=['Conv1d', 'Conv2d'], val=4., bias=5.))
        modeldict.init_weights()
        self.assertTrue(
            torch.equal(
                modeldict['foo_conv_1d'].conv1d.weight,
                torch.full(modeldict['foo_conv_1d'].conv1d.weight.shape, 0.)))
        self.assertTrue(
            torch.equal(
                modeldict['foo_conv_1d'].conv1d.bias,
                torch.full(modeldict['foo_conv_1d'].conv1d.bias.shape, 1.)))
        self.assertTrue(
            torch.equal(
                modeldict['foo_conv_2d'].conv2d.weight,
                torch.full(modeldict['foo_conv_2d'].conv2d.weight.shape, 2.)))
        self.assertTrue(
            torch.equal(
                modeldict['foo_conv_2d'].conv2d.bias,
                torch.full(modeldict['foo_conv_2d'].conv2d.bias.shape, 3.)))


class TestSequential(TestCase):

    def test_sequential_model_weight_init(self):
        seq_model_cfg = [
            dict(
                type='FooConv1d',
                init_cfg=dict(
                    type='Constant', layer='Conv1d', val=0., bias=1.)),
            dict(
                type='FooConv2d',
                init_cfg=dict(
                    type='Constant', layer='Conv2d', val=2., bias=3.)),
        ]
        layers = [build_from_cfg(cfg, COMPONENTS) for cfg in seq_model_cfg]
        seq_model = Sequential(*layers)
        seq_model.init_weights()
        self.assertTrue(
            torch.equal(seq_model[0].conv1d.weight,
                        torch.full(seq_model[0].conv1d.weight.shape, 0.)))
        self.assertTrue(
            torch.equal(seq_model[0].conv1d.bias,
                        torch.full(seq_model[0].conv1d.bias.shape, 1.)))
        self.assertTrue(
            torch.equal(seq_model[1].conv2d.weight,
                        torch.full(seq_model[1].conv2d.weight.shape, 2.)))
        self.assertTrue(
            torch.equal(seq_model[1].conv2d.bias,
                        torch.full(seq_model[1].conv2d.bias.shape, 3.)))
        # inner init_cfg has higher priority
        layers = [build_from_cfg(cfg, COMPONENTS) for cfg in seq_model_cfg]
        seq_model = Sequential(
            *layers,
            init_cfg=dict(
                type='Constant', layer=['Conv1d', 'Conv2d'], val=4., bias=5.))
        seq_model.init_weights()
        self.assertTrue(
            torch.equal(seq_model[0].conv1d.weight,
                        torch.full(seq_model[0].conv1d.weight.shape, 0.)))
        self.assertTrue(
            torch.equal(seq_model[0].conv1d.bias,
                        torch.full(seq_model[0].conv1d.bias.shape, 1.)))
        self.assertTrue(
            torch.equal(seq_model[1].conv2d.weight,
                        torch.full(seq_model[1].conv2d.weight.shape, 2.)))
        self.assertTrue(
            torch.equal(seq_model[1].conv2d.bias,
                        torch.full(seq_model[1].conv2d.bias.shape, 3.)))