# -*- coding: utf-8 -*-
"""
This module contains the Pipeline class which provides a pluggable way to
define the transformations and the order in which they run on the AST.
"""
from __future__ import print_function, division, absolute_import

import ast as ast_module
import logging
import pprint
import random
import types
import llvm.core as lc

# import numba.closures
from numba import PY3
from numba import error
from numba import functions
from numba import transforms
from numba import control_flow
from numba import optimize
from numba import closures
from numba import reporting
from numba import normalize
from numba.codegen import llvmwrapper
from numba import ast_constant_folding as constant_folding
from numba.control_flow import ssa
from numba.codegen import translate
from numba import utils
from numba.missing import FixMissingLocations
from numba.type_inference import infer as type_inference
from numba.asdl import schema
from numba.minivect import minitypes
import numba.visitors

from numba.specialize import comparisons
from numba.specialize import loops
from numba.specialize import exceptions
from numba.specialize import funccalls
from numba.specialize import exttypes

from numba import astsix

logger = logging.getLogger(__name__)

#------------------------------------------------------------------------
# Utilities
#------------------------------------------------------------------------

def get_locals(ast, locals_dict):
    # TODO: Remove this
    if locals_dict is None:
        locals_dict = getattr(ast, "locals_dict", {})
    elif hasattr(ast, "locals_dict"):
        assert ast.locals_dict is locals_dict

    ast.locals_dict = locals_dict
    return locals_dict

def module_name(func):
    if func is None:
        name = "NoneFunc"
        func_id = random.randrange(1000000)
    else:
        name = '%s.%s' % (func.__module__, func.__name__)
        func_id = id(func)

    return 'tmp.module.%s.%x' % (name, func_id)

#------------------------------------------------------------------------
# Entry points
#------------------------------------------------------------------------

def run_pipeline2(env, func, func_ast, func_signature,
                  pipeline=None, **kwargs):
    assert pipeline is None
    assert kwargs.get('order', None) is None
    logger.debug(pprint.pformat(kwargs))
    kwargs['llvm_module'] = lc.Module.new(module_name(func))
    with env.TranslationContext(env, func, func_ast, func_signature,
                                **kwargs) as func_env:
        pipeline = env.get_pipeline(kwargs.get('pipeline_name', None))
        post_ast = pipeline(func_ast, env)
        func_signature = func_env.func_signature
        symtab = func_env.symtab

    return func_env, (func_signature, symtab, post_ast)

def run_env(env, func_env, **kwargs):
    env.translation.push_env(func_env)
    pipeline = env.get_pipeline(kwargs.get('pipeline_name', None))
    try:
        pipeline(func_env.ast, env)
    finally:
        env.translation.pop()

def _infer_types2(env, func, restype=None, argtypes=None, **kwargs):
    ast = functions._get_ast(func)
    func_signature = minitypes.FunctionType(return_type=restype,
                                            args=argtypes)
    return run_pipeline2(env, func, ast, func_signature, **kwargs)

def infer_types2(env, func, restype=None, argtypes=None, **kwargs):
    """
    Like run_pipeline, but takes restype and argtypes instead of a FunctionType
    """
    pipeline, (sig, symtab, ast) = _infer_types2(
        env, func, restype, argtypes, pipeline_name='type_infer', **kwargs)
    return sig, symtab, ast


def compile2(env, func, restype=None, argtypes=None, ctypes=False,
             compile_only=False, **kwds):
    """
    Compile a numba annotated function.

        - decompile function into a Python ast
        - run type inference using the given input types
        - compile the function to LLVM
    """
    # Let the pipeline create a module for the function it is compiling
    # and the user will link that in.
    assert 'llvm_module' not in kwds
    kwds['llvm_module'] = lc.Module.new(module_name(func))
    logger.debug(kwds)
    func_ast = functions._get_ast(func)
    func_signature = minitypes.FunctionType(return_type=restype,
                                            args=argtypes)
    #pipeline, (func_signature, symtab, ast) = _infer_types2(
    #            env, func, restype, argtypes, codegen=True, **kwds)
    with env.TranslationContext(env, func, func_ast, func_signature,
                                need_lfunc_wrapper=not compile_only,
                                **kwds) as func_env:
        pipeline = env.get_pipeline(kwds.get('pipeline_name', None))
        func_ast.pipeline = pipeline
        post_ast = pipeline(func_ast, env)
        func_signature = func_env.func_signature
        symtab = func_env.symtab
        t = func_env.translator
    return func_env


#------------------------------------------------------------------------
# Pipeline refactored code
#------------------------------------------------------------------------

