import re
from code_template import CodeTemplate

import sys
if sys.version_info[0] == 3:
    string_type = str
else:
    string_type = basestring

# temporary things we cannot handle
EXCLUDE_PATTERN = "bernoulli.*|normal.*|exponential.*|random.*|arange.*"
# what has to be done to add a Operation ...
# 1. add virtual dispatch declaration to Type.h and default impl to Type.cpp
TYPE_METHOD_DECLARATION = CodeTemplate("""\
virtual ${return_type} ${method_prefix}${api_name}(${formals}) ;
""")
TYPE_METHOD_DEFINITION = CodeTemplate("""\
${return_type} Type::${method_prefix}${api_name}(${formals}) {
    throw std::runtime_error(std::string("${api_name} is not implemented for type ") + toString());
}
""")
# 2. add virtual override to TypeDerived.h
TYPE_DERIVED_DECLARATION = CodeTemplate("""\
virtual ${return_type} ${method_prefix}${api_name}(${formals}) override;
""")
# 3. add override definition to TypeDerived.cpp
TYPE_DERIVED_DEFINITION = CodeTemplate("""\
${return_type} ${Type}::${method_prefix}${api_name}(${formals}) {
    ${type_definition_body}
}
""")
# 4. add non-virtual declaration to Tensor.h
TENSOR_METHOD_DECLARATION = CodeTemplate("""\
${return_type} ${api_name}(${method_formals})${const_mark};
""")
# 5. add non-virtual declaration to Tensor.cpp
TENSOR_METHOD_DEFINITION = CodeTemplate("""\
inline ${return_type} Tensor::${api_name}(${method_formals})${const_mark} {
    return type().${method_prefix}${api_name}(${method_actuals});
}
""")
# 6. add a method declaration in Functions.h
FUNCTION_DECLARATION = CodeTemplate("""\
static inline ${return_type} ${api_name}(${formals});
""")
# 7. add a method definition in Functions.cpp
FUNCTION_DEFINITION = CodeTemplate("""\
static inline ${return_type} ${api_name}(${formals}) {
    return ${inferred_type}.${api_name}(${actuals});
}
""")

ZERO_DIM_CHECK = CodeTemplate("""\
if(${check_name}.dim() == 0) {
    return ${method_prefix}${api_name}(${zero_dim_actuals});
}""")

SCALAR_EXPAND = CodeTemplate("""\
Tensor ${name}__;
if(${name}_->isScalar()) {
    ${name}__ = ${name}.expand(${other}.sizes());
    ${name}_ = static_cast<${Tensor}*>(${name}__.pImpl);
}
""")


class NYIError(Exception):
    """Indicates we don't support this declaration yet"""

    def __init__(self, reason):
        self.reason = reason


TYPE_FORMAL_GENERIC = {
    'THTensor*': 'Tensor &',
    'THBoolTensor*': 'Tensor &',
    'THIndexTensor*': 'Tensor &',
    'THIntegerTensor*': 'Tensor &',
    'THStorage*': 'Storage &',
    'THGenerator*': 'Generator &',
    'THSize*': 'IntList',
    'THStride*': 'IntList',
    'accreal': 'Scalar',
    'real': 'Scalar',
    'long': 'int64_t',
}

DYNAMIC_TYPE = {
    'THTensor*': 'Tensor',
    'THBoolTensor*': 'BoolTensor',
    'THIndexTensor*': 'IndexTensor',
    'THIntegerTensor*': 'IntegerTensor',
    'THStorage*': 'Storage',
    'THGenerator*': 'Generator',
    'THSize*': 'IntList',
    'THStride*': 'IntList',
    'accreal': 'accreal',
    'real': 'real',
    'long': 'int64_t',
}

