import sys
sys.path.append("../../")
import logging
import numpy as np
import unittest
import paddle
from paddleslim.dygraph.dist import Distill, config2yaml
from paddleslim.common.log_helper import get_logger

_logger = get_logger(
    __name__, logging.INFO, fmt='%(asctime)s-%(levelname)s: %(message)s')

# class TestImperativeDistill(unittest.TestCase):
#     def setUp(self):
#         self.s_model, self.t_model = self.prepare_model()
#         self.t_model.eval()
#         self.distill_configs = self.prepare_config()

#     def prepare_model(self):
#         return paddle.vision.models.MobileNetV1(
#         ), paddle.vision.models.MobileNetV1()

#     def prepare_config(self):
#         self.convert_fn = False
#         distill_configs = [{
#             'loss_function': 'MSELoss',
#             'layers': [
#                 {
#                     "layers_name": ["conv1", "conv1"]
#                 },
#                 {
#                     "layers_name": ["conv2_2", "conv2_2"]
#                 },
#             ]
#         }, {
#             'loss_function': 'CELoss',
#             'temperature': 1.0,
#             'layers': [{
#                 "layers_name": ["fc", "fc"]
#             }, ]
#         }]
#         return distill_configs

#     def test_distill(self):
#         transform = paddle.vision.transforms.Compose([
#             paddle.vision.transforms.Transpose(),
#             paddle.vision.transforms.Normalize([127.5], [127.5])
#         ])

#         train_dataset = paddle.vision.datasets.Cifar10(
#             mode='train', backend='cv2', transform=transform)
#         val_dataset = paddle.vision.datasets.Cifar10(
#             mode='test', backend='cv2', transform=transform)

#         place = paddle.CUDAPlace(0) if paddle.is_compiled_with_cuda(
#         ) else paddle.CPUPlace()
#         train_reader = paddle.io.DataLoader(
#             train_dataset, drop_last=True, places=place, batch_size=64)
#         test_reader = paddle.io.DataLoader(
#             val_dataset, places=place, batch_size=64)

#         def test(model):
#             model.eval()
#             avg_acc = [[], []]
#             for batch_id, data in enumerate(test_reader):
#                 img = paddle.to_tensor(data[0])
#                 label = paddle.to_tensor(data[1])
#                 label = paddle.reshape(label, [-1, 1])
#                 out = model(img)
#                 acc_top1 = paddle.metric.accuracy(input=out, label=label, k=1)
#                 acc_top5 = paddle.metric.accuracy(input=out, label=label, k=5)
#                 avg_acc[0].append(acc_top1.numpy())
#                 avg_acc[1].append(acc_top5.numpy())
#                 if batch_id % 100 == 0:
#                     _logger.info(
#                         "Test | step {}: acc1 = {:}, acc5 = {:}".format(
#                             batch_id, acc_top1.numpy(), acc_top5.numpy()))

#             _logger.info("Test |Average: acc_top1 {}, acc_top5 {}".format(
#                 np.mean(avg_acc[0]), np.mean(avg_acc[1])))

#         def train(model):
#             adam = paddle.optimizer.Adam(
#                 learning_rate=0.001, parameters=model.parameters())

#             for batch_id, data in enumerate(train_reader):
#                 img = paddle.to_tensor(data[0])
#                 label = paddle.to_tensor(data[1])
#                 distill_loss, student_out, teacher_out = model(img)
#                 loss = paddle.nn.functional.loss.cross_entropy(student_out,
#                                                                label)
#                 avg_loss = paddle.mean(loss)
#                 all_loss = avg_loss + distill_loss
#                 all_loss.backward()
#                 adam.step()
#                 adam.clear_grad()
#                 if batch_id % 100 == 0:
#                     _logger.info("Train | At epoch {} step {}: loss = {:}".
#                                  format(str(0), batch_id, all_loss.numpy()))
#             test(self.s_model)
#             self.s_model.train()

#         distill_model = Distill(
#             self.distill_configs,
#             self.s_model,
#             self.t_model,
#             convert_fn=self.convert_fn)
#         train(distill_model)


