#!/usr/bin/env python3
# Copyright (c) 2020 Graphcore Ltd. All rights reserved.

import torch
import pytest
import helpers
import poptorch


class ConstantBuffer(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.register_buffer('stuff', torch.tensor([1, 2, 3],
                                                   dtype=torch.int32))

    def forward(self, x):
        new_stuff = 1.0 + self.stuff
        return torch.sum(x + new_stuff)


def test_constant_buffer():
    model = ConstantBuffer()

    poptorch_model = poptorch.inferenceModel(model)
    assert poptorch_model(torch.tensor([2])) == 15


def test_constant_buffer_repeat():
    model = ConstantBuffer()

    poptorch_model = poptorch.inferenceModel(model)
    assert poptorch_model(torch.tensor([2])) == 15
    assert poptorch_model(torch.tensor([2])) == 15


class UpdatableBuffer(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.register_buffer("buffer_1", torch.ones(1))
        self.register_buffer("buffer_2", torch.ones(1))

    def forward(self, x):
        return torch.sum(x + self.buffer_1 + self.buffer_2)


def test_copy_named_buffer_to_device_single_buffer():
    model = UpdatableBuffer()
    options = poptorch.Options()
    options.updatableNamedBuffers(['buffer_1'])
    poptorch_model = poptorch.inferenceModel(model, options=options)
    x = torch.ones(3).float()

    assert poptorch_model(x) == 9

    poptorch_model.buffer_1.copy_(poptorch_model.buffer_1 + 1)
    poptorch_model.copyNamedBuffersToDevice()

    assert poptorch_model(x) == 12

    poptorch_model.buffer_2.copy_(poptorch_model.buffer_2 + 1)
    poptorch_model.copyNamedBuffersToDevice()

    assert poptorch_model(x) == 12


def test_copy_named_buffer_to_device_two_buffers():
    model = UpdatableBuffer()
    options = poptorch.Options()
    options.updatableNamedBuffers(['buffer_1', 'buffer_2'])
    poptorch_model = poptorch.inferenceModel(model, options=options)
    x = torch.ones(3).float()

    assert poptorch_model(x) == 9

    poptorch_model.buffer_1.copy_(poptorch_model.buffer_1 + 1)
    poptorch_model.copyNamedBuffersToDevice()

    assert poptorch_model(x) == 12

    poptorch_model.buffer_2.copy_(poptorch_model.buffer_2 + 1)
    poptorch_model.copyNamedBuffersToDevice()

    assert poptorch_model(x) == 15


def test_copy_named_buffer_to_device_no_opt():
    model = UpdatableBuffer()
    options = poptorch.Options()

    poptorch_model = poptorch.inferenceModel(model, options=options)
    x = torch.ones(3).float()

    assert poptorch_model(x) == 9

    poptorch_model.buffer_1.copy_(poptorch_model.buffer_1 + 1)
    with pytest.raises(poptorch.poptorch_core.Error):
        poptorch_model.copyNamedBuffersToDevice()

    assert poptorch_model(x) == 9


def test_copy_named_buffer_to_device_invalid_opt():
    model = UpdatableBuffer()
    options = poptorch.Options()
    options.updatableNamedBuffers(['non_existing_buffer'])

    poptorch_model = poptorch.inferenceModel(model, options=options)
    x = torch.ones(3).float()

    with pytest.raises(poptorch.poptorch_core.Error):
        poptorch_model(x)


def test_training_then_inference():
    momentum = 0.1

    class Model(torch.nn.Module):
        def __init__(self):
            super().__init__()
            self.bn = torch.nn.BatchNorm1d(10, momentum=momentum)
            self.loss = torch.nn.MSELoss()

        def forward(self, x, target):
            y = self.bn(x)
            return y, self.loss(y, target)

    model = Model()

    input = torch.ones([4, 10], dtype=torch.float32)
    target = torch.ones([4, 10], dtype=torch.float32) + 1

    training_model = poptorch.trainingModel(model)

    training_model.compile(input, target)

    inference_model = poptorch.inferenceModel(model)

    inference_model.compile(input, target)


def test_buffer_implicit_copy():
    momentum = 0.1

    class Model(torch.nn.Module):
        def __init__(self):
            super().__init__()
            self.bn = torch.nn.BatchNorm1d(10, momentum=momentum)
            self.loss = torch.nn.MSELoss()

        def forward(self, x, target):
            y = self.bn(x)
            return y, self.loss(y, target)

    model = Model()

    input = torch.ones([4, 10], dtype=torch.float32)
    target = torch.ones([4, 10], dtype=torch.float32) + 1

    poptorch_model = poptorch.trainingModel(model)

    poptorch_model(input, target)
    helpers.assert_allclose(actual=model.bn.running_mean,
                            expected=input[0, :] * momentum)

    poptorch_model.copyWeightsToHost()
    helpers.assert_allclose(actual=model.bn.running_mean,
                            expected=input[0, :] * momentum)


def test_error_on_remove_buffer():
    class Model(torch.nn.Module):
        def __init__(self):
            super().__init__()
            self.register_buffer('y', torch.tensor([2]))

        def forward(self, x):
            x = x + 1
            if 'y' in self._buffers:
                del self._buffers['y']
            return x

    model = Model()

    poptorch_model = poptorch.inferenceModel(model)

    error_msg = (r"Buffer y is removed from the model when calling the "
                 r"forward method\.")
    with pytest.raises(poptorch.Error, match=error_msg):
        poptorch_model(torch.tensor([5.0]))


def test_error_on_redefine_buffer():
    class Model(torch.nn.Module):
        def __init__(self):
            super().__init__()
            self.register_buffer('y', torch.tensor([2]))

        def forward(self, x):
            x = x + 1
            # pylint: disable=attribute-defined-outside-init
            self.y = x

    model = Model()

    poptorch_model = poptorch.inferenceModel(model)
    error_msg = (r"Buffer y is reassigned within the model when calling the "
                 r"forward method\. This is not supported\. Consider using "
                 r"self\.y\.copy_\(src\) to copy data "
                 r"from a source tensor, where src is the name of the "
                 r"source tensor\.")

    with pytest.raises(poptorch.Error, match=error_msg):
        poptorch_model(torch.tensor([5.0]))


class BufferUpdatingModel(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.conv = torch.nn.Conv2d(2, 2, 1, padding=0)
        self.register_buffer("test_buff", torch.zeros([2],
                                                      dtype=torch.float32))

        self.loss = torch.nn.L1Loss()

    def forward(self, inp, target):
        x = self.conv(inp)

        with torch.no_grad():
            # pylint: disable=attribute-defined-outside-init
            self.test_buff += self.conv.bias[0]

        return x, self.loss(x, target)


@pytest.mark.parametrize("device_iterations", [1, 3, 5])
@pytest.mark.parametrize("gradient_accumulation", [1, 3, 5])
def test_buffer_update_with_param(device_iterations, gradient_accumulation):
    model = BufferUpdatingModel()
    model.conv.weight.data = torch.ones_like(model.conv.weight.data)
    model.conv.bias.data = torch.ones_like(model.conv.bias.data)
    opt = torch.optim.SGD(model.parameters(), lr=0.1)

    times_to_run = 10
    dummy_input = torch.ones([2, 2, 2, 2])
    dummy_target = torch.zeros_like(dummy_input)

    for _ in range(times_to_run * device_iterations):
        opt.zero_grad()
        for _ in range(gradient_accumulation):
            _, loss = model(dummy_input, dummy_target)

            # Match mean gradient_accumulation
            loss /= gradient_accumulation

            loss.backward()

        opt.step()

    model_bias = model.conv.bias.clone()
    model_test_buff = model.test_buff.clone()

    # pylint: disable=attribute-defined-outside-init
    model.test_buff = torch.zeros([2], dtype=torch.float32)
    model.conv.weight.data = torch.ones_like(model.conv.weight.data)
    model.conv.bias.data = torch.ones_like(model.conv.bias.data)

    # Check for proper cloning
    with pytest.raises(AssertionError):
        helpers.assert_allclose(expected=model_bias, actual=model.conv.bias)
    with pytest.raises(AssertionError):
        helpers.assert_allclose(expected=model_test_buff,
                                actual=model.test_buff)

    opts = poptorch.Options()
    opts.deviceIterations(device_iterations)
    opts.Training.gradientAccumulation(gradient_accumulation)

    dummy_input = torch.ones(
        [2 * device_iterations * gradient_accumulation, 2, 2, 2])
    dummy_target = torch.zeros_like(dummy_input)

    poptorch_model = poptorch.trainingModel(model,
                                            optimizer=torch.optim.SGD(
                                                model.parameters(), lr=0.1),
                                            options=opts)

    for _ in range(times_to_run):
        dummy_target = torch.zeros_like(dummy_input)
        poptorch_model(dummy_input, dummy_target)

    helpers.assert_allclose(expected=model_bias,
                            actual=poptorch_model.conv.bias)
    helpers.assert_allclose(expected=model_test_buff,
                            actual=poptorch_model.test_buff)


def test_failing_on_replicas():
    model = BufferUpdatingModel()

    opts = poptorch.Options()
    opts.replicationFactor(2)

    poptorch_model = poptorch.trainingModel(model,
                                            optimizer=torch.optim.SGD(
                                                model.parameters(), lr=0.1),
                                            options=opts)

    dummy_input = torch.ones([4, 2, 2, 2])
    dummy_target = torch.zeros_like(dummy_input)

    error_msg = (r"PopTorch does not support broadcasting buffers. "
                 r"If your model is able to tolerate buffers becoming "
                 r"out of sync between replicas, you can disable "
                 r"buffer broadcasting using "
                 r"poptorch.Options.broadcastBuffers\(False\).")

    with pytest.raises(poptorch.Error, match=error_msg):
        poptorch_model(dummy_input, dummy_target)


def test_constant_buffer_with_replicas():
    # This should not have an error as the buffer is constant
    model = ConstantBuffer()

    opts = poptorch.Options()
    opts.replicationFactor(2)

    poptorch_model = poptorch.inferenceModel(model, opts)
    poptorch_model(torch.tensor([1, 2]))


def test_no_input_but_one_buffer():
    class Model(torch.nn.Module):
        def __init__(self):
            super().__init__()
            self.register_buffer("x", torch.tensor([1.], dtype=torch.float))

        def forward(self):
            # pylint: disable=attribute-defined-outside-init,no-member
            self.x += 1.0
            return self.x

    model = Model()
    poptorch_model = poptorch.inferenceModel(model)

    assert poptorch_model() == 2.
    assert poptorch_model() == 3.
    assert poptorch_model() == 4.
    assert poptorch_model() == 5.


def test_unsynchronised_replicated_buffers():
    class ReplicaBufferModel(torch.nn.Module):
        def __init__(self):
            super().__init__()
            self.register_buffer("buffer", torch.zeros(1, 2))

        def forward(self, x):
            buffer_update = self.buffer + x
            self.buffer.copy_(buffer_update)
            return poptorch.identity_loss(self.buffer, reduction='none')

    num_replica = 2
    torch.manual_seed(43)
    opts = poptorch.Options()
    opts.replicationFactor(num_replica)
    opts.deviceIterations(1)
    opts.broadcastBuffers(False)

    model = ReplicaBufferModel()
    model.float()
    poptorch_model = poptorch.inferenceModel(model, opts)

    x = torch.tensor([[9], [2]])

    # Each replica update its buffer in place with a random value 50 times.
    for _ in range(50):
        y = poptorch_model(x)

    assert y[0][-1] == x[0] * 50
    assert y[1][-1] == x[1] * 50