TYPE_RETURN = {
    'THTensor*': 'Tensor',
    'THIndexTensor*': 'Tensor',
    'THBoolTensor*': 'Tensor',
    'THIntegerTensor*': 'Tensor',
    'real': 'Scalar',
    'accreal': 'Scalar',
    'long': 'int64_t',
}
CHECKED_CAST = {
    'THTensor*': CodeTemplate('checked_cast<${Tensor}>(${arg_name}.pImpl,"${arg_name}",${arg_pos}, ${null_okay})'),
    'THBoolTensor*':
        CodeTemplate('checked_cast<${Backend}ByteTensor>(${arg_name}.pImpl,"${arg_name}",${arg_pos}, ${null_okay})'),
    'THIndexTensor*':
        CodeTemplate('checked_cast<${Backend}LongTensor>(${arg_name}.pImpl,"${arg_name}",${arg_pos}, ${null_okay})'),
    'THIntegerTensor*':
        CodeTemplate('checked_cast<${Backend}IntTensor>(${arg_name}.pImpl,"${arg_name}",${arg_pos}, ${null_okay})'),
    'THStorage*': CodeTemplate('checked_cast<${Storage}>(&${arg_name},"${arg_name}",${arg_pos}, false)'),
    'THGenerator*': CodeTemplate('check_generator(&${arg_name})'),
    'THSize*': CodeTemplate('THLongStorageView::make(${arg_name},true)'),
    'THStride*': CodeTemplate('THLongStorageView::make(${arg_name},true)'),
    'real': CodeTemplate('${arg_name}.to${ScalarName}()'),
    'accreal': CodeTemplate('${arg_name}.to${AccScalarName}()'),
    'TensorList': CodeTemplate('tensor_list_checked_cast<${Tensor}, Tensor, '
                               '${THTensor}>(${arg_name},"${arg_name}",${arg_pos})'),
}

CHECKED_USE = {
    'THTensor*': '{}_->tensor',
    'THIndexTensor*': '{}_->tensor',
    'THBoolTensor*': '{}_->tensor',
    'THIntegerTensor*': '{}_->tensor',
    'THStorage*': '{}_->storage',
    'THGenerator*': '{}_->generator',
    'TensorList': "{0}_.data(), {0}_.size()",
}

CHECKED_USE_NULLABLE = CodeTemplate('${arg_name}_ ? ${usage} : NULL')

ALLOC_WRAP = {
    'THTensor*': 'new ${Tensor}(context)',
    'THBoolTensor*': 'new ${Backend}ByteTensor(context)',
    'THIndexTensor*': 'new ${Backend}LongTensor(context)',
    'THIntegerTensor*': 'new ${Backend}IntTensor(context)',
}

CONSTANT_REPLACEMENTS = [
    ('AS_REAL', '${AS_REAL}'),
    ('THPDefaultGenerator->cdata',
     'dynamic_cast<${Generator}&>(context->defaultGenerator(backend())).generator'),
    ('__storage_size.get\\(\\)',
     'THLongStorageView::make(static_cast<int64_t>(storage.size()))'),
    ('__last_dim', 'self.ndimension()-1'),
]


class nested_dict(object):
    def __init__(self, base, parent):
        self.base, self.parent = base, parent

    def __getitem__(self, x):
        r = self.base.get(x)
        if r is not None:
            return r
        return self.parent[x]


def is_real_argument_to_wrapper(argument):
    return not argument.get('output', False) and\
        argument['type'] != 'CONSTANT' and\
        argument['type'] != 'argument'


def is_mutable_formal_argument(argument, option):
    return argument.get('output') or option['inplace'] and argument['name'] == 'self'


def to_return_type(arg, option):
    t = arg['type']
    rt = TYPE_RETURN.get(t, t)
    if rt == 'Tensor' and not arg.get('allocate'):
        rt = rt + ' &'
        if not is_mutable_formal_argument(arg, option):
            rt = 'const ' + rt
    return {
        'type': rt,
        'dynamic_type': DYNAMIC_TYPE.get(arg['type'], arg['type']),
    }