class TestImperativeDistillCase1(unittest.TestCase):
    def setUp(self):
        self.s_model, self.t_model = self.prepare_model()
        self.t_model.eval()
        self.distill_configs = self.prepare_config()

    def test_distill(self):
        transform = paddle.vision.transforms.Compose([
            paddle.vision.transforms.Transpose(),
            paddle.vision.transforms.Normalize([127.5], [127.5])
        ])

        train_dataset = paddle.vision.datasets.Cifar10(
            mode='train', backend='cv2', transform=transform)
        val_dataset = paddle.vision.datasets.Cifar10(
            mode='test', backend='cv2', transform=transform)

        place = paddle.CUDAPlace(0) if paddle.is_compiled_with_cuda(
        ) else paddle.CPUPlace()
        train_reader = paddle.io.DataLoader(
            train_dataset, drop_last=True, places=place, batch_size=64)
        test_reader = paddle.io.DataLoader(
            val_dataset, places=place, batch_size=64)

        def test(model):
            model.eval()
            avg_acc = [[], []]
            for batch_id, data in enumerate(test_reader):
                img = paddle.to_tensor(data[0])
                label = paddle.to_tensor(data[1])
                label = paddle.reshape(label, [-1, 1])
                out = model(img)
                acc_top1 = paddle.metric.accuracy(input=out, label=label, k=1)
                acc_top5 = paddle.metric.accuracy(input=out, label=label, k=5)
                avg_acc[0].append(acc_top1.numpy())
                avg_acc[1].append(acc_top5.numpy())
                if batch_id % 100 == 0:
                    _logger.info(
                        "Test | step {}: acc1 = {:}, acc5 = {:}".format(
                            batch_id, acc_top1.numpy(), acc_top5.numpy()))

            _logger.info("Test |Average: acc_top1 {}, acc_top5 {}".format(
                np.mean(avg_acc[0]), np.mean(avg_acc[1])))

        def train(model):
            adam = paddle.optimizer.Adam(
                learning_rate=0.001, parameters=model.parameters())

            for batch_id, data in enumerate(train_reader):
                img = paddle.to_tensor(data[0])
                label = paddle.to_tensor(data[1])
                distill_loss, student_out, teacher_out = model(img)
                loss = paddle.nn.functional.loss.cross_entropy(student_out,
                                                               label)
                avg_loss = paddle.mean(loss)
                all_loss = avg_loss + distill_loss
                all_loss.backward()
                adam.step()
                adam.clear_grad()
                if batch_id % 100 == 0:
                    _logger.info("Train | At epoch {} step {}: loss = {:}".
                                 format(str(0), batch_id, all_loss.numpy()))
            test(self.s_model)
            self.s_model.train()

        distill_model = Distill(
            self.distill_configs,
            self.s_model,
            self.t_model,
            convert_fn=self.convert_fn)
        train(distill_model)

    def prepare_model(self):
        class convbn(paddle.nn.Layer):
            def __init__(self):
                super(convbn, self).__init__()
                self.conv = paddle.nn.Conv2D(3, 3, 3, padding=1)
                self.bn = paddle.nn.BatchNorm(3)

            def forward(self, x):
                conv_out = self.conv(x)
                bn_out = self.bn(conv_out)
                return tuple([conv_out, bn_out])

        class Model(paddle.nn.Layer):
            def __init__(self):
                super(Model, self).__init__()
                self.conv1 = paddle.nn.Conv2D(3, 3, 3, padding=1)
                self.conv2 = convbn()
                self.conv3 = paddle.nn.Conv2D(3, 3, 3, padding=1)
                self.fc = paddle.nn.Linear(3072, 10)

            def forward(self, x):
                self.conv1_out = self.conv1(x)
                conv2_out = self.conv2(self.conv1_out)
                self.conv3_out = self.conv3(conv2_out[0])
                out = paddle.reshape(self.conv3_out, shape=[x.shape[0], -1])
                out = paddle.nn.functional.softmax(out)
                out = self.fc(out)
                return out

        return Model(), Model()

    def prepare_config(self):
        self.convert_fn = True
        distill_configs = [{
            'loss_function': 'MSELoss',
            'layers': [
                {
                    "layers_name": ["conv1", "conv1"],
                    'align_params': {
                        'align_type': '1x1conv',
                        'in_channel': 3,
                        'out_channel': 3
                    }
                },
                {
                    "layers_name": ["conv2", "conv3"],
                    'io': ["input", "output"],
                    'align_params': {
                        'align_type': '3x3conv',
                        'in_channel': 3,
                        'out_channel': 3
                    }
                },
                {
                    "layers_name": ["conv2", "conv3"],
                    'io': ["output", "output"],
                    'idx': [1, None],
                    'align_params': {
                        'align_type': '1x1conv+bn',
                        'in_channel': 3,
                        'out_channel': 3
                    }
                },
                {
                    "layers_name": ["conv2", "conv3"],
                    'io': ["output", "output"],
                    'idx': [1, None],
                    'align_params': {
                        'align_type': '3x3conv+bn',
                        'in_channel': 3,
                        'out_channel': 3,
                        'transpose_model': 'student'
                    }
                },
            ]
        }, {
            'loss_function': 'CELoss',
            'temperature': 1.0,
            'layers': [{
                "layers_name": ["fc", "fc"],
                'align_params': {
                    'align_type': 'linear',
                    'in_channel': 10,
                    'out_channel': 10,
                    'weight_init': {
                        'initializer': 'Normal',
                        'mean': 0.0,
                        'std': 0.02
                    },
                }
            }, ]
        }]
        config2yaml(distill_configs, 'test.yaml')
        return './test.yaml'


if __name__ == '__main__':
    unittest.main()
