from test_pytorch_common import TestCase, run_tests, skipIfNoLapack, flatten
import test_onnx_common

import torch
import torch.onnx
from torch.autograd import Variable, Function
from torch.nn import Module
import torch.nn as nn

import onnx
import onnx.checker
import onnx.helper

import google.protobuf.text_format

import itertools
import io
import unittest
import inspect
import argparse
import glob
import os
import shutil
import sys
import common
from onnx import numpy_helper

_onnx_test = False


def export_to_string(model, inputs, *args, **kwargs):
    f = io.BytesIO()
    with torch.no_grad():
        torch.onnx.export(model, inputs, f, *args, **kwargs)
    return f.getvalue()


class FuncModule(Module):
    def __init__(self, f, params=tuple()):
        super(FuncModule, self).__init__()
        self.f = f
        self.params = nn.ParameterList(list(params))

    def forward(self, *args):
        return self.f(*itertools.chain(args, self.params))


class TestOperators(TestCase):

    def assertONNXExpected(self, binary_pb, subname=None):
        model_def = onnx.ModelProto.FromString(binary_pb)
        onnx.checker.check_model(model_def)
        # doc_string contains stack trace in it, strip it
        onnx.helper.strip_doc_string(model_def)
        self.assertExpected(google.protobuf.text_format.MessageToString(model_def, float_format='.15g'), subname)
        return model_def

    def assertONNX(self, f, args, params=tuple(), **kwargs):
        if isinstance(f, nn.Module):
            m = f
        else:
            m = FuncModule(f, params)
        onnx_model_pb = export_to_string(m, args, **kwargs)
        model_def = self.assertONNXExpected(onnx_model_pb)
        if _onnx_test:
            test_function = inspect.stack()[1][0].f_code.co_name
            test_name = test_function[0:4] + "_operator" + test_function[4:]
            output_dir = os.path.join(test_onnx_common.pytorch_operator_dir, test_name)
            # Assume:
            #     1) the old test should be delete before the test.
            #     2) only one assertONNX in each test, otherwise will override the data.
            assert not os.path.exists(output_dir), "{} should not exist!".format(output_dir)
            os.makedirs(output_dir)
            with open(os.path.join(output_dir, "model.onnx"), 'wb') as file:
                file.write(model_def.SerializeToString())
            data_dir = os.path.join(output_dir, "test_data_set_0")
            os.makedirs(data_dir)
            if isinstance(args, Variable):
                args = (args,)
            for index, var in enumerate(flatten(args)):
                tensor = numpy_helper.from_array(var.data.numpy())
                with open(os.path.join(data_dir, "input_{}.pb".format(index)), 'wb') as file:
                    file.write(tensor.SerializeToString())
            outputs = m(*args)
            if isinstance(outputs, Variable):
                outputs = (outputs,)
            for index, var in enumerate(flatten(outputs)):
                tensor = numpy_helper.from_array(var.data.numpy())
                with open(os.path.join(data_dir, "output_{}.pb".format(index)), 'wb') as file:
                    file.write(tensor.SerializeToString())

    def assertONNXRaises(self, err, f, args, params=tuple(), **kwargs):
        if isinstance(f, nn.Module):
            m = f
        else:
            m = FuncModule(f, params)
        self.assertExpectedRaises(err, lambda: export_to_string(m, args, **kwargs))

    def assertONNXRaisesRegex(self, err, reg, f, args, params=tuple(), **kwargs):
        if isinstance(f, nn.Module):
            m = f
        else:
            m = FuncModule(f, params)
        with self.assertRaisesRegex(err, reg):
            export_to_string(m, args, **kwargs)

    def test_basic(self):
        x = Variable(torch.Tensor([0.4]), requires_grad=True)
        y = Variable(torch.Tensor([0.7]), requires_grad=True)
        self.assertONNX(lambda x, y: -torch.sigmoid(torch.tanh(x * (x + y))), (x, y))

    def test_view(self):
        x = Variable(torch.Tensor([0]), requires_grad=True)
        self.assertONNX(lambda x: x.view(1, 1), x)

    def test_index(self):
        x = Variable(torch.Tensor([[0]]), requires_grad=True)
        self.assertONNX(lambda x: x[0], x)

    def test_type_as(self):
        x = Variable(torch.Tensor([0]), requires_grad=True)
        self.assertONNX(lambda x: x.type_as(x), x)

    def test_addconstant(self):
        x = Variable(torch.DoubleTensor(2, 3), requires_grad=True)
        self.assertONNX(lambda x: x + 1, x)

    def test_add_broadcast(self):
        x = Variable(torch.DoubleTensor(2, 3), requires_grad=True)
        y = Variable(torch.DoubleTensor(3), requires_grad=True)
        self.assertONNX(lambda x, y: x + y, (x, y))

    def test_add_left_broadcast(self):
        x = Variable(torch.DoubleTensor(3), requires_grad=True)
        y = Variable(torch.DoubleTensor(2, 3), requires_grad=True)
        self.assertONNXRaisesRegex(RuntimeError,
                                   r"ONNX export failed: Could not export a broadcasted operation.*",
                                   lambda x, y: x + y, (x, y), verbose=True)

    def test_add_size1_broadcast(self):
        x = Variable(torch.DoubleTensor(2, 3), requires_grad=True)
        y = Variable(torch.DoubleTensor(2, 1), requires_grad=True)
        self.assertONNX(lambda x, y: x + y, (x, y))

    def test_add_size1_right_broadcast(self):
        x = Variable(torch.DoubleTensor(2, 3), requires_grad=True)
        y = Variable(torch.DoubleTensor(3), requires_grad=True)
        self.assertONNX(lambda x, y: x + y, (x, y))

    def test_add_size1_singleton_broadcast(self):
        x = Variable(torch.DoubleTensor(2, 3), requires_grad=True)
        y = Variable(torch.DoubleTensor(1, 3), requires_grad=True)
        self.assertONNX(lambda x, y: x + y, (x, y))

    def test_transpose(self):
        x = Variable(torch.Tensor([[0, 1], [2, 3]]), requires_grad=True)
        self.assertONNX(lambda x: x.transpose(0, 1).transpose(1, 0), x)

    def test_chunk(self):
        x = Variable(torch.Tensor([0, 1, 2]), requires_grad=True)
        self.assertONNX(lambda x: x.chunk(2), x)

    def test_concat2(self):
        x = Variable(torch.randn(2, 3))
        y = Variable(torch.randn(2, 3))
        self.assertONNX(lambda inputs: torch.cat(inputs, 1), ((x, y),))

    def test_mm(self):
        m1 = Variable(torch.randn(2, 3), requires_grad=True)
        m2 = Variable(torch.randn(3, 4), requires_grad=True)
        self.assertONNX(torch.mm, (m1, m2))

    def test_addmm(self):
        m1 = Variable(torch.randn(2, 3), requires_grad=True)
        m2 = Variable(torch.randn(3, 4), requires_grad=True)
        m3 = Variable(torch.randn(4), requires_grad=True)
        self.assertONNX(lambda x, y, z: torch.addmm(torch.addmm(z, x, y), x, y), (m1, m2, m3))

    def test_permute2(self):
        x = Variable(torch.Tensor([[[[[[0]]]]]]), requires_grad=True)
        self.assertONNX(lambda x: x.permute(0, 1, 4, 2, 5, 3), x)

    def test_pad(self):
        x = Variable(torch.Tensor([[[[0, 1, 1, 1], [2, 3, 7, 7]]]]), requires_grad=True)
        self.assertONNX(nn.ReflectionPad2d((2, 3, 0, 1)), x)

    def test_params(self):
        x = Variable(torch.Tensor([[1, 2], [3, 4]]), requires_grad=True)
        y = nn.Parameter(torch.Tensor([[1, 2], [3, 4]]), requires_grad=True)
        self.assertONNX(lambda x, y: -torch.sigmoid(torch.tanh(x * (x + y))), x, params=(y, ))

    def test_non_float_params(self):
        x = Variable(torch.LongTensor([[1, 2], [3, 4]]), requires_grad=True)
        y = nn.Parameter(torch.LongTensor([[1, 2], [3, 4]]), requires_grad=True)
        self.assertONNX(lambda x, y: x * (x + y), x, params=(y, ))

    def test_symbolic_mismatch(self):
        class MyFun(Function):
            @staticmethod
            def symbolic(g, x):
                # The inside of this function should never be invoked, because
                # we will fail due to an argument mismatch first.
                assert False

            @staticmethod
            def forward(ctx, x, y):
                return x + y

        x = Variable(torch.randn(2, 2).fill_(1.0))
        y = Variable(torch.randn(2, 2).fill_(1.0))
        # NB: Don't use expect test here, the type error wobbles depending
        # on Python version
        with self.assertRaisesRegex(TypeError, "occurred when translating MyFun"):
            export_to_string(FuncModule(MyFun().apply), (x, y))

    # TODO: Do an nn style test for these
    def test_batchnorm(self):
        x = Variable(torch.randn(2, 2, 2, 2).fill_(1.0), requires_grad=True)
        self.assertONNX(nn.BatchNorm2d(2), x)

    def test_batchnorm_1d(self):
        x = Variable(torch.randn(2, 2).fill_(1.0), requires_grad=True)
        self.assertONNX(nn.BatchNorm1d(2), x)

    def test_batchnorm_training(self):
        x = Variable(torch.randn(2, 2, 2, 2).fill_(1.0), requires_grad=True)
        self.assertONNX(nn.BatchNorm2d(2), x, training=True)

    def test_conv(self):
        x = Variable(torch.randn(20, 16, 50, 40).fill_(1.0), requires_grad=True)
        self.assertONNX(nn.Conv2d(16, 13, 3, bias=False), x)

    def test_convtranspose(self):
        x = Variable(torch.randn(2, 3, 4, 5).fill_(1.0), requires_grad=True)
        self.assertONNX(nn.ConvTranspose2d(3, 3, 3, stride=3, bias=False,
                                           padding=1, output_padding=2), x)

    def test_maxpool(self):
        x = Variable(torch.randn(20, 16, 50))
        self.assertONNX(nn.MaxPool1d(3, stride=2), x)

    def test_at_op(self):
        x = Variable(torch.randn(3, 4))

        class MyFun(Function):

            @staticmethod
            def symbolic(g, x):
                return g.at("add", x, x)

            @staticmethod
            def forward(ctx, x):
                return x + x

        class MyModule(Module):
            def forward(self, x):
                return MyFun.apply(x)

        self.assertONNX(MyModule(), x)

    def test_clip(self):
        x = Variable(torch.randn(3, 4), requires_grad=True)
        self.assertONNX(lambda x: torch.clamp(x, min=-0.5, max=0.5), x)

    def test_max(self):
        x = Variable(torch.randn(3, 4), requires_grad=True)
        y = Variable(torch.randn(3, 4), requires_grad=True)
        self.assertONNX(lambda x, y: torch.max(x, y), (x, y))

    def test_min(self):
        x = Variable(torch.randn(3, 4), requires_grad=True)
        y = Variable(torch.randn(3, 4), requires_grad=True)
        self.assertONNX(lambda x, y: torch.min(x, y), (x, y))

    def test_mean(self):
        x = Variable(torch.randn(1, 2, 3, 4), requires_grad=True)
        self.assertONNX(lambda x: torch.mean(x), x)

    def test_reduced_mean(self):
        x = Variable(torch.randn(1, 2, 3, 4), requires_grad=True)
        self.assertONNX(lambda x: torch.mean(x, dim=2), x)

    def test_reduced_mean_keepdim(self):
        x = Variable(torch.randn(1, 2, 3, 4), requires_grad=True)
        self.assertONNX(lambda x: torch.mean(x, dim=2, keepdim=True), x)

    def test_sum(self):
        x = Variable(torch.randn(1, 2, 3, 4), requires_grad=True)
        self.assertONNX(lambda x: torch.sum(x), x)

    def test_reduced_sum(self):
        x = Variable(torch.randn(1, 2, 3, 4), requires_grad=True)
        self.assertONNX(lambda x: torch.sum(x, dim=2), x)

    def test_reduced_sum_keepdim(self):
        x = Variable(torch.randn(1, 2, 3, 4), requires_grad=True)
        self.assertONNX(lambda x: torch.sum(x, dim=2, keepdim=True), x)

    def test_prod(self):
        x = Variable(torch.randn(1, 2, 3, 4), requires_grad=True)
        self.assertONNX(lambda x: torch.prod(x), x)

    def test_reduced_prod(self):
        x = Variable(torch.randn(1, 2, 3, 4), requires_grad=True)
        self.assertONNX(lambda x: torch.prod(x, dim=2), x)

    def test_reduced_prod_keepdim(self):
        x = Variable(torch.randn(1, 2, 3, 4), requires_grad=True)
        self.assertONNX(lambda x: torch.prod(x, dim=2, keepdim=True), x)

    def test_sqrt(self):
        x = Variable(torch.randn(3, 4), requires_grad=True)
        self.assertONNX(lambda x: torch.sqrt(x), x)

    def test_equal(self):
        x = Variable(torch.randn(3, 4).int(), requires_grad=True)
        y = Variable(torch.randn(3, 4).int(), requires_grad=True)
        self.assertONNX(lambda x, y: x == y, (x, y))

    def test_lt(self):
        x = Variable(torch.randn(3, 4).int(), requires_grad=True)
        y = Variable(torch.randn(3, 4).int(), requires_grad=True)
        self.assertONNX(lambda x, y: x < y, (x, y))

    def test_gt(self):
        x = Variable(torch.randn(3, 4).int(), requires_grad=True)
        y = Variable(torch.randn(3, 4).int(), requires_grad=True)
        self.assertONNX(lambda x, y: x > y, (x, y))

    def test_le(self):
        x = Variable(torch.randn(3, 4).int(), requires_grad=True)
        y = Variable(torch.randn(3, 4).int(), requires_grad=True)
        self.assertONNX(lambda x, y: x <= y, (x, y))

    def test_ge(self):
        x = Variable(torch.randn(3, 4).int(), requires_grad=True)
        y = Variable(torch.randn(3, 4).int(), requires_grad=True)
        self.assertONNX(lambda x, y: x >= y, (x, y))

    def test_exp(self):
        x = Variable(torch.randn(3, 4), requires_grad=True)
        self.assertONNX(lambda x: x.exp(), x)

    def test_flatten(self):
        # Flatten is a special case of Reshape when the output is a 2-D tensor.
        x = Variable(torch.randn(1, 2, 3, 4), requires_grad=True)
        self.assertONNX(lambda x: x.view(x.size()[0], x.numel() // x.size()[0]), x)

    def test_logsoftmax(self):
        x = Variable(torch.randn(1, 2, 3, 4), requires_grad=True)
        self.assertONNX(nn.LogSoftmax(dim=2), x)

    def test_pow(self):
        x = Variable(torch.randn(1, 2, 3, 4), requires_grad=True)
        y = Variable(torch.randn(1, 2, 3, 4), requires_grad=True)
        self.assertONNX(lambda x, y: x.pow(y), (x, y))

    def test_selu(self):
        x = Variable(torch.randn(1, 2, 3, 4), requires_grad=True)
        self.assertONNX(nn.SELU(), x)

    def test_repeat(self):
        x = Variable(torch.randn(1, 2, 3, 4), requires_grad=True)
        self.assertONNX(lambda x: x.repeat(1, 2, 3, 4), x)

    def test_repeat_dim_overflow(self):
        x = Variable(torch.randn(1, 2), requires_grad=True)
        self.assertONNX(lambda x: x.repeat(1, 2, 3, 4), x)

    def test_symbolic_override(self):
        """Lifted from fast-neural-style: custom implementation of instance norm
        to be mapped to ONNX operator"""

        class CustomInstanceNorm(torch.nn.Module):
            def __init__(self, dim, eps=1e-9):
                super(CustomInstanceNorm, self).__init__()
                self.scale = nn.Parameter(torch.FloatTensor(dim).uniform_())
                self.shift = nn.Parameter(torch.FloatTensor(dim).zero_())
                self.eps = eps

            def forward(self, x):
                return self._run_forward(x, self.scale, self.shift, eps=self.eps)

            @staticmethod
            @torch.onnx.symbolic_override(
                lambda g, x, scale, shift, eps: g.op(
                    'InstanceNormalization', x, scale, shift, epsilon_f=eps)
            )
            def _run_forward(x, scale, shift, eps):
                # since we hand-roll instance norm it doesn't perform well all in fp16
                n = x.size(2) * x.size(3)
                t = x.view(x.size(0), x.size(1), n)
                mean = torch.mean(t, 2).unsqueeze(2).unsqueeze(3).expand_as(x)
                # Calculate the biased var. torch.var returns unbiased var
                var = torch.var(t, 2).unsqueeze(2).unsqueeze(3).expand_as(x) * ((float(n) - 1) / float(n))
                scale_broadcast = scale.unsqueeze(1).unsqueeze(1).unsqueeze(0)
                scale_broadcast = scale_broadcast.expand_as(x)
                shift_broadcast = shift.unsqueeze(1).unsqueeze(1).unsqueeze(0)
                shift_broadcast = shift_broadcast.expand_as(x)
                out = (x - mean) / torch.sqrt(var + eps)
                out = out * scale_broadcast + shift_broadcast
                return out

        instnorm = CustomInstanceNorm(10)
        x = Variable(torch.randn(2, 10, 32, 32))
        self.assertONNX(instnorm, x)

    """
    def test_rnn(self):
        rnn = nn.RNN(30, 20, 2)
        input = Variable(torch.randn(10, 32, 30))
        output, hidden = rnn(input)
        self.assertONNX(rnn, input)
    """

    def test_symbolic_override_nested(self):
        def symb(g, x, y):
            assert isinstance(x, torch._C.Value)
            assert isinstance(y[0], torch._C.Value)
            assert isinstance(y[1], torch._C.Value)
            return g.op('Sum', x, y[0], y[1]), (
                g.op('Neg', x), g.op('Neg', y[0]))

        @torch.onnx.symbolic_override_first_arg_based(symb)
        def foo(x, y):
            return x + y[0] + y[1], (-x, -y[0])

        class BigModule(torch.nn.Module):
            def forward(self, x, y):
                return foo(x, y)

        inp = (Variable(torch.FloatTensor([1])),
               (Variable(torch.FloatTensor([2])),
                Variable(torch.FloatTensor([3]))))
        BigModule()(*inp)
        self.assertONNX(BigModule(), inp)


if __name__ == '__main__':
    onnx_test_flag = '--onnx-test'
    _onnx_test = onnx_test_flag in common.UNITTEST_ARGS
    if onnx_test_flag in common.UNITTEST_ARGS:
        common.UNITTEST_ARGS.remove(onnx_test_flag)
    if _onnx_test:
        for d in glob.glob(os.path.join(test_onnx_common.pytorch_operator_dir, "test_operator_*")):
            shutil.rmtree(d)
    run_tests()