def create_generic(top_env, declarations):

    # change from THTensor* to Tensor & so we get how it will appear
    # in the aten argument list...
    def translate_formal(argument, option):
        type_str = TYPE_FORMAL_GENERIC.get(argument['type'], argument['type'])
        if type_str == 'Tensor &' and not is_mutable_formal_argument(argument, option):
            type_str = 'const ' + type_str
        translated = {
            'name': argument['name'],
            'type': type_str,
            'dynamic_type': DYNAMIC_TYPE.get(argument['type'], argument['type']),
        }
        if argument.get('output'):
            translated['output'] = True
        return translated

    def get_formals(option):
        seen = set()
        result = []

        def insert(argument):
            if argument['name'] not in seen:
                seen.add(argument['name'])
                result.append(argument)
        for argument in option['arguments']:
            if argument['type'] == 'THSTensor*':
                raise NYIError("Sparse Tensor")
            if is_real_argument_to_wrapper(argument):
                insert(argument)
        for argument in option['arguments']:
            if argument.get('output') and not argument.get('allocate', False):
                insert(argument)

        return [translate_formal(argument, option) for argument in result]

    def get_return_types(option):
        ret = option['return']
        if ret['kind'] == 'arguments':
            argument_indices = ret['arguments']
            if len(argument_indices) == 1:
                the_arg = option['arguments'][argument_indices[0]]
                return [to_return_type(the_arg, option)]
            else:
                return [to_return_type(option['arguments'][idx], option)
                        for idx in argument_indices]
        elif ret['kind'] == 'type':
            return [{
                'type': TYPE_RETURN.get(ret['type'], ret['type']),
                'dynamic_type': DYNAMIC_TYPE.get(ret['type'], ret['type']),
            }]
        else:
            raise Exception("format_return_type")

    def format_return_type(return_types):
        if len(return_types) == 1:
            return return_types[0]['type']
        return "std::tuple<{}>".format(','.join(r['type'] for r in return_types))
        return return_types

    def find_first_tensor(formals):
        for formal in formals:
            if 'Tensor' == formal['dynamic_type'] or 'TensorList' == formal['dynamic_type']:
                return formal['name']
        return None

    def format_formal(f):
        return '{} {}'.format(f['type'], f['name'])

    def process_option(option, output_options):
        option['inplace'] = re.search(
            '(^__i|[^_]_$)', option['api_name']) is not None

        if re.match(EXCLUDE_PATTERN, option['name']):
            print("Excluding {}".format(option['name']))
            raise NYIError("NYI")
        # print(yaml.dump(option))
        formals = get_formals(option)
        option['formals_list'] = formals
        option['formals'] = [format_formal(f) for f in formals]
        option['returns'] = get_return_types(option)
        option['actuals'] = [f['name'] for f in formals]

        option['method_formals'] = [format_formal(f) for f in formals
                                    if f['name'] != 'self']
        option['method_actuals'] = [
            f['name'] if f['name'] != 'self' else '*this' for f in formals]
        option['return_type'] = format_return_type(option['returns'])

        option['const_mark'] = '' if option['inplace'] else ' const'

        is_method = 'method' in option['variants']
        is_function = 'function' in option['variants']

        # method-only things are prefixed with m_ in Type so that
        # another function-only variant can exist without the name colliding
        option['method_prefix'] = 'm_' if is_method and not is_function else ''
        env = nested_dict(option, top_env)
        top_env['type_method_declarations'].append(
            TYPE_METHOD_DECLARATION.substitute(env))
        top_env['type_method_definitions'].append(
            TYPE_METHOD_DEFINITION.substitute(env))

        if is_method:
            top_env['tensor_method_declarations'].append(
                TENSOR_METHOD_DECLARATION.substitute(env))
            top_env['tensor_method_definitions'].append(
                TENSOR_METHOD_DEFINITION.substitute(env))
            output_options.append({
                'name': option['name'],
                'arguments': [f for f in formals if f['name'] != 'self'],
                'method_of': 'Tensor',
                'returns': option['returns'],
                'inplace': option['inplace'],
            })

        if is_function:
            first_tensor = find_first_tensor(formals)
            output_option = {
                'name': option['name'],
                'arguments': formals,
                'returns': option['returns'],
                'inplace': option['inplace'],
            }
            if first_tensor is not None:
                option['inferred_type'] = 'infer_type({})'.format(first_tensor)
                top_env['function_declarations'].append(
                    FUNCTION_DECLARATION.substitute(env))
                top_env['function_definitions'].append(
                    FUNCTION_DEFINITION.substitute(env))
            else:
                output_option['method_of'] = 'Type'
            output_options.append(output_option)

    output_declarations = []
    for declaration in declarations:
        output_options = []
        for option in declaration['options']:
            try:
                process_option(option, output_options)
            except NYIError:
                option['skip'] = True
        if len(output_options) > 0:
            output_declarations.append({
                'name': output_options[0]['name'],
                'options': output_options,
            })
    return output_declarations


