# Copyright (c) 2020 Graphcore Ltd. All rights reserved
import datetime
import sys

import re

from popgen import registry, transform


# emit_handlers(namespace, aten, handlers, f=sys.stdout)
#
# Emits the C++ handlers for one operator.
# Parameters:
#   namespace - namespace the operator is in
#   aten - name of the operator
#   handlers - list of handlers. must differ in arity.
#   f - output stream
def emit_handlers(namespace, aten, handlers, f=sys.stdout):
    values = dict()
    opname = get_camel_case_op_name(aten)
    emit_arity_check = len(handlers) > 1

    decl = "torch::jit::Node *" + opname +  "Handler(" + \
           "torch::jit::Graph *graph, " + "torch::jit::Node *node) {"
    if len(decl) <= 80:
        f.write(decl + "\n")
    else:
        decl = "torch::jit::Node *" + opname + "Handler("
        f.write(decl + "torch::jit::Graph *graph,\n")
        f.write(" " * len(decl))
        f.write("torch::jit::Node *node) {\n")

    arities = set()
    for handler in handlers:
        assert handler.graph_arity not in arities, \
               aten + " has multiple handlers with the same arity"
        arities.add(handler.graph_arity)
        values.clear()
        handler = transform.generate_complex_ops(handler)
        handler = transform.value_numbering(handler)
        handler = transform.generate_typed_constants(handler)
        handler.annotate("// " + handler.render())

        if emit_arity_check:
            f.write("  if (node->inputs().size() == " +
                    str(handler.graph_arity) + ") {\n")
            handler.emit(values, 0, "    ", f, True)
            f.write("  }\n")
        else:
            handler.emit(values, 0, "  ", f, True)

    if emit_arity_check:
        arity_list = sorted(list(arities))
        expect_str = "Expecting " + str(arity_list[0])

        for i in range(1, len(arity_list) - 1):
            expect_str += ', ' + str(arity_list[i])

        if len(arity_list) > 1:
            expect_str += ' or ' + str(arity_list[-1])

        if len(arity_list) > 1 or arity_list[0] > 1:
            expect_str += " operands, "
        else:
            expect_str += " operand, "

        f.write('\n  std::stringstream errmsg;\n')
        f.write('  errmsg << "Incorrect number of arguments for operator ";\n')
        f.write('  errmsg << "' + namespace + '::' + aten + '. ";\n')
        f.write('  errmsg << "' + expect_str + '";\n')
        f.write(
            '  errmsg << "got " << node->inputs().size() << " operand(s).";\n')
        f.write("  ERROR(&errmsg);\n")
        f.write("  return nullptr;\n")

    f.write("}\n\n")


# generate(script, namespace, filename, global_symbols)
#
# Generate a file containg C++ implementation of handlers
# Parameters:
#   script - name of top-level script
#   namespace - the namespace the operators are in
#   filename - the output fil
#   global_symbols - dictionary of globals from top-level
def generate(script, namespace, filename, global_symbols):
    f = open(filename, 'w')

    now = datetime.datetime.now()
    f.write('// DO NOT EDIT! Generated by ' + script + '\n')
    f.write('// Copyright (c) ' + str(now.year) +
            ' Graphcore Ltd. All rights reserved.\n\n')

    f.write('#include "../PoptorchStaticInit.hpp"\n')
    f.write('#include "../PoptorchSymbols.hpp"\n')
    f.write('#include "PopartCanonicalizationUtils.hpp"\n')
    f.write('#include "poptorch/OpBuilder.hpp"\n')
    f.write('#include "poptorch/Utils.hpp"\n')
    f.write('#include "poptorch_logging/Error.hpp"\n')
    f.write('#include "poptorch_logging/Logging.hpp"\n')

    f.write("\nnamespace poptorch {\n")
    f.write("\nnamespace {\n\n")

    registry.add_implicit_handlers(global_symbols)
    for (aten, handler) in sorted(registry.handlers.items()):
        emit_handlers(namespace, aten, handler, f)

    f.write("} // namespace\n")

    f.write("\n__attribute__((constructor(HANDLER_INIT_PRIORITY))) ")
    f.write("static void registration() {\n")

    for (source, _) in registry.forwardings.items():
        transform.validate_forwarding(source)

    to_register = sorted(
        list(registry.handlers.keys()) + list(registry.forwardings.keys()))
    for aten in to_register:
        opname = get_camel_case_op_name(registry.forwardings.get(aten, aten))
        reg_handler_line = ("  registerHandler(" + namespace + "::" + aten +
                            ", " + opname + "Handler);\n")
        if len(reg_handler_line) > 81:
            reg_handler_line = reg_handler_line.replace(
                ", ", ",\n                  ")
        f.write(reg_handler_line)
    f.write("}\n\n")

    f.write("} // namespace poptorch\n")
    f.close()

    registry.clear()


# get_op_name(aten)
#
# Returns the name of the C++ handler function for an operator
# Parameters:
#   aten - the name of the operator
def get_op_name(aten):
    opname = aten.split(':')[-1]
    return opname


def snake_to_camel_case(snake_case):
    rx = re.compile(r"_[A-z]")
    return re.sub(rx, lambda s: s.group(0).upper()[1], snake_case)


def get_camel_case_op_name(aten):
    rx = re.compile(r"_$")
    return re.sub(rx, "InPlace", snake_to_camel_case(get_op_name(aten)))