from __future__ import absolute_import
import pytest
import numpy as np
import torch
import mindspore
from mindspore import nn
from mindspore import Tensor, context
from mindspore.nn.optim_ex import AdamW
from mindspore.nn.lr_scheduler import StepLR


class Network(nn.Cell):
    def __init__(self, conv_weight, conv_bias, lin_weight, lin_bias):
        super().__init__()
        self.conv = nn.Conv1d(4, 2, 1, stride=2, weight_init=conv_weight, has_bias=True, bias_init=conv_bias)
        self.lin = nn.Dense(2, 3, weight_init=lin_weight, bias_init=lin_bias)
        self.relu = nn.ReLU()

    def construct(self, x):
        out = self.conv(x)
        out = self.lin(out)
        out = self.relu(out)
        return out


class NetworkPt(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.conv = torch.nn.Conv1d(4, 2, 1, stride=2, bias=False)
        self.lin = torch.nn.Linear(2, 3)
        self.relu = torch.nn.ReLU()

    def forward(self, x):
        out = self.conv(x)
        out = self.lin(out)
        out = self.relu(out)
        return out


class AdamWFactory():
    def __init__(self, group=True, lr_dynamic=False, if_change=False, dtype=np.float32):
        super().__init__()
        self.conv_weight_np = np.random.randn(2, 4, 1).astype(dtype)
        self.conv_bias_np = np.random.randn(2,).astype(dtype)
        self.lin_weight_np = np.random.randn(3, 2).astype(dtype)
        self.lin_bias_np = np.random.randn(3,).astype(dtype)

        self.data = np.random.rand(1, 4, 4).astype(np.float32)
        self.label = np.random.rand(1, 2, 3).astype(np.float32)

        self.group = group
        self.lr_dynamic = lr_dynamic
        self.if_change = if_change
        self.epochs = 4
        self.steps = 1
        self.lr = 0.002
        self.betas = (0.9, 0.999)
        self.eps = 1e-8
        self.weight_decay = 0
        self.amsgrad = False
        self.maximize = False

    def forward_pytorch_impl(self):
        conv_weight = torch.Tensor(self.conv_weight_np.copy())
        conv_bias = torch.Tensor(self.conv_bias_np.copy())
        lin_weight = torch.Tensor(self.lin_weight_np.copy())
        lin_bias = torch.Tensor(self.lin_bias_np.copy())

        model = NetworkPt()
        model.conv.weight = torch.nn.Parameter(conv_weight)
        model.conv.bias = torch.nn.Parameter(conv_bias)
        model.lin.weight = torch.nn.Parameter(lin_weight)
        model.lin.bias = torch.nn.Parameter(lin_bias)

        data = torch.from_numpy(self.data.copy())
        label = torch.from_numpy(self.label.copy())

        if not self.group:
            optimizer = torch.optim.AdamW(model.parameters(), self.lr)
        else:
            conv_params, no_conv_params = [], []
            for param in model.named_parameters():
                if "conv" in param[0]:
                    conv_params.append(param[1])
                else:
                    no_conv_params.append(param[1])
            group_params = [{'params': conv_params, 'weight_decay': 0.01, 'lr': 0.9, "betas": (0.88, 0.8)},
                            {'params': no_conv_params, 'lr': 0.66, "amsgrad": True}]
            optimizer = torch.optim.AdamW(params=group_params, lr=self.lr)

        criterion = torch.nn.L1Loss(reduction='mean')
        lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, 2, gamma=0.5, last_epoch=-1)

        for _ in range(self.epochs):
            for _ in range(self.steps):
                optimizer.zero_grad()
                loss = criterion(model(data), label)
                loss.backward()
                optimizer.step()
            if self.lr_dynamic:
                lr_scheduler.step()
            if self.if_change:
                optimizer.param_groups[1]["betas"] = (0.77, 0.7)
                optimizer.param_groups[1]["amsgrad"] = False

        output = model(data)
        return output.detach().numpy()


    def forward_mindspore_impl(self):
        conv_weight = Tensor(self.conv_weight_np.copy())
        conv_bias = Tensor(self.conv_bias_np.copy())
        lin_weight = Tensor(self.lin_weight_np.copy())
        lin_bias = Tensor(self.lin_bias_np.copy())
        model = Network(conv_weight, conv_bias, lin_weight, lin_bias)

        data = Tensor(self.data)
        label = Tensor(self.label)

        if not self.group:
            optimizer = AdamW(model.trainable_params(), self.lr)
        else:
            conv_params = list(filter(lambda x: 'conv' in x.name, model.trainable_params()))
            no_conv_params = list(filter(lambda x: 'conv' not in x.name, model.trainable_params()))
            group_params = [{'params': conv_params, 'weight_decay': 0.01, 'lr': 0.9, "betas": (0.88, 0.8)},
                            {'params': no_conv_params, 'lr': 0.66, "amsgrad": True}]
            optimizer = AdamW(params=group_params, lr=self.lr)

        criterion = nn.MAELoss(reduction="mean")
        lr_scheduler = StepLR(optimizer, 2, gamma=0.5, last_epoch=-1)

        def forward_fn(data, label):
            logits = model(data)
            loss = criterion(logits, label)
            return loss, logits

        grad_fn = mindspore.value_and_grad(forward_fn, None, optimizer.parameters, has_aux=True)

        def train_step(data, label):
            (loss, _), grads = grad_fn(data, label)
            optimizer(grads)
            return loss

        def train(epochs, steps, lr_dynamic, if_change):
            for _ in range(epochs):
                for _ in range(steps):
                    train_step(data, label)
                if lr_dynamic:
                    lr_scheduler.step()
                if if_change:
                    optimizer.param_groups[1]["betas"] = (0.77, 0.7)
                    optimizer.param_groups[1]["amsgrad"] = False

        train(self.epochs, self.steps, self.lr_dynamic, self.if_change)
        output = model(data)
        return output.asnumpy()

    def result_cmp(self):
        loss_expect = self.forward_pytorch_impl()
        loss_out = self.forward_mindspore_impl()
        allclose_nparray(loss_expect, loss_out, 0.005, 0.005)