def create_derived(backend_type_env, declarations):
    type_object_declarations = []
    type_object_definitions = []

    def requires_checked_cast(argument):
        return argument['type'] in CHECKED_CAST

    def nullable_argument(argument):
        return (argument['type'] == 'THTensor*' and
                argument.get('default', '') == 'nullptr')

    def bool_option_is_string(argument):
        return 'if_true' in argument and isinstance(argument['if_true'], string_type)

    def get_argument(argument, option):
        if requires_checked_cast(argument):
            checked_use = CHECKED_USE.get(argument['type'], '{}_').format(argument['name'])
            if nullable_argument(argument):
                checked_use = CHECKED_USE_NULLABLE.substitute(
                    env={}, arg_name=argument['name'], usage=checked_use)
            return checked_use
        elif argument['type'] == 'bool' and 'if_true' in argument:
            if bool_option_is_string(argument):
                tpl = '({}) ? "{}" : "{}"'
            else:
                tpl = '({}) ? {} : {}'
            return tpl.format(argument['name'],
                              argument['if_true'], argument['if_false'])
        elif argument['type'] == "CONSTANT":
            if bool_option_is_string(argument):  # this is a bool that is actually a string...
                return '"{}"'.format(argument['name'])
            v = str(argument['name'])
            for pattern, replacement in CONSTANT_REPLACEMENTS:
                v = re.sub(pattern, replacement, v)
            return CodeTemplate(v).substitute(backend_type_env)
        # e.g. argument 0, i.e. repeat the 0th argument in this position...
        elif argument['type'] == 'argument':
            index = int(argument['name'])
            return get_argument(option['arguments'][index], option)
        else:
            return argument['name']

    def drop_argument(argument, option):
        return backend_type_env['Backend'] == 'CUDA' and (
            (option['mode'] == 'TH' and argument['type'] == 'THGenerator*') or
            argument['name'] == 'THPDefaultGenerator->cdata')

    def get_arguments(option):
        return [get_argument(argument, option)
                for argument in option['arguments'] if not drop_argument(argument, option)]

    def is_actual_return_long(ret):
        if ret['type'] == 'long':
            return True
        if ret['type'] == 'real':
            return backend_type_env['ScalarName'] == 'Long'
        if ret['type'] == 'accreal':
            return backend_type_env['AccScalarName'] == 'Long'
        return False

    def handle_zero_dim(env, option):
        if 'zero_dim_dispatch_when_scalar' not in option:
            return []
        check_name = option['zero_dim_dispatch_when_scalar']
        zero_dim_actuals = [arg['name']
                            if arg['name'] != check_name else "Scalar({})".format(arg['name'])
                            for arg in option['formals_list']]
        return [ZERO_DIM_CHECK.substitute(env, check_name=check_name, zero_dim_actuals=zero_dim_actuals)]

    def emit_body(env, option):
        body = []
        body += handle_zero_dim(env, option)
        # arguments are potentially duplicated because of one argument
        # referencing another
        seen_names = set()
        count = 0
        is_cuda = backend_type_env['Backend'] == 'CUDA'

        # scalar_check is the heuristic conditions when a result may be a scalar_check
        # if there is a THSize* argument, then its dimensions are used to determine scalar.
        # otherwise, it is true if all the input tensors are scalars,
        scalar_check_is_from_size = False
        scalar_check = None
        for arg in option['arguments']:
            if is_real_argument_to_wrapper(arg):
                count += 1
            if arg['type'] == 'THSize*':
                scalar_check_is_from_size = True
                scalar_check = '{}.size() == 0'.format(arg['name'])
            # only generated checked casts the first time we see it
            if not arg['name'] in seen_names and requires_checked_cast(arg):
                seen_names.add(arg['name'])

                # make a new allocation of TensorImpl, then wrap a Tensor around it.
                if arg.get('allocate', False):
                    allocation = CodeTemplate(
                        ALLOC_WRAP[arg['type']]).substitute(env)
                    body.append('auto {}_ = {};'.format(
                        arg['name'], allocation))
                    body.append('auto {} = Tensor({}_,false);'.format(
                        arg['name'], arg['name']))
                # extract the TensorImpl from an existing tensor (or Storage, etc.)
                else:
                    # special case where we allow undefined Tensors, and thus
                    # the checked cast succeeds even if the Tensor is not
                    # defined
                    null_okay = 'true' if nullable_argument(arg) else 'false'

                    check_cast = CHECKED_CAST[arg['type']].substitute(
                        env, arg_name=arg['name'], arg_pos=count,
                        null_okay=null_okay)
                    body.append("auto {}_ = {};".format(
                        arg['name'], check_cast))
                if drop_argument(arg, option):
                    body.append("(void) {}_; //silence unused warning".format(arg['name']))
                # resize tensors for special ops that require it
                if 'resize' in arg:
                    resize = arg['resize']
                    if isinstance(resize, str):
                        body.append("{}.resize_({}.sizes());".format(
                            arg['name'], resize))
                    else:
                        dims = ['{}.size({})'.format(name, dim)
                                for name, dim in resize]
                        body.append("{}.resize_({{ {} }});".format(
                            arg['name'], ','.join(dims)))
                # also special handling where we zero some outputs.
                if arg.get('cpu_zero', False) and not is_cuda:
                    body.append("{}.zero_();".format(arg['name']))

                # handle scalars that occur on LHS of things like a - b
                if 'broadcast' in arg and 'inplace' not in arg['broadcast']:
                    other = arg['broadcast'].split(' ')[0].split(',')[0]
                    body.append(SCALAR_EXPAND.substitute(env,
                                                         name=arg['name'],
                                                         other=other))

                # dim() == 0 of all input tensors is and'd to form
                # the test for whether the output is also a scalar
                if (not arg.get('output') and 'Tensor' in arg['type'] and
                        'TensorList' not in arg['type'] and not scalar_check_is_from_size):
                    check = '{}.dim() == 0'.format(arg['name'])
                    scalar_check = (check if scalar_check is None
                                    else scalar_check + ' && ' + check)

        option['derived_actuals'] = get_arguments(option)
        is_nn = option['mode'] == 'NN'
        if is_cuda or is_nn:
            option['derived_actuals'] = ['context->thc_state'] + option['derived_actuals']

        if is_nn:
            prefix = 'THNN_{}'.format(env['THType'])
        else:
            prefix = env['THTensor'] + '_'

        call = prefix + CodeTemplate("${cname}(${derived_actuals})").substitute(env)
        ret = option['return']

        if ret['kind'] == 'arguments':
            if 'aten_custom_call' in option:
                scalar_check = None  # all aten_custom_call bodies handle settings on their own.
                body.append(CodeTemplate(option['aten_custom_call']).substitute(env))
            else:
                body.append(call + ";")
            arguments_indices = ret['arguments']
            arguments = [option['arguments'][argi]
                         for argi in arguments_indices]
            if scalar_check is not None:
                if len(arguments) > 1:
                    body.append("bool maybe_scalar = {};".format(scalar_check))
                    scalar_check = 'maybe_scalar'
                for arg in arguments:
                    body.append("{}_->maybeScalar({});".format(arg['name'], scalar_check))
            if len(arguments_indices) == 1:
                arg = arguments[0]
                body.append("return {};".format(arg['name']))
            else:
                types = [to_return_type(arg, option)['type'] for arg in arguments]
                # TODO: check for move semantics...
                names = [arg['name'] for arg in arguments]
                body.append(CodeTemplate("return std::tuple<${types}>(${names});").substitute(
                    types=types, names=names))
        elif ret['kind'] == 'type':
            if ret['type'] == 'THTensor*':
                maybe_scalar = "->maybeScalar({})".format(scalar_check) \
                               if scalar_check is not None \
                               else ""
                return_tensor = "return Tensor((new ${Tensor}(context,${arg_name}))${maybe_scalar},false);"
                body.append(CodeTemplate(return_tensor).substitute(env, arg_name=call, maybe_scalar=maybe_scalar))
            else:
                # we using int64_t for long in the API, so correct it here...
                if is_actual_return_long(ret):
                    call = "static_cast<int64_t>({})".format(call)
                body.append("return {};".format(call))
        else:
            raise Exception("NYI - return handling")
        return body

    def process_option(option):
        pair = (backend_type_env['Backend'],
                backend_type_env['ScalarName'])
        if pair in option['backend_type_pairs']:
            env = nested_dict(option, backend_type_env)
            body = emit_body(env, option)
            option['type_definition_body'] = body
            type_object_declarations.append(
                TYPE_DERIVED_DECLARATION.substitute(env))
            type_object_definitions.append(
                TYPE_DERIVED_DEFINITION.substitute(env))

    for declaration in declarations:
        for option in declaration['options']:
            if not option.get('skip', False):
                try:
                    process_option(option)
                except NYIError:
                    pass
    return type_object_declarations, type_object_definitions
