# -*- coding: utf-8 -*-

# (C) Copyright 2020, 2021, 2022 IBM. All Rights Reserved.
#
# This code is licensed under the Apache License, Version 2.0. You may
# obtain a copy of this license in the LICENSE.txt file in the root directory
# of this source tree or at http://www.apache.org/licenses/LICENSE-2.0.
#
# Any modifications or derivative works of this code must retain this
# copyright notice, and modified files need to carry a notice indicating
# that they have been altered from the originals.

# pylint: disable=too-many-locals, too-many-public-methods, no-member
"""Test for different utility functionality."""

from tempfile import TemporaryFile
from copy import deepcopy
from unittest import SkipTest

from numpy import array
from numpy.random import rand
from numpy.testing import assert_array_almost_equal, assert_raises
from torch import Tensor, load, save, device, manual_seed
from torch import abs as torch_abs
from torch.nn import Module, Sequential
from torch.nn import Linear as torch_linear
from torch.nn.functional import mse_loss
from torch.optim import SGD

from aihwkit.nn import (
    AnalogConv2d, AnalogConv2dMapped, AnalogSequential, AnalogLinearMapped
)
from aihwkit.optim import AnalogSGD
from aihwkit.simulator.configs import SingleRPUConfig, InferenceRPUConfig, FloatingPointRPUConfig
from aihwkit.simulator.configs.devices import ConstantStepDevice, LinearStepDevice
from aihwkit.simulator.configs.utils import IOParameters, UpdateParameters, MappingParameter
from aihwkit.exceptions import TileError, ModuleError
from aihwkit.nn.conversion import convert_to_analog

from .helpers.decorators import parametrize_over_layers
from .helpers.layers import (
    Conv2d, Conv2dCuda, Linear, LinearCuda, LinearMapped,
    LinearMappedCuda, Conv2dMapped, Conv2dMappedCuda
)
from .helpers.testcases import ParametrizedTestCase, SKIP_CUDA_TESTS
from .helpers.tiles import FloatingPoint, ConstantStep, Inference, InferenceLearnOutScaling


