# Copyright (c) 2016-present, Facebook, Inc.
#
# 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.
##############################################################################

# TODO(jiayq): as more and more tests are moving to hypothesis test, we
# can gradually remove this test script. DO NOT ADD MORE TESTS TO THIS
# FILE.

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from __future__ import unicode_literals
import numpy as np
from caffe2.python import (
    brew,
    core,
    device_checker,
    gradient_checker,
    model_helper,
    test_util,
    workspace,
)
from caffe2.python.gradient_checker import NetGradientChecker
from caffe2.python.net_builder import ops, NetBuilder
from caffe2.proto import caffe2_pb2

import unittest


if workspace.has_gpu_support and workspace.NumCudaDevices() > 0:
    gpu_device_option = caffe2_pb2.DeviceOption()
    gpu_device_option.device_type = caffe2_pb2.CUDA
    cpu_device_option = caffe2_pb2.DeviceOption()
    gpu_device_checker = device_checker.DeviceChecker(
        0.01, [gpu_device_option]
    )
    device_checker = device_checker.DeviceChecker(
        0.01, [gpu_device_option, cpu_device_option]
    )
    gpu_gradient_checkers = [
        gradient_checker.GradientChecker(
            0.005, 0.05, gpu_device_option, "gpu_checker_ws"
        ),
    ]
    gradient_checkers = [
        gradient_checker.GradientChecker(
            0.005, 0.05, gpu_device_option, "gpu_checker_ws"
        ),
        gradient_checker.GradientChecker(
            0.01, 0.05, cpu_device_option, "cpu_checker_ws"
        ),
    ]
else:
    cpu_device_option = caffe2_pb2.DeviceOption()
    gpu_device_option = None
    gpu_device_checker = device_checker.DeviceChecker(
        0.01, []
    )
    device_checker = device_checker.DeviceChecker(0.01, [cpu_device_option])

    gradient_checkers = [
        gradient_checker.GradientChecker(
            0.01, 0.05, cpu_device_option, "cpu_checker_ws"
        )
    ]
    gpu_gradient_checkers = []


class TestLRN(test_util.TestCase):

    def setUp(self):
        self.test_configs = [(6, 10), (3, 13), ]