class PipelineStage(object):

    is_composed = False

    def check_preconditions(self, ast, env):
        return True

    def check_postconditions(self, ast, env):
        return True

    def transform(self, ast, env):
        raise NotImplementedError('%r does not implement transform!' %
                                  type(self))

    def make_specializer(self, cls, ast, env, **kws):
        crnt = env.translation.crnt
        kws = kws.copy()
        kws.update(func_signature=crnt.func_signature,
                   nopython=env.translation.nopython,
                   symtab=crnt.symtab,
                   func_name=crnt.func_name,
                   llvm_module=crnt.llvm_module,
                   func_globals=crnt.function_globals,
                   locals=crnt.locals,
                   allow_rebind_args=env.translation.allow_rebind_args,
                   warn=env.translation.crnt.warn,
                   is_closure=crnt.is_closure,
                   closures=crnt.closures,
                   closure_scope=crnt.closure_scope,
                   env=env)
        return cls(env.context, crnt.func, ast, **kws)

    def __call__(self, ast, env):
        if env.stage_checks: self.check_preconditions(ast, env)

        if self.is_composed:
            ast = self.transform(ast, env)
        else:
            try:
                ast = self.transform(ast, env)
            except error.NumbaError as e:
                func_env = env.translation.crnt
                error_env = func_env.error_env
                if func_env.is_closure:
                    flags, parent_func_env = env.translation.stack[-2]
                    error_env.merge_in(parent_func_env.error_env)
                else:
                    reporting.report(env, exc=e)
                raise

        env.translation.crnt.ast = ast
        if env.stage_checks: self.check_postconditions(ast, env)
        return ast

class SimplePipelineStage(PipelineStage):

    transformer = None

    def transform(self, ast, env):
        transform = self.make_specializer(self.transformer, ast, env)
        return transform.visit(ast)


class AST3to2(PipelineStage):

    def transform(self, ast, env):
        if not PY3:
            return ast
        return astsix.AST3to2().visit(ast)


def ast3to2(ast, env):
    if not PY3:
        return ast
    return astsix.AST3to2().visit(ast)


def resolve_templates(ast, env):
    # TODO: Unify with decorators module
    crnt = env.translation.crnt
    if crnt.template_signature is not None:
        from numba import typesystem

        argnames = [name.id for name in ast.args.args]
        argtypes = list(crnt.func_signature.args)

        typesystem.resolve_templates(crnt.locals, crnt.template_signature,
                                     argnames, argtypes)
        crnt.func_signature = minitypes.FunctionType(
            return_type=crnt.func_signature.return_type,
            args=tuple(argtypes))

    return ast


def validate_signature(tree, env):
    arg_types = env.translation.crnt.func_signature.args
    if (isinstance(tree, ast_module.FunctionDef) and
        len(arg_types) != len(tree.args.args)):
        raise error.NumbaError(
            "Incorrect number of types specified in @jit() for function %r" %
            env.crnt.func_name)

    return tree

def update_signature(tree, env):
    func_env = env.translation.crnt
    func_signature = func_env.func_signature

    restype = func_signature.return_type
    if restype and (restype.is_struct or restype.is_complex):
        # Change signatures returning complex numbers or structs to
        # signatures taking a pointer argument to a complex number
        # or struct
        func_signature = func_signature.return_type(*func_signature.args)
        func_signature.struct_by_reference = True
        func_env.func_signature = func_signature

    return tree


def get_lfunc(env, func_env):
    lfunc = func_env.llvm_module.add_function(
        func_env.func_signature.to_llvm(env.context),
        func_env.mangled_name)
    return lfunc


def create_lfunc(tree, env):
    """
    Update the FunctionEnvironment with an LLVM function if the signature
    is known (try this before type inference to support recursion).
    """
    func_env = env.translation.crnt

    if (not func_env.lfunc and func_env.func_signature and
            func_env.func_signature.return_type):
        assert func_env.llvm_module is not None
        lfunc = get_lfunc(env, func_env)

        func_env.lfunc = lfunc
        if func_env.func:
            env.specializations.register_specialization(func_env)

    return tree

def create_lfunc1(tree, env):
    func_env = env.translation.crnt
    if not func_env.is_closure:
        create_lfunc(tree, env)

    return tree

def create_lfunc2(tree, env):
    func_env = env.translation.crnt
    assert func_env.func_signature and func_env.func_signature.return_type
    return create_lfunc1(tree, env)

def create_lfunc3(tree, env):
    func_env = env.translation.crnt
    create_lfunc(tree, env)
    return tree