@parametrize_over_layers(
    layers=[Linear, Conv2d, LinearMapped, LinearCuda, LinearMappedCuda,
            Conv2dCuda, Conv2dMapped, Conv2dMappedCuda],
    tiles=[FloatingPoint, ConstantStep, Inference, InferenceLearnOutScaling],
    biases=['analog', 'digital', None],
)
class SerializationTest(ParametrizedTestCase):
    """Tests for serialization."""

    @staticmethod
    def train_model(model, loss_func, x_b, y_b):
        """Train the model."""
        opt = AnalogSGD(model.parameters(), lr=0.5)
        # opt.regroup_param_groups(model)
        epochs = 3
        for _ in range(epochs):
            opt.zero_grad()
            pred = model(x_b)

            loss = loss_func(pred, y_b)
            loss.backward()
            opt.step()
        return loss

    @staticmethod
    def get_layer_and_tile_weights(model):
        """Return the weights and biases of the model and the tile and whether
        it automatically syncs"""

        if isinstance(model, AnalogLinearMapped):
            weight, bias = model.get_weights()
            return weight, bias, weight, bias, True

        if isinstance(model, AnalogConv2dMapped):
            weight, bias = model.get_weights()
            return weight, bias, weight, bias, True

        weight = model.weight.data.detach().cpu().numpy()
        if model.use_bias:
            bias = model.bias.data.detach().cpu().numpy()
        else:
            bias = None

        analog_weight, analog_bias = model.get_weights()
        analog_weight = analog_weight.detach().cpu().numpy()
        if analog_bias is not None:
            analog_bias = analog_bias.detach().cpu().numpy()

        return weight, bias, analog_weight.reshape(weight.shape), analog_bias, True

    @staticmethod
    def get_analog_tile(model):
        """ Return a (python) analog tile of the model"""
        return list(model.analog_tiles())[0]

    def test_save_load_state_dict_train(self):
        """Test saving and loading using a state dict after training."""
        model = self.get_layer()

        # Perform an update in order to modify tile weights and biases.
        loss_func = mse_loss
        if isinstance(model, (AnalogConv2d, AnalogConv2dMapped)):
            input_x = Tensor(rand(2, 2, 3, 3))*0.2
            input_y = Tensor(rand(2, 3, 5, 5))*0.2
        else:
            input_x = Tensor(rand(2, model.in_features))*0.2
            input_y = Tensor(rand(2, model.out_features))*0.2

        if self.use_cuda:
            input_x = input_x.cuda()
            input_y = input_y.cuda()

        self.train_model(model, loss_func, input_x, input_y)

        # Keep track of the current weights and biases for comparing.
        (model_weights, model_biases,
         tile_weights, tile_biases, sync) = self.get_layer_and_tile_weights(model)

        # now the tile weights should be out of sync
        if not sync:
            assert_raises(AssertionError, assert_array_almost_equal, model_weights, tile_weights)
            if self.bias:
                assert_raises(AssertionError, assert_array_almost_equal, model_biases, tile_biases)

        # Save the model to a file.
        with TemporaryFile() as file:
            save(model.state_dict(), file)
            # Create a new model and load its state dict.
            file.seek(0)
            new_model = self.get_layer()
            new_model.load_state_dict(load(file))

        # Compare the new model weights and biases. they should now be in sync
        (new_model_weights, new_model_biases,
         new_tile_weights, new_tile_biases, _) = self.get_layer_and_tile_weights(new_model)

        assert_array_almost_equal(tile_weights, new_model_weights)
        assert_array_almost_equal(tile_weights, new_tile_weights)
        if self.bias:
            assert_array_almost_equal(tile_biases, new_model_biases)
            assert_array_almost_equal(tile_biases, new_tile_biases)

    def test_save_load_state_dict_train_after(self):
        """Test saving and loading using a state dict and training after load."""
        model = self.get_layer()

        # Perform an update in order to modify tile weights and biases.
        loss_func = mse_loss
        if isinstance(model, (AnalogConv2d, AnalogConv2dMapped)):
            input_x = Tensor(rand(2, 2, 3, 3))*0.2
            input_y = Tensor(rand(2, 3, 5, 5))*0.2
        else:
            input_x = Tensor(rand(2, model.in_features))*0.2
            input_y = Tensor(rand(2, model.out_features))*0.2

        if self.use_cuda:
            input_x = input_x.cuda()
            input_y = input_y.cuda()

        self.train_model(model, loss_func, input_x, input_y)

        # Keep track of the current weights and biases for comparing.
        _, _, tile_weights, tile_biases, _ = self.get_layer_and_tile_weights(model)

        # Save the model to a file.
        with TemporaryFile() as file:
            save(model.state_dict(), file)
            # Create a new model and load its state dict.
            file.seek(0)
            new_model = self.get_layer()
            new_model.load_state_dict(load(file))

        # Compare the new model weights and biases. they should now be in sync
        (new_model_weights, new_model_biases,
         new_tile_weights, new_tile_biases, _) = self.get_layer_and_tile_weights(new_model)

        assert_array_almost_equal(tile_weights, new_model_weights)
        assert_array_almost_equal(tile_weights, new_tile_weights)
        if self.bias:
            assert_array_almost_equal(tile_biases, new_model_biases)
            assert_array_almost_equal(tile_biases, new_tile_biases)

        new_loss = self.train_model(new_model, loss_func, input_x, input_y)

        loss = self.train_model(model, loss_func, input_x, input_y)
        if self.tile_class != ConstantStep:
            self.assertTensorAlmostEqual(loss, new_loss)

    def test_save_load_state_dict_train_after_old_model(self):
        """Test saving and loading using a state dict and training after load with old model."""
        model = self.get_layer()

        # Perform an update in order to modify tile weights and biases.
        loss_func = mse_loss
        if isinstance(model, (AnalogConv2d, AnalogConv2dMapped)):
            input_x = Tensor(rand(2, 2, 3, 3))*0.2
            input_y = Tensor(rand(2, 3, 5, 5))*0.2
        else:
            input_x = Tensor(rand(2, model.in_features))*0.2
            input_y = Tensor(rand(2, model.out_features))*0.2

        if self.use_cuda:
            input_x = input_x.cuda()
            input_y = input_y.cuda()

        self.train_model(model, loss_func, input_x, input_y)

        # Keep track of the current weights and biases for comparing.
        _, _, tile_weights, tile_biases, _ = self.get_layer_and_tile_weights(model)

        # Save the model to a file.
        with TemporaryFile() as file:
            save(model.state_dict(), file)
            # Create a new model and load its state dict.
            file.seek(0)
            model.load_state_dict(load(file))

        # Compare the new model weights and biases. they should now be in sync
        (new_model_weights, new_model_biases,
         new_tile_weights, new_tile_biases, _) = self.get_layer_and_tile_weights(model)

        assert_array_almost_equal(tile_weights, new_model_weights)
        assert_array_almost_equal(tile_weights, new_tile_weights)
        if self.bias:
            assert_array_almost_equal(tile_biases, new_model_biases)
            assert_array_almost_equal(tile_biases, new_tile_biases)

        self.train_model(model, loss_func, input_x, input_y)

    def test_save_load_train_after(self):
        """Test saving and loading using a state dict and training after load."""
        model = self.get_layer()

        # Perform an update in order to modify tile weights and biases.
        loss_func = mse_loss
        if isinstance(model, (AnalogConv2d, AnalogConv2dMapped)):
            input_x = Tensor(rand(2, 2, 3, 3))*0.2
            input_y = Tensor(rand(2, 3, 5, 5))*0.2
        else:
            input_x = Tensor(rand(2, model.in_features))*0.2
            input_y = Tensor(rand(2, model.out_features))*0.2

        if self.use_cuda:
            input_x = input_x.cuda()
            input_y = input_y.cuda()

        # Keep track of the current weights and biases for comparing.
        _, _, tile_weights, tile_biases, _ = self.get_layer_and_tile_weights(model)

        # Save the model to a file.
        with TemporaryFile() as file:
            save(model, file)
            # Create a new model and load its state dict.
            file.seek(0)
            new_model = load(file)

        # Compare the new model weights and biases. they should now be in sync
        (new_model_weights, new_model_biases,
         new_tile_weights, new_tile_biases, _) = self.get_layer_and_tile_weights(new_model)

        assert_array_almost_equal(tile_weights, new_model_weights)
        assert_array_almost_equal(tile_weights, new_tile_weights)
        if self.bias:
            assert_array_almost_equal(tile_biases, new_model_biases)
            assert_array_almost_equal(tile_biases, new_tile_biases)

        new_loss = self.train_model(new_model, loss_func, input_x, input_y)
        loss = self.train_model(model, loss_func, input_x, input_y)
        if self.tile_class != ConstantStep:
            self.assertTensorAlmostEqual(loss, new_loss)

    def test_save_load_model(self):
        """Test saving and loading a model directly."""
        model = self.get_layer()

        # Keep track of the current weights and biases for comparing.
        (model_weights, model_biases,
         tile_weights, tile_biases, _) = self.get_layer_and_tile_weights(model)
        assert_array_almost_equal(model_weights, tile_weights)
        if self.bias:
            assert_array_almost_equal(model_biases, tile_biases)

        # Save the model to a file.
        with TemporaryFile() as file:
            save(model, file)
            # Load the model.
            file.seek(0)
            new_model = load(file)

        # Compare the new model weights and biases.
        (new_model_weights, new_model_biases,
         new_tile_weights, new_tile_biases, _) = self.get_layer_and_tile_weights(new_model)

        assert_array_almost_equal(model_weights, new_model_weights)
        assert_array_almost_equal(tile_weights, new_tile_weights)
        if self.bias:
            assert_array_almost_equal(model_biases, new_model_biases)
            assert_array_almost_equal(tile_biases, new_tile_biases)

        # Asserts over the AnalogContext of the new model.
        new_analog_tile = self.get_analog_tile(new_model)
        analog_tile = self.get_analog_tile(model)
        self.assertTrue(hasattr(new_analog_tile.analog_ctx, 'analog_tile'))
        self.assertIsInstance(new_analog_tile.analog_ctx.analog_tile,
                              analog_tile.__class__)
        self.assertTrue(new_analog_tile.is_cuda == analog_tile.is_cuda)

    def test_save_load_model_cross_device(self):
        """Test saving and loading a model directly."""

        if SKIP_CUDA_TESTS:
            raise SkipTest('CUDA not available.')

        model = self.get_layer()

        map_location = 'cuda'
        if self.use_cuda:
            map_location = 'cpu'

        # Keep track of the current weights and biases for comparing.
        (model_weights, model_biases,
         tile_weights, tile_biases, _) = self.get_layer_and_tile_weights(model)
        assert_array_almost_equal(model_weights, tile_weights)
        if self.bias:
            assert_array_almost_equal(model_biases, tile_biases)

        # Save the model to a file.
        with TemporaryFile() as file:
            save(model, file)
            # Load the model.
            file.seek(0)
            new_model = load(file, map_location=device(map_location))

        # Compare the new model weights and biases.
        (new_model_weights, new_model_biases,
         new_tile_weights, new_tile_biases, _) = self.get_layer_and_tile_weights(new_model)

        assert_array_almost_equal(model_weights, new_model_weights)
        assert_array_almost_equal(tile_weights, new_tile_weights)
        if self.bias:
            assert_array_almost_equal(model_biases, new_model_biases)
            assert_array_almost_equal(tile_biases, new_tile_biases)

        # Asserts over the AnalogContext of the new model.
        new_analog_tile = self.get_analog_tile(new_model)
        analog_tile = self.get_analog_tile(model)

        self.assertTrue(hasattr(new_analog_tile.analog_ctx, 'analog_tile'))
        self.assertIsInstance(new_analog_tile.analog_ctx.analog_tile,
                              analog_tile.__class__)

        self.assertTrue(new_analog_tile.is_cuda != analog_tile.is_cuda)

        if analog_tile.shared_weights is not None:
            self.assertTrue(new_analog_tile.shared_weights.device.type == map_location)

    def test_save_load_meta_parameter(self):
        """Test saving and loading a device with custom parameters."""
        # Create the device and the array.
        rpu_config = SingleRPUConfig(
            forward=IOParameters(inp_noise=0.321),
            backward=IOParameters(inp_noise=0.456),
            update=UpdateParameters(desired_bl=78),
            device=ConstantStepDevice(w_max=0.987)
        )

        model = self.get_layer(rpu_config=rpu_config)

        # Save the model to a file.
        with TemporaryFile() as file:
            save(model, file)
            # Load the model.
            file.seek(0)
            new_model = load(file)

        # Assert over the new model tile parameters.
        new_analog_tile = self.get_analog_tile(new_model)
        analog_tile = self.get_analog_tile(model)

        parameters = new_analog_tile.tile.get_parameters()
        self.assertAlmostEqual(parameters.forward_io.inp_noise, 0.321)
        self.assertAlmostEqual(parameters.backward_io.inp_noise, 0.456)
        self.assertAlmostEqual(parameters.update.desired_bl, 78)
        self.assertTrue(new_analog_tile.is_cuda == analog_tile.is_cuda)

    def test_save_load_hidden_parameters(self):
        """Test saving and loading a device with hidden parameters."""
        # Create the device and the array.
        model = self.get_layer()
        analog_tile = self.get_analog_tile(model)
        hidden_parameters = analog_tile.tile.get_hidden_parameters()

        # Save the model to a file.
        with TemporaryFile() as file:
            save(model, file)
            # Load the model.
            file.seek(0)
            new_model = load(file)

        # Assert over the new model tile parameters.
        new_analog_tile = self.get_analog_tile(new_model)
        new_hidden_parameters = new_analog_tile.tile.get_hidden_parameters()
        assert_array_almost_equal(hidden_parameters, new_hidden_parameters)

    def test_save_load_out_scaling_alpha(self):
        """Test saving and loading a device with out_scaling_alpha."""
        # Create the device and the array.
        model = self.get_layer()
        alpha = 2.0
        analog_tile = self.get_analog_tile(model)
        analog_tile.out_scaling_alpha = Tensor([alpha])

        # Save the model to a file.
        with TemporaryFile() as file:
            save(model, file)
            # Load the model.
            file.seek(0)
            new_model = load(file)

        # Assert over the new model tile parameters.
        new_analog_tile = self.get_analog_tile(new_model)
        alpha_new = new_analog_tile.get_scales().detach().cpu()
        assert_array_almost_equal(array(alpha), array(alpha_new))

    def test_save_load_shared_weights(self):
        """Test saving and loading a device with shared_weights."""

        if isinstance(self.get_rpu_config(), FloatingPointRPUConfig):
            raise SkipTest('Not available for FP')

        # Create the device and the array.
        model = self.get_layer()
        analog_tile = self.get_analog_tile(model)

        shared_weights = None
        if analog_tile.shared_weights is not None:
            shared_weights = analog_tile.shared_weights.detach().cpu().numpy()

        # Save the model to a file.
        with TemporaryFile() as file:
            save(model, file)
            # Load the model.
            file.seek(0)
            new_model = load(file)

        # Assert over the new model tile parameters.
        new_analog_tile = self.get_analog_tile(new_model)
        if shared_weights is not None:
            new_shared_weights = new_analog_tile.shared_weights
            assert_array_almost_equal(shared_weights, new_shared_weights.detach().cpu().numpy())

    def test_save_load_weight_scaling_omega(self):
        """Test saving and loading a device with weight scaling omega."""
        # Create the device and the array.
        rpu_config = SingleRPUConfig(mapping=MappingParameter(weight_scaling_omega=0.4))
        model = self.get_layer(rpu_config=rpu_config)
        analog_tile = self.get_analog_tile(model)
        alpha = analog_tile.get_scales().detach().cpu()
        self.assertNotEqual(alpha, 1.0)

        # Save the model to a file.
        with TemporaryFile() as file:
            save(model, file)
            # Load the model.
            file.seek(0)
            new_model = load(file)

        # Assert over the new model tile parameters.
        new_analog_tile = self.get_analog_tile(new_model)
        alpha_new = new_analog_tile.get_scales().detach().cpu()
        assert_array_almost_equal(array(alpha), array(alpha_new))

    def test_remapping(self):
        """Test remapping of the weights."""
        # Create the device and the array.
        rpu_config = InferenceRPUConfig(mapping=MappingParameter(
            weight_scaling_omega=0.4,
            weight_scaling_columnwise=True))

        model = self.get_layer(rpu_config=rpu_config)
        analog_tile = self.get_analog_tile(model)

        user_weights = Tensor(model.out_features, model.in_features).uniform_(-0.1, 0.1)
        if self.bias:
            user_biases = Tensor(model.out_features).uniform_(-0.1, 0.1)
        else:
            user_biases = None
        model.set_weights(user_weights, user_biases, apply_weight_scaling=True,
                          force_exact=True)

        alpha_initial = analog_tile.get_scales().detach().cpu()
        self.assertNotEqual(alpha_initial.max(), 1.0)

        # remap module
        model.remap_weights(weight_scaling_omega=1.0)

        weights, _ = model.get_weights(apply_weight_scaling=True, force_exact=True)
        analog_weights, _ = model.get_weights(apply_weight_scaling=False, force_exact=True)

        assert_array_almost_equal(array(user_weights), array(weights))

        alpha = analog_tile.get_scales().detach().cpu()
        assert_raises(AssertionError, assert_array_almost_equal, array(alpha), array(alpha_initial))

        self.assertEqual(torch_abs(analog_weights).max(), 1.0)

    def test_save_load_state_dict_hidden_parameters(self):
        """Test saving and loading via state_dict with hidden parameters."""
        # Create the device and the array.
        model = self.get_layer()
        analog_tile = self.get_analog_tile(model)
        hidden_parameters = analog_tile.tile.get_hidden_parameters()

        # Save the model to a file.
        with TemporaryFile() as file:
            save(model.state_dict(), file)
            # Load the model.
            file.seek(0)
            new_model = self.get_layer()
            new_model.load_state_dict(load(file))

        # Assert over the new model tile parameters.
        new_analog_tile = self.get_analog_tile(new_model)
        new_hidden_parameters = new_analog_tile.tile.get_hidden_parameters()
        assert_array_almost_equal(hidden_parameters, new_hidden_parameters)

    def test_state_dict_children_layers_sequential(self):
        """Test using the state_dict with children analog layers via Sequential."""
        children_layer = self.get_layer()
        model = Sequential(children_layer)

        # Keep track of the current weights and biases for comparing.
        (model_weights, model_biases,
         tile_weights, tile_biases, _) = self.get_layer_and_tile_weights(children_layer)

        state_dict = model.state_dict()
        lst = [key for key in state_dict.keys() if key.startswith('0.')
               and children_layer.ANALOG_STATE_PREFIX in key]
        self.assertTrue(len(lst) > 0)

        # Update the state_dict of a new model.
        new_children_layer = self.get_layer()
        new_model = Sequential(new_children_layer)
        new_model.load_state_dict(model.state_dict())

        # Compare the new model weights and biases.
        (new_model_weights, new_model_biases, new_tile_weights, new_tile_biases, _) = \
            self.get_layer_and_tile_weights(new_children_layer)

        assert_array_almost_equal(model_weights, new_model_weights)
        assert_array_almost_equal(tile_weights, new_tile_weights)
        if self.bias:
            assert_array_almost_equal(model_biases, new_model_biases)
            assert_array_almost_equal(tile_biases, new_tile_biases)

    def test_state_dict_children_layers_subclassing(self):
        """Test using the state_dict with children analog layers via subclassing."""

        class CustomModule(Module):
            """Module that defines its children layers via custom attributes."""
            # pylint: disable=abstract-method
            def __init__(self, layer: Module):
                super().__init__()
                self.custom_child = layer

        children_layer = self.get_layer()
        model = CustomModule(children_layer)

        # Keep track of the current weights and biases for comparing.
        (model_weights, model_biases, tile_weights, tile_biases, _) = \
            self.get_layer_and_tile_weights(children_layer)

        state_dict = model.state_dict()
        lst = [key for key in state_dict.keys() if key.startswith('custom_child.')
               and children_layer.ANALOG_STATE_PREFIX in key]
        self.assertTrue(len(lst) > 0)

        # Update the state_dict of a new model.
        new_children_layer = self.get_layer()
        new_model = CustomModule(new_children_layer)
        new_model.load_state_dict(model.state_dict())

        # Compare the new model weights and biases.
        (new_model_weights, new_model_biases, new_tile_weights, new_tile_biases, _) = \
            self.get_layer_and_tile_weights(new_children_layer)

        assert_array_almost_equal(model_weights, new_model_weights)
        assert_array_almost_equal(tile_weights, new_tile_weights)
        if self.bias:
            assert_array_almost_equal(model_biases, new_model_biases)
            assert_array_almost_equal(tile_biases, new_tile_biases)

    def test_state_dict_analog_strict(self):
        """Test the `strict` flag for analog layers."""
        model = self.get_layer()
        state_dict = model.state_dict()

        # Remove the analog key from the state dict.
        lst = [key for key in state_dict.keys() if key.startswith(model.ANALOG_STATE_PREFIX)]
        del state_dict[lst[0]]

        # Check that it fails when using `strict`.
        with self.assertRaises(RuntimeError) as context:
            model.load_state_dict(state_dict, strict=True)
        self.assertIn('Missing key', str(context.exception))

        # Check that it passes when not using `strict`.
        model.load_state_dict(state_dict, strict=False)

    def test_state_dict(self):
        """Test creating a new model using a state dict, without saving to disk."""
        model = self.get_layer()
        state_dict = model.state_dict()
        analog_tile = self.get_analog_tile(model)

        new_model = self.get_layer()
        new_model.load_state_dict(state_dict)
        new_analog_tile = self.get_analog_tile(new_model)

        # Asserts over the AnalogContext of the new model.
        self.assertTrue(hasattr(new_analog_tile.analog_ctx, 'analog_tile'))
        self.assertIsInstance(new_analog_tile.analog_ctx.analog_tile,
                              analog_tile.__class__)

    def test_hidden_parameter_mismatch(self):
        """Test for error if tile structure mismatches."""
        model = self.get_layer()
        state_dict = model.state_dict()
        analog_tile = self.get_analog_tile(model)

        # Create the device and the array.
        rpu_config = SingleRPUConfig(
            device=LinearStepDevice()  # different hidden structure
        )

        new_model = self.get_layer(rpu_config=rpu_config)
        new_analog_tile = self.get_analog_tile(new_model)
        if new_analog_tile.__class__.__name__ != analog_tile.__class__.__name__:
            with self.assertRaises(TileError):
                new_model.load_state_dict(state_dict)

    def test_load_state_load_rpu_config(self):
        """Test creating a new model using a state dict, while using a different RPU config."""

        # Create the device and the array.
        rpu_config_org = self.get_rpu_config()

        # Skipped for FP
        if isinstance(rpu_config_org, FloatingPointRPUConfig):
            raise SkipTest('Not available for FP')

        rpu_config_org.forward.is_perfect = False
        old_value = 0.11
        rpu_config_org.forward.inp_noise = old_value

        model = self.get_layer(rpu_config=rpu_config_org)
        state_dict = model.state_dict()

        rpu_config = deepcopy(rpu_config_org)
        new_value = 0.51
        rpu_config.forward.inp_noise = new_value

        # Test restore_rpu_config=False
        new_model = self.get_layer(rpu_config=rpu_config)
        new_model.load_state_dict(state_dict, load_rpu_config=False)
        new_analog_tile = self.get_analog_tile(new_model)

        parameters = new_analog_tile.tile.get_parameters()
        self.assertAlmostEqual(parameters.forward_io.inp_noise, new_value)

        # Test restore_rpu_config=True
        new_model = self.get_layer(rpu_config=rpu_config)
        new_model.load_state_dict(state_dict, load_rpu_config=True)
        new_analog_tile = self.get_analog_tile(new_model)

        parameters = new_analog_tile.tile.get_parameters()
        self.assertAlmostEqual(parameters.forward_io.inp_noise, old_value)

    def test_load_state_load_rpu_config_sequential(self):
        """Test creating a new model using a state dict, while using a different RPU config."""

        # Create the device and the array.
        rpu_config_org = self.get_rpu_config()

        # Skipped for FP
        if isinstance(rpu_config_org, FloatingPointRPUConfig):
            raise SkipTest('Not available for FP')

        rpu_config_org.forward.is_perfect = False
        old_value = 0.11
        rpu_config_org.forward.inp_noise = old_value

        model = AnalogSequential(self.get_layer(rpu_config=rpu_config_org))
        state_dict = model.state_dict()

        rpu_config = deepcopy(rpu_config_org)
        new_value = 0.51
        rpu_config.forward.inp_noise = new_value

        # Test restore_rpu_config=False
        new_model = AnalogSequential(self.get_layer(rpu_config=rpu_config))
        new_model.load_state_dict(state_dict, load_rpu_config=False)
        new_analog_tile = self.get_analog_tile(new_model[0])

        parameters = new_analog_tile.tile.get_parameters()
        self.assertAlmostEqual(parameters.forward_io.inp_noise, new_value)

        # Test restore_rpu_config=True
        new_model = AnalogSequential(self.get_layer(rpu_config=rpu_config))
        new_model.load_state_dict(state_dict, load_rpu_config=True)
        new_analog_tile = self.get_analog_tile(new_model[0])

        parameters = new_analog_tile.tile.get_parameters()
        self.assertAlmostEqual(parameters.forward_io.inp_noise, old_value)

    def test_load_state_load_rpu_config_wrong(self):
        """Test creating a new model using a state dict, while using a different RPU config."""

        # Skipped for FP
        if isinstance(self.get_rpu_config(), FloatingPointRPUConfig):
            raise SkipTest('Not available for FP')

        # Create the device and the array.
        model = self.get_layer()
        state_dict = model.state_dict()

        rpu_config = FloatingPointRPUConfig()

        new_model = self.get_layer(rpu_config=rpu_config)
        assert_raises(ModuleError, new_model.load_state_dict, state_dict, load_rpu_config=False)


