import numpy

from einops import rearrange, parse_shape, reduce
from tests import collect_test_backends
from tests.test_ops import imp_op_backends


def test_rearrange_examples():
    def test1(x):
        # transpose
        y = rearrange(x, 'b c h w -> b h w c')
        assert y.shape == (10, 30, 40, 20)
        return y

    def test2(x):
        # view / reshape
        y = rearrange(x, 'b c h w -> b (c h w)')
        assert y.shape == (10, 20 * 30 * 40)
        return y

    def test3(x):
        # depth-to-space
        y = rearrange(x, 'b (c h1 w1) h w -> b c (h h1) (w w1)', h1=2, w1=2)
        assert y.shape == (10, 5, 30 * 2, 40 * 2)
        return y

    def test4(x):
        # space-to-depth
        y = rearrange(x, 'b c (h h1) (w w1) -> b (h1 w1 c) h w', h1=2, w1=2)
        assert y.shape == (10, 20 * 4, 30 // 2, 40 // 2)
        return y

    def test5(x):
        # simple transposition
        y = rearrange(x, 'b1 sound b2 letter -> b1 b2 sound letter')
        assert y.shape == (10, 30, 20, 40)
        return y

    def test6(x):
        # parsing parameters
        t = rearrange(x, 'b c h w -> (b h w) c')
        t = t[:, ::2]  # replacement for dot-product, just changes size of second axis
        assert t.shape == (10 * 30 * 40, 10)

        y = rearrange(t, '(b h w) c2 -> b c2 h w', **parse_shape(x, 'b _ h w'))
        assert y.shape == (10, 10, 30, 40)
        return y

    def test7(x):
        # split of embedding into groups
        y1, y2 = rearrange(x, 'b (c g) h w -> g b c h w', g=2)
        assert y1.shape == (10, 10, 30, 40)
        assert y2.shape == (10, 10, 30, 40)
        return y1 + y2  # only one tensor is expected in output

    def test8(x):
        # max-pooling
        y = reduce(x, 'b c (h h1) (w w1) -> b c h w', reduction='max', h1=2, w1=2)
        assert y.shape == (10, 20, 30 // 2, 40 // 2)
        return y

    def test9(x):
        # squeeze - unsqueeze
        y = reduce(x, 'b c h w -> b c () ()', reduction='max')
        assert y.shape == (10, 20, 1, 1)
        y = rearrange(y, 'b c () () -> c b')
        assert y.shape == (20, 10)
        return y

    def test10(x):
        # stack
        tensors = list(x + 0)  # 0 is needed https://github.com/tensorflow/tensorflow/issues/23185
        tensors = rearrange(tensors, 'b c h w -> b h w c')
        assert tensors.shape == (10, 30, 40, 20)
        return tensors

    def test11(x):
        # concatenate
        tensors = list(x + 0)  # 0 is needed https://github.com/tensorflow/tensorflow/issues/23185
        tensors = rearrange(tensors, 'b c h w -> h (b w) c')
        assert tensors.shape == (30, 10 * 40, 20)
        return tensors

    def shufflenet(x, convolve, c1, c2):
        # shufflenet reordering example
        x = convolve(x)
        x = rearrange(x, 'b (c1 c2) h w-> b (c2 c1) h w', c1=c1, c2=c2)
        x = convolve(x)
        return x

    def convolve_strided_1d(x, stride, usual_convolution):
        x = rearrange(x, 'b c t1 t2 -> b c (t1 t2)')  # reduce dimensionality
        x = rearrange(x, 'b c (t stride) -> (stride b) c t', stride=stride)
        x = usual_convolution(x)
        x = rearrange(x, '(stride b) c t -> b c (t stride)', stride=stride)
        return x

    def convolve_strided_2d(x, h_stride, w_stride, usual_convolution):
        x = rearrange(x, 'b c (h hs) (w ws) -> (hs ws b) c h w', hs=h_stride, ws=w_stride)
        x = usual_convolution(x)
        x = rearrange(x, '(hs ws b) c h w -> b c (h hs) (w ws)', hs=h_stride, ws=w_stride)
        return x

    def unet_like_1d(x, usual_convolution):
        # u-net like steps for increasing / reducing dimensionality
        x = rearrange(x, 'b c t1 t2 -> b c (t1 t2)')  # reduce dimensionality
        y = rearrange(x, 'b c (t dt) -> b (dt c) t', dt=2)
        y = usual_convolution(y)
        x = x + rearrange(y, 'b (dt c) t -> b c (t dt)', dt=2)
        return x

    # mock for convolution (works for all backends)
    convolve_mock = lambda x: x

    tests = [test1, test2, test3, test4, test5, test6, test7, test8, test9, test10, test11,
             lambda x: shufflenet(x, convolve=convolve_mock, c1=4, c2=5),
             lambda x: convolve_strided_1d(x, stride=2, usual_convolution=convolve_mock),
             lambda x: convolve_strided_2d(x, h_stride=2, w_stride=2, usual_convolution=convolve_mock),
             lambda x: unet_like_1d(x, usual_convolution=convolve_mock),
             ]

    for backend in imp_op_backends:
        print('testing source_examples for ', backend.framework_name)
        for test in tests:
            x = numpy.arange(10 * 20 * 30 * 40).reshape([10, 20, 30, 40])
            result1 = test(x)
            result2 = backend.to_numpy(test(backend.from_numpy(x)))
            assert numpy.array_equal(result1, result2)

            # now with strides
            x = numpy.arange(10 * 2 * 20 * 3 * 30 * 1 * 40).reshape([10 * 2, 20 * 3, 30 * 1, 40 * 1])
            # known torch bug - torch doesn't support negative steps
            last_step = -1 if (backend.framework_name != 'torch' and backend.framework_name != 'oneflow') else 1
            indexing_expression = numpy.index_exp[::2, ::3, ::1, ::last_step]
            result1 = test(x[indexing_expression])
            result2 = backend.to_numpy(test(backend.from_numpy(x)[indexing_expression]))
            assert numpy.array_equal(result1, result2)


def tensor_train_example_numpy():
    # kept here just for a collection, only tested for numpy
    # https://arxiv.org/pdf/1509.06569.pdf, (5)
    x = numpy.ones([3, 4, 5, 6])
    rank = 4
    if numpy.__version__ < '1.15.0':
        # numpy.einsum fails here, skip test
        return
    # creating appropriate Gs
    Gs = [numpy.ones([d, d, rank, rank]) for d in x.shape]
    Gs[0] = Gs[0][:, :, :1, :]
    Gs[-1] = Gs[-1][:, :, :, :1]

    # einsum way
    y = x.reshape((1,) + x.shape)
    for G in Gs:
        # taking partial results left-to-right
        # y = numpy.einsum('i j alpha beta, alpha i ...  -> beta ... j', G, y)
        y = numpy.einsum('i j a b, a i ...  -> b ... j', G, y)
    y1 = y.reshape(-1)

    # alternative way
    y = x.reshape(-1)
    for G in Gs:
        i, j, alpha, beta = G.shape
        y = rearrange(y, '(i rest alpha) -> rest (alpha i)', alpha=alpha, i=i)
        y = y @ rearrange(G, 'i j alpha beta -> (alpha i) (j beta)')
        y = rearrange(y, 'rest (beta j) -> (beta rest j)', beta=beta, j=j)
    y2 = y
    assert numpy.allclose(y1, y2)

    # yet another way
    y = x
    for G in Gs:
        i, j, alpha, beta = G.shape
        y = rearrange(y, 'i ... (j alpha) -> ... j (alpha i)', alpha=alpha, i=i)
        y = y @ rearrange(G, 'i j alpha beta -> (alpha i) (j beta)')
    y3 = y.reshape(-1)
    assert numpy.allclose(y1, y3)


def test_pytorch_yolo_fragment():
    if not any(b.framework_name == 'torch' for b in collect_test_backends(symbolic=False, layers=False)):
        return
    import torch

    def old_way(input, num_classes, num_anchors, anchors, stride_h, stride_w):
        # https://github.com/BobLiu20/YOLOv3_PyTorch/blob/c6b483743598b5f64d520d81e7e5f47ba936d4c9/nets/yolo_loss.py#L28-L44
        bs = input.size(0)
        in_h = input.size(2)
        in_w = input.size(3)
        scaled_anchors = [(a_w / stride_w, a_h / stride_h) for a_w, a_h in anchors]

        prediction = input.view(bs, num_anchors,
                                5 + num_classes, in_h, in_w).permute(0, 1, 3, 4, 2).contiguous()
        # Get outputs
        x = torch.sigmoid(prediction[..., 0])  # Center x
        y = torch.sigmoid(prediction[..., 1])  # Center y
        w = prediction[..., 2]  # Width
        h = prediction[..., 3]  # Height
        conf = torch.sigmoid(prediction[..., 4])  # Conf
        pred_cls = torch.sigmoid(prediction[..., 5:])  # Cls pred.

        # https://github.com/BobLiu20/YOLOv3_PyTorch/blob/c6b483743598b5f64d520d81e7e5f47ba936d4c9/nets/yolo_loss.py#L70-L92
        FloatTensor = torch.cuda.FloatTensor if x.is_cuda else torch.FloatTensor
        LongTensor = torch.cuda.LongTensor if x.is_cuda else torch.LongTensor
        # Calculate offsets for each grid
        grid_x = torch.linspace(0, in_w - 1, in_w).repeat(in_w, 1).repeat(
            bs * num_anchors, 1, 1).view(x.shape).type(FloatTensor)
        grid_y = torch.linspace(0, in_h - 1, in_h).repeat(in_h, 1).t().repeat(
            bs * num_anchors, 1, 1).view(y.shape).type(FloatTensor)
        # Calculate anchor w, h
        anchor_w = FloatTensor(scaled_anchors).index_select(1, LongTensor([0]))
        anchor_h = FloatTensor(scaled_anchors).index_select(1, LongTensor([1]))
        anchor_w = anchor_w.repeat(bs, 1).repeat(1, 1, in_h * in_w).view(w.shape)
        anchor_h = anchor_h.repeat(bs, 1).repeat(1, 1, in_h * in_w).view(h.shape)
        # Add offset and scale with anchors
        pred_boxes = FloatTensor(prediction[..., :4].shape)
        pred_boxes[..., 0] = x.data + grid_x
        pred_boxes[..., 1] = y.data + grid_y
        pred_boxes[..., 2] = torch.exp(w.data) * anchor_w
        pred_boxes[..., 3] = torch.exp(h.data) * anchor_h
        # Results
        _scale = torch.Tensor([stride_w, stride_h] * 2).type(FloatTensor)
        output = torch.cat((pred_boxes.view(bs, -1, 4) * _scale,
                            conf.view(bs, -1, 1), pred_cls.view(bs, -1, num_classes)), -1)
        return output

    def new_way(input, num_classes, num_anchors, anchors, stride_h, stride_w):
        raw_predictions = rearrange(input, ' b (anchor prediction) h w -> prediction b anchor h w', anchor=num_anchors)

        anchors = torch.FloatTensor(anchors).to(input.device)
        anchor_sizes = rearrange(anchors, 'anchor dim -> dim () anchor () ()')

        _, _, _, in_h, in_w = raw_predictions.shape
        grid_h = rearrange(torch.arange(in_h).float(), 'h -> () () h ()').to(input.device)
        grid_w = rearrange(torch.arange(in_w).float(), 'w -> () () () w').to(input.device)

        predicted_bboxes = torch.zeros_like(raw_predictions)
        predicted_bboxes[0] = (raw_predictions[0].sigmoid() + grid_h) * stride_h  # center y
        predicted_bboxes[1] = (raw_predictions[1].sigmoid() + grid_w) * stride_w  # center x
        predicted_bboxes[2:4] = (raw_predictions[2:4].exp()) * anchor_sizes  # bbox width and height
        predicted_bboxes[4] = raw_predictions[4].sigmoid()  # confidence
        predicted_bboxes[5:] = raw_predictions[5:].sigmoid()  # class predictions
        # only to match results of original code, not needed
        return rearrange(predicted_bboxes, 'prediction b anchor h w -> b anchor h w prediction')

    stride_h = 4
    stride_w = 4
    batch_size = 5
    num_classes = 12
    anchors = [[50, 100], [100, 50], [75, 75]]
    num_anchors = len(anchors)

    input = torch.randn([batch_size, num_anchors * (5 + num_classes), 1, 1])
    result1 = old_way(input=input, num_anchors=num_anchors, num_classes=num_classes,
                      stride_h=stride_h, stride_w=stride_w, anchors=anchors)
    result2 = new_way(input=input, num_anchors=num_anchors, num_classes=num_classes,
                      stride_h=stride_h, stride_w=stride_w, anchors=anchors)
    result1 = result1.reshape(result2.shape)
    assert torch.allclose(result1, result2)