#!/usr/bin/env python3
#
# Copyright (c) 2017-2018 Intel Corporation
#
# 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.


# To add new kernel please add a .cl file to kernels directory
# the database name will be the part of the file name up to first '.' character
# the trailing characters are a tag to allow multiple primitive implementations

from __future__ import print_function
from string import Template
import os
import sys
import functions_def

struct_tpl = Template('''// $struct_name function autogenerated file
#pragma once
#include "functions_base.hpp"
${use_includes}

namespace iclgpu { namespace functions {

struct $struct_name
{
$fields
};

} } // namespace iclgpu::functions
''')


def hpp_struct(struct):
    return struct_tpl.substitute(
        struct_name=struct['name'],
        fields='\n'.join(['    {} {};'.format(f['c_type'], f['name']) for f in struct['fields']]),
        use_includes='\n'.join(['#include "struct_{}.hpp'.format(use) for use in struct['uses']])
    )


function_tpl = Template(r'''// $func_name function autogenerated file
#pragma once
#include "functions_base.hpp"
${use_includes}

namespace iclgpu { namespace functions {

struct $func_name
{
    static constexpr size_t params_num = $params_num;
    static constexpr const char* name() { return "$func_name"; }

    struct params
    {
        $params_defs
    };

    struct score : public function_score<params_num>
    {
        $scores_defs

        score() : function_score()
            , $scores_init
        {}
    };

    using impl = function_impl_${impl_type}<$func_name>;
    using selector = selector_${selector_type}<$func_name>;
};

namespace implementations
{
$impls_defs
}

template<>
inline implementations_list<$func_name> get_implementations<$func_name>(const std::shared_ptr<context>& ctx)
{
    return {
        $impls_makers
    };
}

} } // namespace iclgpu::functions
// End of $func_name autogenerated
''')

impl_tpls = {
    'execute': Template('''
struct ${func_name}_${impl_name} : ${func_name}::impl
{
    using ${func_name}::impl::impl;
    const char* name() const override { return "${impl_name}"; }
    const char* full_name() const override { return "${func_name}_${impl_name}"; }
    bool accept(const ${func_name}::params& params, ${func_name}::score& score) override;
    event execute(const ${func_name}::params& params, const std::vector<event>& dep_events) override;
};
'''),

    'command': Template('''
struct ${func_name}_${impl_name} : ${func_name}::impl
{
    using ${func_name}::impl::impl;
    const char* name() const override { return "${impl_name}"; }
    const char* full_name() const override { return "${func_name}_${impl_name}"; }
    bool accept(const ${func_name}::params& params, ${func_name}::score& score) override;
    command_builder selected() override;
};
''')
}

impl_maker_tpl = Template('ctx->get<implementations::${func_name}_${impl_name}>(),')


def hpp_function(func):
    assert func['type'] == 'function'
    func_name = func['name']

    params_defs = []
    scores_defs = []
    scores_init = []
    params_num = len(func['params'])
    assert params_num > 0
    for i, p in enumerate(func['params']):
        params_defs.append('{} {};'.format(p['c_type'], p['name']))
        scores_defs.append('float& {};'.format(p['name']))
        scores_init.append('{}(_data[{}])'.format(p['name'], i))

    impls_defs = []
    impls_makers = []
    impl_type = func.get('impl_type', 'execute')
    impl_tpl = impl_tpls[impl_type]
    selector_type = 'accept'
    for impl_name in func['implementations']:
        impls_defs.append(impl_tpl.substitute(func_name=func_name, impl_name=impl_name))
        impls_makers.append(impl_maker_tpl.substitute(func_name=func_name, impl_name=impl_name))

    return function_tpl.substitute(
        func_name=func_name,
        params_num=params_num,
        use_includes='\n'.join(['#include "struct_{}.hpp"'.format(use) for use in func['uses']]),
        params_defs='\n        '.join(params_defs),
        scores_defs='\n        '.join(scores_defs),
        scores_init='\n            , '.join(scores_init),
        impls_defs=''.join(impls_defs),
        impls_makers='\n        '.join(impls_makers),
        impl_type=impl_type,
        selector_type=selector_type
    )


cl_impl_tpl = Template('''
__kernel void ${func_name}_${impl_name}(${params})
{
    // TODO Add implementation code here:
}
''')

