# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you 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.
"""Test task extraction for autotvm"""
import tvm.relay.testing
from tvm import relay
from tvm import autotvm

def get_network(name, batch_size):
    """Get the symbol definition and random weight of a network"""
    input_shape = (batch_size, 3, 224, 224)

    if name == 'resnet-18':
        mod, params = relay.testing.resnet.get_workload(num_layers=18, batch_size=batch_size)
    elif name == 'mobilenet':
        mod, params = relay.testing.mobilenet.get_workload(batch_size=batch_size)
    elif name == 'dcgan':
        mod, params = relay.testing.dcgan.get_workload(batch_size=batch_size)
        input_shape = (batch_size, 100)
    else:
        raise ValueError("Unsupported network: " + name)

    return mod, params, input_shape

def test_task_extraction():
    target = 'llvm'
    mod_list = []
    params_list = []

    mod, params, _ = get_network('resnet-18', batch_size=1)
    tasks = autotvm.task.extract_from_program(mod["main"], target=target,
                                              params=params,
                                              ops=(relay.op.nn.conv2d,))
    assert len(tasks) == 12

    mod, params, _ = get_network('resnet-18', batch_size=1)
    tasks = autotvm.task.extract_from_program(mod["main"], target=target,
                                              params=params,
                                              ops=(relay.op.nn.dense,))
    assert len(tasks) == 1

    mod, params, _ = get_network('resnet-18', batch_size=1)
    mod_list.append(mod)
    params_list.append(params)
    tasks = autotvm.task.extract_from_program(mod["main"], target=target,
                                              params=params,
                                              ops=(relay.op.nn.conv2d, relay.op.nn.dense))
    assert len(tasks) == 13

    mod, params, _ = get_network('mobilenet', batch_size=1)
    mod_list.append(mod)
    params_list.append(params)
    tasks = autotvm.task.extract_from_program(mod["main"], target=target,
                                              params=params,
                                              ops=(relay.op.nn.conv2d, relay.op.nn.dense))
    assert len(tasks) == 20

    mod, params, _ = get_network('dcgan', batch_size=1)
    tasks = autotvm.task.extract_from_program(mod["main"], target=target,
                                              params=params,
                                              ops=(relay.op.nn.conv2d_transpose,))
    assert len(tasks) == 4

    tasks = autotvm.task.extract_from_multiple_program([m['main'] for m in mod_list], params_list,
                                                       target=target,
                                                       ops=(relay.op.nn.conv2d,))
    assert len(tasks) == 31

def test_template_key_provided():
    """test task extraction using non-'direct' template_key"""
    target = 'llvm'

    import topi
    template_keys = {
        # topi.nn.conv2d - is left blank to test fallback logic
        topi.nn.dense: 'direct_nopack',
        topi.nn.depthwise_conv2d_nchw: 'direct',
    }

    mod, params, _ = get_network('mobilenet', batch_size=1)
    tasks = autotvm.task.extract_from_program(mod['main'], target=target,
                                              params=params,
                                              ops=(relay.op.nn.conv2d, relay.op.nn.dense),
                                              template_keys=template_keys)
    for task in tasks:
        if 'dense' in task.name:
            assert task.config_space.template_key == 'direct_nopack'
        else:
            assert task.config_space.template_key == 'direct'

def test_template_key_empty():
    """test task extraction using empty template_key"""
    target = 'llvm'
    mod, params, _ = get_network('mobilenet', batch_size=1)
    tasks = autotvm.task.extract_from_program(mod['main'], target=target,
                                              params=params,
                                              ops=(relay.op.nn.conv2d, relay.op.nn.dense),
                                              template_keys=None)
    for task in tasks:
        assert task.config_space.template_key == 'direct'

def test_template_key_default():
    """test task extraction without template_key"""
    target = 'llvm'
    mod, params, _ = get_network('mobilenet', batch_size=1)
    tasks = autotvm.task.extract_from_program(mod['main'], target=target,
                                              params=params,
                                              ops=(relay.op.nn.conv2d, relay.op.nn.dense))
    for task in tasks:
        assert task.config_space.template_key == 'direct'

if __name__ == '__main__':
    test_task_extraction()
    test_template_key_provided()
    test_template_key_empty()
    test_template_key_default()
