#   Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
#
# 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 os
import sys
import time
import unittest
import numpy as np
import tempfile

EPOCH_NUM = 1
BATCH_SIZE = 1024


def train_func_base(epoch_id, train_loader, model, cost, optimizer):

    total_step = len(train_loader)
    epoch_start = time.time()
    for batch_id, (images, labels) in enumerate(train_loader()):
        # forward
        outputs = model(images)
        loss = cost(outputs, labels)
        # backward and optimize
        loss.backward()
        optimizer.step()
        optimizer.clear_grad()
        print("Epoch [{}/{}], Step [{}/{}], Loss: {}".format(
            epoch_id + 1, EPOCH_NUM, batch_id + 1, total_step, loss.numpy()))
    epoch_end = time.time()
    print(
        f"Epoch ID: {epoch_id+1}, FP32 train epoch time: {(epoch_end - epoch_start) * 1000} ms"
    )


def train_func_ampo1(epoch_id, train_loader, model, cost, optimizer, scaler):
    import paddle

    total_step = len(train_loader)
    epoch_start = time.time()
    for batch_id, (images, labels) in enumerate(train_loader()):
        # forward
        with paddle.amp.auto_cast(
                custom_black_list={"flatten_contiguous_range", "greater_than"},
                level='O1'):
            outputs = model(images)
            loss = cost(outputs, labels)
        # backward and optimize
        scaled = scaler.scale(loss)
        scaled.backward()
        scaler.minimize(optimizer, scaled)
        optimizer.clear_grad()
        print("Epoch [{}/{}], Step [{}/{}], Loss: {}".format(
            epoch_id + 1, EPOCH_NUM, batch_id + 1, total_step, loss.numpy()))
    epoch_end = time.time()
    print(
        f"Epoch ID: {epoch_id+1}, AMPO1 train epoch time: {(epoch_end - epoch_start) * 1000} ms"
    )


def test_func(epoch_id, test_loader, model, cost):
    import paddle

    # evaluation every epoch finish
    model.eval()
    avg_acc = [[], []]
    for batch_id, (images, labels) in enumerate(test_loader()):
        # forward
        outputs = model(images)
        loss = cost(outputs, labels)
        # accuracy
        acc_top1 = paddle.metric.accuracy(input=outputs, label=labels, k=1)
        acc_top5 = paddle.metric.accuracy(input=outputs, label=labels, k=5)
        avg_acc[0].append(acc_top1.numpy())
        avg_acc[1].append(acc_top5.numpy())
    model.train()
    print(
        f"Epoch ID: {epoch_id+1}, Top1 accurary: {np.array(avg_acc[0]).mean()}, Top5 accurary: {np.array(avg_acc[1]).mean()}"
    )


