# Copyright 2015 Leon Sixt
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import os
import tempfile
from unittest.mock import Mock, MagicMock
import math
from keras.layers.convolutional import Convolution2D
from keras.models import Sequential
from seya.data_utils import floatX
import theano

from beras.layers.attention import RotationTransformer
import numpy as np


def test_rotates_images():
    bs = 3
    img = theano.shared(np.zeros((bs, 1, 8, 8), dtype=floatX))

    angle = np.asarray([0, math.pi / 2, math.pi], dtype=floatX)
    angle = theano.shared(angle)

    locnet = Mock()
    locnet.layers = [Mock()]
    locnet.layers[0].input = img
    locnet.get_output = Mock(return_value=angle)
    assert locnet.get_output() == angle
    rot_layer = RotationTransformer(locnet, return_theta=True)
    rot_layer.input = img
    theta = rot_layer.get_output()
    np_theta = theta.eval()
    np.testing.assert_almost_equal(
        np_theta,
        np.asarray([
            [[1, 0, 0],
             [0, 1, 0]],

            [[0, -1, 0],
             [1, 0, 0]],
            [[-1, 0, 0],
             [0, -1, 0]],
        ]),
        verbose=True
    )


def get_net():
    rotnet = Sequential()
    rotnet.add(Convolution2D(1, 2, 2, input_shape=(1, 8, 8)))
    rotnet.add(Convolution2D(1, 2, 2))

    net = Sequential()
    net.add(RotationTransformer(rotnet, return_theta=True,
                                input_shape=(1, 8, 8)))
    print(net.layers[0].output_shape)
    net.add(Convolution2D(1, 2, 2))
    return net, rotnet


def test_save_and_loads_weights():
    dirname = tempfile.mkdtemp()
    save_path = os.path.join(dirname, "test_net.hdf5")
    net_save, rotnet_save = get_net()
    net_load, rotnet_load = get_net()
    net_save.save_weights(save_path)
    net_load.load_weights(save_path)
    for s, l in zip(net_save.params, net_load.params):
        assert (s.get_value() == l.get_value()).all()
    for s, l in zip(rotnet_save.params, rotnet_load.params):
        assert (s.get_value() == l.get_value()).all()

@@ -575,16 +575,30 @@
 port os%0A
+import shutil%0A
 import t
@@ -712,16 +712,70 @@
 ution2D%0A
+from keras.layers.core import Dense, Flatten, Dropout%0A
 from ker
@@ -1817,24 +1817,53 @@
 (1, 8, 8)))%0A
+    rotnet.add(Dropout(0.5))%0A
     rotnet.a
@@ -1893,243 +1893,413 @@
 2))%0A
-%0A    net = Sequential()%0A    net.add(RotationTransformer(rotnet, return_theta=True,%0A                                input_shape=(1, 8, 8)))%0A    print(net.layers%5B0%5D.output_shape)%0A    net.add(Convolution2D(1, 2, 2))%0A    return net, rotnet
+    rotnet.add(Dropout(0.5))%0A    rotnet.add(Flatten())%0A    rotnet.add(Dense(1))%0A%0A    net = Sequential()%0A    net.add(RotationTransformer(rotnet, input_shape=(1, 8, 8)))%0A    net.add(Convolution2D(1, 2, 2))%0A    net.add(Flatten())%0A    net.add(Dense(10))%0A    return net, rotnet%0A%0A%0Adef test_rotnet_predict():%0A    net, _ = get_net()%0A    net.compile(%22adam%22, %22mse%22)%0A    net.predict(np.random.sample((128, 1, 8, 8)))
 %0A%0A%0Ad
@@ -2365,16 +2365,29 @@
 dtemp()%0A
+    try:%0A    
     save
@@ -2433,16 +2433,20 @@
 .hdf5%22)%0A
+    
     net_
@@ -2467,32 +2467,36 @@
 ave = get_net()%0A
+    
     net_load, ro
@@ -2513,24 +2513,28 @@
 = get_net()%0A
+    
     net_save
@@ -2558,24 +2558,28 @@
 e_path)%0A    
+   
+ 
 net_load.loa
@@ -2595,24 +2595,28 @@
 (save_path)%0A
+    
     for s, l
@@ -2658,32 +2658,36 @@
 arams):%0A        
+ 
+   
 assert (s.get_va
@@ -2712,24 +2712,28 @@
 ue()).all()%0A
+    
     for s, l
@@ -2781,32 +2781,36 @@
 arams):%0A        
+    
 assert (s.get_va
@@ -2819,28 +2819,71 @@
 e() == l.get_value()).all()%0A
+    finally:%0A        shutil.rmtree(dirname)

refactor rotnet tests