@parametrize_over_layers(
    layers=[Linear, LinearCuda],
    tiles=[FloatingPoint],
    biases=[None],
)
class SerializationTestExtended(ParametrizedTestCase):
    """Tests for serialization."""

    @staticmethod
    def train_model_torch(model, loss_func, x_b, y_b):
        """Train the model with torch SGD."""
        opt = SGD(model.parameters(), lr=0.5)
        epochs = 100
        for _ in range(epochs):
            opt.zero_grad()
            pred = model(x_b)
            loss = loss_func(pred, y_b)

            loss.backward()
            opt.step()

    @staticmethod
    def get_torch_model(use_cuda: bool):
        """ Returns a torch model."""
        manual_seed(4321)
        torch_model = Sequential(
            torch_linear(4, 3),
            torch_linear(3, 3),
            Sequential(
                torch_linear(3, 1),
                torch_linear(1, 1)
            )
        )
        if use_cuda:
            torch_model.cuda()
        return torch_model

    def test_load_state_dict_conversion(self):
        """Test loading and setting conversion with alpha."""

        # Create the device and the array.

        x_b = Tensor([[0.1, 0.2, 0.3, 0.4], [0.2, 0.4, 0.3, 0.1]])
        y_b = Tensor([[0.3], [0.6]])
        if self.use_cuda:
            x_b = x_b.cuda()
            y_b = y_b.cuda()

        model = self.get_torch_model(self.use_cuda)

        self.train_model_torch(model, mse_loss, x_b, y_b)

        analog_model = convert_to_analog(model, self.get_rpu_config())
        analog_loss = mse_loss(analog_model(x_b), y_b)

        with TemporaryFile() as file:
            save(analog_model.state_dict(), file)
            # Load the model.
            file.seek(0)
            model = self.get_torch_model(self.use_cuda)
            new_analog_model = convert_to_analog(model, self.get_rpu_config())
            state_dict = load(file)
            new_analog_model.load_state_dict(state_dict, load_rpu_config=True)

        new_state_dict = new_analog_model.state_dict()
        for key in new_state_dict.keys():

            if not key.endswith('analog_tile_state'):
                continue

            state1 = new_state_dict[key]
            state2 = state_dict[key]
            assert_array_almost_equal(state1['analog_tile_weights'],
                                      state2['analog_tile_weights'])
            assert_array_almost_equal(state1['analog_alpha_scale'],
                                      state2['analog_alpha_scale'])

        new_analog_loss = mse_loss(new_analog_model(x_b), y_b)
        self.assertTensorAlmostEqual(new_analog_loss, analog_loss)