def _count_unequal_element(data_expected, data_me, rtol, atol):
    assert data_expected.shape == data_me.shape
    total_count = len(data_expected.flatten())
    error = np.abs(data_expected - data_me)
    greater = np.greater(error, atol + np.abs(data_me) * rtol)
    loss_count = np.count_nonzero(greater)
    assert (loss_count / total_count) < rtol, \
        "\ndata_expected_std:{0}\ndata_me_error:{1}\nloss:{2}". \
            format(data_expected[greater], data_me[greater], error[greater])


def allclose_nparray(data_expected, data_me, rtol, atol, equal_nan=True):
    if np.any(np.isnan(data_expected)) or np.any(np.isnan(data_me)):
        assert np.allclose(data_expected, data_me, rtol, atol, equal_nan=equal_nan)
    elif not np.allclose(data_expected, data_me, rtol, atol, equal_nan=equal_nan):
        _count_unequal_element(data_expected, data_me, rtol, atol)
    else:
        assert np.array(data_expected).shape == np.array(data_me).shape


@pytest.mark.level1
@pytest.mark.platform_x86_gpu_training
@pytest.mark.platform_arm_ascend_training
@pytest.mark.platform_x86_ascend_training
@pytest.mark.env_onecard
@pytest.mark.parametrize('mode', [context.GRAPH_MODE, context.PYNATIVE_MODE])
def test_adamw_basic(mode):
    """
    Feature: Test adamw.
    Description: Test adamw with default parameter.
    Expectation: success.
    """
    fact = AdamWFactory(False, False)
    fact.result_cmp()


@pytest.mark.level1
@pytest.mark.platform_x86_gpu_training
@pytest.mark.platform_arm_ascend_training
@pytest.mark.platform_x86_ascend_training
@pytest.mark.env_onecard
@pytest.mark.parametrize('mode', [context.GRAPH_MODE, context.PYNATIVE_MODE])
def test_adamw_group(mode):
    """
    Feature: Test adamw.
    Description: Test adamw with grouped params.
    Expectation: success.
    """
    fact = AdamWFactory(True, False)
    fact.result_cmp()


@pytest.mark.level1
@pytest.mark.platform_x86_gpu_training
@pytest.mark.platform_arm_ascend_training
@pytest.mark.platform_x86_ascend_training
@pytest.mark.env_onecard
@pytest.mark.parametrize('mode', [context.GRAPH_MODE, context.PYNATIVE_MODE])
def test_adamw_lr_dynamic(mode):
    """
    Feature: Test adamw.
    Description: Test adamw when lr is dynamic.
    Expectation: success.
    """
    fact = AdamWFactory(False, True)
    fact.result_cmp()


@pytest.mark.level1
@pytest.mark.platform_x86_gpu_training
@pytest.mark.platform_arm_ascend_training
@pytest.mark.platform_x86_ascend_training
@pytest.mark.env_onecard
@pytest.mark.parametrize('mode', [context.GRAPH_MODE, context.PYNATIVE_MODE])
def test_adamw_group_lr_dynamic(mode):
    """
    Feature: Test adamw.
    Description: Test adamw with grouped params when lr is dynamic.
    Expectation: success.
    """
    fact = AdamWFactory(True, True)
    fact.result_cmp()


@pytest.mark.level1
@pytest.mark.platform_x86_gpu_training
@pytest.mark.platform_arm_ascend_training
@pytest.mark.platform_x86_ascend_training
@pytest.mark.env_onecard
@pytest.mark.parametrize('mode', [context.PYNATIVE_MODE])
def test_adamw_group_lr_dynamic_change_param(mode):
    """
    Feature: Test adamw.
    Description: Test adamw with grouped params when optimizer params are changed.
    Expectation: success.
    """
    fact = AdamWFactory(True, True, True)
    fact.result_cmp()