class TestCustomCPUPlugin(unittest.TestCase):

    def setUp(self):
        # compile so and set to current path
        cur_dir = os.path.dirname(os.path.abspath(__file__))
        self.temp_dir = tempfile.TemporaryDirectory()
        cmd = 'cd {} \
            && git clone {} \
            && cd PaddleCustomDevice \
            && git fetch origin \
            && git checkout {} -b dev \
            && cd backends/custom_cpu \
            && mkdir build && cd build && cmake .. && make -j8'.format(
            self.temp_dir.name, os.getenv('PLUGIN_URL'),
            os.getenv('PLUGIN_TAG'))
        os.system(cmd)

        # set environment for loading and registering compiled custom kernels
        # only valid in current process
        os.environ['CUSTOM_DEVICE_ROOT'] = os.path.join(
            cur_dir, '{}/PaddleCustomDevice/backends/custom_cpu/build'.format(
                self.temp_dir.name))

    def tearDown(self):
        self.temp_dir.cleanup()

    def test_custom_cpu_plugin(self):
        self._test_to_static()
        self._test_amp_o1()

    def _test_to_static(self):
        import paddle

        class LeNet5(paddle.nn.Layer):

            def __init__(self):
                super(LeNet5, self).__init__()
                self.fc = paddle.nn.Linear(in_features=1024, out_features=10)
                self.relu = paddle.nn.ReLU()
                self.fc1 = paddle.nn.Linear(in_features=10, out_features=10)

            def forward(self, x):
                out = paddle.flatten(x, 1)
                out = self.fc(out)
                out = self.relu(out)
                out = self.fc1(out)
                return out

        # set device
        paddle.set_device('custom_cpu')

        # model
        model = LeNet5()

        # cost and optimizer
        cost = paddle.nn.CrossEntropyLoss()
        optimizer = paddle.optimizer.Adam(learning_rate=0.001,
                                          parameters=model.parameters())

        # convert to static model
        build_strategy = paddle.static.BuildStrategy()
        mnist = paddle.jit.to_static(model, build_strategy=build_strategy)

        # data loader
        transform = paddle.vision.transforms.Compose([
            paddle.vision.transforms.Resize((32, 32)),
            paddle.vision.transforms.ToTensor(),
            paddle.vision.transforms.Normalize(mean=(0.1307, ), std=(0.3081, ))
        ])
        train_dataset = paddle.vision.datasets.MNIST(mode='train',
                                                     transform=transform,
                                                     download=True)
        test_dataset = paddle.vision.datasets.MNIST(mode='test',
                                                    transform=transform,
                                                    download=True)
        train_loader = paddle.io.DataLoader(train_dataset,
                                            batch_size=BATCH_SIZE,
                                            shuffle=True,
                                            drop_last=True,
                                            num_workers=2)
        test_loader = paddle.io.DataLoader(test_dataset,
                                           batch_size=BATCH_SIZE,
                                           shuffle=True,
                                           drop_last=True,
                                           num_workers=2)

        # train and eval
        for epoch_id in range(EPOCH_NUM):
            train_func_base(epoch_id, train_loader, model, cost, optimizer)
            test_func(epoch_id, test_loader, model, cost)

    def _test_amp_o1(self):
        import paddle

        class LeNet5(paddle.nn.Layer):

            def __init__(self):
                super(LeNet5, self).__init__()
                self.fc = paddle.nn.Linear(in_features=1024, out_features=10)
                self.relu = paddle.nn.ReLU()
                self.fc1 = paddle.nn.Linear(in_features=10, out_features=10)

            def forward(self, x):
                out = paddle.flatten(x, 1)
                out = self.fc(out)
                out = self.relu(out)
                out = self.fc1(out)
                return out

        # set device
        paddle.set_device('custom_cpu')

        # model
        model = LeNet5()

        # cost and optimizer
        cost = paddle.nn.CrossEntropyLoss()
        optimizer = paddle.optimizer.Adam(learning_rate=0.001,
                                          parameters=model.parameters())

        # convert to static model
        scaler = paddle.amp.GradScaler(init_loss_scaling=1024)
        model, optimizer = paddle.amp.decorate(models=model,
                                               optimizers=optimizer,
                                               level='O1')

        # data loader
        transform = paddle.vision.transforms.Compose([
            paddle.vision.transforms.Resize((32, 32)),
            paddle.vision.transforms.ToTensor(),
            paddle.vision.transforms.Normalize(mean=(0.1307, ), std=(0.3081, ))
        ])
        train_dataset = paddle.vision.datasets.MNIST(mode='train',
                                                     transform=transform,
                                                     download=True)
        test_dataset = paddle.vision.datasets.MNIST(mode='test',
                                                    transform=transform,
                                                    download=True)
        train_loader = paddle.io.DataLoader(train_dataset,
                                            batch_size=BATCH_SIZE,
                                            shuffle=True,
                                            drop_last=True,
                                            num_workers=2)
        test_loader = paddle.io.DataLoader(test_dataset,
                                           batch_size=BATCH_SIZE,
                                           shuffle=True,
                                           drop_last=True,
                                           num_workers=2)

        # train and eval
        for epoch_id in range(EPOCH_NUM):
            train_func_ampo1(epoch_id, train_loader, model, cost, optimizer,
                             scaler)
            test_func(epoch_id, test_loader, model, cost)


if __name__ == '__main__':
    if os.name == 'nt' or sys.platform.startswith('darwin'):
        # only support Linux now
        exit()
    unittest.main()