cpp_impl_tpls = {
    'execute': Template('''#include "functions/${func_name}.hpp"

static const char* module_name = "${func_name}_${impl_name}";
static const char* kernel_name = "${func_name}_${impl_name}";

namespace iclgpu { namespace functions { namespace implementations {

bool ${func_name}_${impl_name}::accept(const ${func_name}::params& params, ${func_name}::score& score)
{
    // TODO Add implementation code here:
    return false;
}

event ${func_name}_${impl_name}::execute(const ${func_name}::params& params, const std::vector<event>& dep_events)
{
    // TODO Modify implementation code here:
    auto engine = context()->get_engine();
    auto kernel = engine->get_kernel(kernel_name, module_name);
    size_t buf_size = 1;
${set_kernel_args}

    auto gws = nd_range(1);
    auto lws = null_range;
    auto options = kernel_options(gws, lws);
    kernel->set_options(options);

    return kernel->submit(dep_events);
}

} } } // namespace iclgpu::functions::implementations
'''),

    'command': Template('''#include "functions/${func_name}.hpp"

static const char* module_name = "${func_name}_${impl_name}";
static const char* kernel_name = "${func_name}_${impl_name}";

namespace iclgpu { namespace functions { namespace implementations {

bool ${func_name}_${impl_name}::accept(const ${func_name}::params& params, ${func_name}::score& score)
{
    // TODO Add implementation code here:
    return false;
}

${func_name}_${impl_name}::command_builder ${func_name}_${impl_name}::selected()
{
  // TODO Modify implementation code here:
  auto engine = context()->get_engine();
  auto kernel = engine->get_kernel(kernel_name, module_name);
  return [=](const ${func_name}::params& params)
  {
    size_t buf_size = 1;
${set_kernel_args}

    auto gws = nd_range(1);
    auto lws = null_range;
    auto options = kernel_options(gws, lws);
    kernel->set_options(options);
    return kernel;
  };
}

} } } // namespace iclgpu::functions::implementations
''')
}

set_kernel_args_tpls = {
    'scalar': Template('''
    kernel->set_arg(${idx}, params.${param_name});'''),

    'blob': Template('''
    auto buf_${param_name} = engine->get_${param_inout}_buffer(params.${param_name}, buf_size);
    kernel->set_arg(${idx}, buf_${param_name});''')
}


def create_function_sources(hpp_dir, impl_dir, func):
    if not os.path.exists(hpp_dir):
        os.makedirs(hpp_dir)

    if not os.path.exists(impl_dir):
        os.makedirs(impl_dir)

    func_name = func['name']
    hpp_file_path = os.path.join(hpp_dir, func_name + '.hpp')
    print(hpp_file_path.replace('\\', '\\\\'), file=sys.stderr)
    hpp_file_content = hpp_function(func)
    create_hpp = True
    if os.path.exists(hpp_file_path):
        with open(hpp_file_path, 'r') as f:
            create_hpp = (hpp_file_content != f.read())

    if create_hpp:
        with open(hpp_file_path, 'w') as f:
            f.write(hpp_file_content)

    impl_type = func.get('impl_type', 'execute')
    cpp_impl_tpl = cpp_impl_tpls[impl_type]
    for impl_name in func['implementations']:
        impl_filename = func_name + '_' + impl_name

        params = []
        set_kernel_args = []
        read_buffers = []
        for i, param in enumerate(func['params']):
            params.append(('__global ' if param['blob'] else '') + param['cl_type'] + ' ' + param['name'])
            param_kind = 'blob' if param.get('blob', False) else 'scalar'
            set_kernel_args_tpl = set_kernel_args_tpls[param_kind]
            set_kernel_args.append(set_kernel_args_tpl.substitute(
                param_name=param['name'],
                param_inout=param['inout_type'],
                idx=i
            ))

        cpp_file_path = os.path.join(impl_dir, impl_filename + '.cpp')
        cl_file_path = os.path.join(impl_dir, impl_filename + '.cl')

        # Autogenerate impl files if .cpp does not exist
        if not os.path.exists(cpp_file_path):
            # Write func_impl.cpp file
            with open(cpp_file_path, 'w') as f:
                f.write(cpp_impl_tpl.substitute(
                    func_name=func_name,
                    impl_name=impl_name,
                    set_kernel_args=''.join(set_kernel_args),
                    read_buffers=''.join(read_buffers)
                ))
            # Write func_impl.cl file
            with open(cl_file_path, 'w') as f:
                f.write(cl_impl_tpl.substitute(
                    func_name=func_name,
                    impl_name=impl_name,
                    params=', '.join(params)
                ))

        # Escape '\' for cmake compatibility
        print(cpp_file_path.replace('\\', '\\\\'))

        # Print func_impl.cl file name if it exists
        if os.path.exists(cl_file_path):
            print(cl_file_path.replace('\\', '\\\\'))



def create_struct_sources(hpp_dir, impl_dir, struct):
    os.makedirs(hpp_dir, exist_ok=True)
    hpp_file_path = os.path.join(hpp_dir, 'struct_' + struct['name'] + '.hpp')
    print(hpp_file_path.replace('\\', '\\\\'), file=sys.stderr)
    hpp_file_content = hpp_struct(struct)
    create_hpp = True
    if os.path.exists(hpp_file_path):
        with open(hpp_file_path, 'r') as f:
            create_hpp = (hpp_file_content != f.read())

    if create_hpp:
        with open(hpp_file_path, 'w') as f:
            f.write(hpp_file_content)


functions_def_file = sys.argv[1] if len(sys.argv) > 1 else os.path.join(os.getcwd(), 'functions.fdef')
impls_root = sys.argv[2] if len(sys.argv) > 2 else os.path.dirname(functions_def_file)
hpp_root = sys.argv[3] if len(sys.argv) > 3 else impls_root

try:
    with open(functions_def_file) as f:
        items = [x for x in functions_def.FunctionsDef(f, functions_def_file)]
except SyntaxError as err:
    sys.exit(err)

for item in items:
    if item['type'] == 'function':
        create_function_sources(hpp_root, os.path.join(impls_root, item['name']), item)
    elif item['type'] == 'struct':
        create_struct_sources(hpp_root, os.path.join(impls_root, item['name']), item)