class NormalizeASTStage(PipelineStage):
    def transform(self, ast, env):
        transform = self.make_specializer(normalize.NormalizeAST, ast, env)
        return transform.visit(ast)

class ControlFlowAnalysis(PipelineStage):
    _pre_condition_schema = None

    @property
    def pre_condition_schema(self):
        if self._pre_condition_schema is None:
            self._pre_condition_schema = schema.load('Python.asdl')
        return self._pre_condition_schema

    def check_preconditions(self, ast, env):
        self.pre_condition_schema.verify(ast)  # raises exception on error
        return True

    def transform(self, ast, env):
        transform = self.make_specializer(control_flow.ControlFlowAnalysis,
                                          ast, env)
        ast = transform.visit(ast)
        env.translation.crnt.symtab = transform.symtab
        ast.flow = transform.flow
        env.translation.crnt.ast.cfg_transform = transform
        return ast


def dump_cfg(ast, env):
    if env.translation.crnt.cfg_transform.graphviz:
        env.translation.crnt.cfg_transform.render_gv(ast)
    return ast


class ConstFolding(PipelineStage):
    def check_preconditions(self, ast, env):
        assert not hasattr(env.crnt, 'constvars')
        return super(ConstFolding, self).check_preconditions(ast, env)

    def check_postconditions(self, ast, env):
        assert hasattr(env.crnt, 'constvars')
        return super(ConstFolding, self).check_postconditions(ast, env)

    def transform(self, ast, env):
        const_marker = self.make_specializer(constant_folding.ConstantMarker,
                                             ast, env)
        const_marker.visit(ast)
        constvars = const_marker.get_constants()
        # FIXME: Make constvars a property of the FunctionEnvironment,
        # or nix this transformation pass.
        env.translation.crnt.constvars = constvars
        const_folder = self.make_specializer(constant_folding.ConstantFolder,
                                             ast, env, constvars=constvars)
        return const_folder.visit(ast)


class TypeInfer(PipelineStage):
    def check_preconditions(self, ast, env):
        assert env.translation.crnt.symtab is not None
        return super(TypeInfer, self).check_preconditions(ast, env)

    def transform(self, ast, env):
        crnt = env.translation.crnt
        type_inferer = self.make_specializer(type_inference.TypeInferer,
                                             ast, env, **crnt.kwargs)
        type_inferer.infer_types()
        crnt.func_signature = type_inferer.func_signature
        logger.debug("signature for %s: %s", crnt.func_name,
                     crnt.func_signature)
        crnt.symtab = type_inferer.symtab
        return ast


class TypeSet(PipelineStage):
    def transform(self, ast, env):
        visitor = self.make_specializer(type_inference.TypeSettingVisitor, ast,
                                        env)
        visitor.visit(ast)
        return ast


class ClosureTypeInference(PipelineStage):
    def transform(self, ast, env):
        type_inferer = self.make_specializer(
                            numba.closures.ClosureTypeInferer, ast, env)
        return type_inferer.visit(ast)


class TransformFor(PipelineStage):
    def transform(self, ast, env):
        transform = self.make_specializer(loops.TransformForIterable, ast,
                                          env)
        return transform.visit(ast)

#----------------------------------------------------------------------------
# Specializing/Lowering Transforms
#----------------------------------------------------------------------------

class Specialize(PipelineStage):
    def transform(self, ast, env):
        return ast

class RewriteArrayExpressions(PipelineStage):
    def transform(self, ast, env):
        from numba import array_expressions

        transformer = self.make_specializer(
            array_expressions.ArrayExpressionRewriteNative, ast, env)
        return transformer.visit(ast)

class SpecializeComparisons(PipelineStage):
    def transform(self, ast, env):
        transform = self.make_specializer(comparisons.SpecializeComparisons,
                                          ast, env)
        return transform.visit(ast)


class SpecializeSSA(PipelineStage):
    def transform(self, ast, env):
        ssa.specialize_ssa(ast)
        return ast

class SpecializeClosures(SimplePipelineStage):

    transformer = closures.ClosureSpecializer


class Optimize(PipelineStage):
    def transform(self, ast, env):
        return ast


class Preloader(PipelineStage):
    def transform(self, ast, env):
        transform = self.make_specializer(optimize.Preloader, ast, env)
        return transform.visit(ast)


class SpecializeLoops(PipelineStage):
    def transform(self, ast, env):
        transform = self.make_specializer(loops.SpecializeObjectIteration, ast,
                                          env)
        return transform.visit(ast)


class LateSpecializer(PipelineStage):
    def transform(self, ast, env):
        specializer = self.make_specializer(transforms.LateSpecializer, ast,
                                            env)
        return specializer.visit(ast)

class ExtensionTypeLowerer(PipelineStage):
    def transform(self, ast, env):
        specializer = self.make_specializer(exttypes.ExtensionTypeLowerer,
                                            ast, env)
        return specializer.visit(ast)

class SpecializeFunccalls(PipelineStage):
    def transform(self, ast, env):
        transform = self.make_specializer(funccalls.FunctionCallSpecializer,
                                          ast, env)
        return transform.visit(ast)


class SpecializeExceptions(PipelineStage):
    def transform(self, ast, env):
        transform = self.make_specializer(exceptions.ExceptionSpecializer, ast,
                                          env)
        return transform.visit(ast)


def cleanup_symtab(ast, env):
    "Pop original variables from the symtab"
    for var in env.translation.crnt.symtab.values():
        if not var.parent_var and var.renameable:
            env.translation.crnt.symtab.pop(var.name, None)
    return ast


class FixASTLocations(PipelineStage):
    def transform(self, ast, env):
        fixer = self.make_specializer(FixMissingLocations, ast, env)
        fixer.visit(ast)
        return ast


class CodeGen(PipelineStage):
    def transform(self, ast, env):
        func_env = env.translation.crnt
        func_env.translator = self.make_specializer(
            translate.LLVMCodeGenerator, ast, env,
            **func_env.kwargs)

        func_env.translator.translate()
        func_env.lfunc = func_env.translator.lfunc
        return ast

class LinkingStage(PipelineStage):
    """
    Link the resulting LLVM function into the global fat module.
    """

    def transform(self, ast, env):
        func_env = env.translation.crnt

        # Link libraries into module
        env.context.intrinsic_library.link(func_env.lfunc.module)
        # env.context.cbuilder_library.link(func_env.lfunc.module)
        env.constants_manager.link(func_env.lfunc.module)

        if func_env.link:
            # Link function into fat LLVM module
            func_env.lfunc = env.llvm_context.link(func_env.lfunc)
            func_env.translator.lfunc = func_env.lfunc

        func_env.lfunc_pointer = func_env.translator.lfunc_pointer

        return ast

class WrapperStage(PipelineStage):
    """
    Build a wrapper LLVM function around the compiled numba function to call
    it from Python.
    """

    def transform(self, ast, env):
        func_env = env.translation.crnt
        if func_env.is_closure:
            wrap = func_env.need_closure_wrapper
        else:
            wrap = func_env.wrap

        if wrap:
            numbawrapper, lfuncwrapper, _ = (
                llvmwrapper.build_wrapper_function(env))
            func_env.numba_wrapper_func = numbawrapper
            func_env.llvm_wrapper_func = lfuncwrapper

            # Set pointer to function for external code and numba.addressof()
            numbawrapper.lfunc_pointer = func_env.lfunc_pointer

        return ast

class ErrorReporting(PipelineStage):
    "Sort and issue warnings and errors"
    def transform(self, ast, env):
        reporting.report(env)
        return ast

class ComposedPipelineStage(PipelineStage):

    is_composed = True

    def __init__(self, stages=None):
        if stages is None:
            stages = []
        self.stages = [self.check_stage(stage)[1] for stage in stages]

    @staticmethod
    def check_stage(stage):
        def _check_stage_object(stage_obj):
            if (isinstance(stage_obj, type) and
                    issubclass(stage_obj, PipelineStage)):
                stage_obj = stage_obj()
            return stage_obj

        if isinstance(stage, str):
            name = stage
            def _stage(ast, env):
                stage_obj = getattr(env.pipeline_stages, name)
                return _check_stage_object(stage_obj)(ast, env)
            _stage.__name__ = name
            stage = _stage
        else:
            name = stage.__name__
            stage = _check_stage_object(stage)

        return name, stage

    def transform(self, ast, env):
        logger.debug('Running composed stages: %s', self.stages)
        for stage in self.stages:
            if env.debug:
                stage_tuple = (stage, utils.ast2tree(ast))
                logger.debug(pprint.pformat(stage_tuple))
            ast = stage(ast, env)
        return ast

    @classmethod
    def compose(cls, stage0, stage1):
        if isinstance(stage0, ComposedPipelineStage):
            stage0s = stage0.stages
        else:
            stage0s = [check_stage(stage0)[1]]
        if isinstance(stage1, ComposedPipelineStage):
            stage1s = stage1.stages
        else:
            stage1s = [check_stage(stage1)[1]]
        return cls(stage0s + stage1s)
