
# pylint: disable=invalid-name, unused-import
"""A parser for Relay's text format."""
from __future__ import absolute_import

import sys

from collections import deque
from typing import TypeVar, Deque, Tuple, Optional, Union, NamedTuple, List, Callable, Any, Dict

import tvm

from . import module
from .base import Span, SourceName
from . import expr
from . import ty
from . import op


class ParseError(Exception):
    """Exception type for parse errors."""

    def __init__(self, message):
        # type: (str) -> None
        super(ParseError, self).__init__()
        self.message = message

PYTHON_VERSION = sys.version_info.major
try:
    if PYTHON_VERSION == 2:
        from .grammar.py2.RelayVisitor import RelayVisitor
        from .grammar.py2.RelayParser import RelayParser
        from .grammar.py2.RelayLexer import RelayLexer
    else:
        from .grammar.py3.RelayVisitor import RelayVisitor
        from .grammar.py3.RelayParser import RelayParser
        from .grammar.py3.RelayLexer import RelayLexer
except ImportError:
    raise ParseError("Couldn't find ANTLR parser. Try building with USE_ANTLR=ON.")

try:
    from antlr4 import ParserRuleContext, InputStream, CommonTokenStream
    from antlr4.tree.Tree import TerminalNode
except ImportError:
    raise ParseError("Couldn't find ANTLR runtime." +
                     "Try running `pip{} install antlr4-python{}-runtime`."
                     .format(PYTHON_VERSION, PYTHON_VERSION))

BINARY_OPS = {
    RelayParser.MUL: op.multiply,
    RelayParser.DIV: op.divide,
    RelayParser.ADD: op.add,
    RelayParser.SUB: op.subtract,
    RelayParser.LT:  op.less,
    RelayParser.GT:  op.greater,
    RelayParser.LE:  op.less_equal,
    RelayParser.GE:  op.greater_equal,
    RelayParser.EQ:  op.equal,
    RelayParser.NE:  op.not_equal,
}

TYPE_PREFIXES = [
    "int",
    "uint",
    "float",
    "bool",
]

T = TypeVar("T")
Scope = Deque[Tuple[str, T]]
Scopes = Deque[Scope[T]]

def lookup(scopes, name):
    # type: (Scopes[T], str) -> Optional[T]
    """Look up `name` in `scopes`."""

    for scope in scopes:
        for key, val in scope:
            if key == name:
                return val
    return None

def spanify(f):
    """A decorator which attaches span information
       to the value returned by calling `f`.

       Intended for use with the below AST visiting
       methods. The idea is that after we do the work
       of constructing the AST we attach Span information.
    """

    def _wrapper(*args, **kwargs):
        # Assumes 0th arg is self and gets source_name from object.
        sn = args[0].source_name
        # Assumes 1st arg is an ANTLR parser context.
        ctx = args[1]
        ast = f(*args, **kwargs)
        line, col = ctx.getSourceInterval()
        sp = Span(sn, line, col)
        ast.set_span(sp)
        return ast
    return _wrapper

# TODO(@jmp): Use https://stackoverflow.com/q/13889941
# to figure out how to get ANTLR4 to be more unhappy about syntax errors
class ParseTreeToRelayIR(RelayVisitor):
    """Parse Relay text format into Relay IR."""

    def __init__(self, source_name):
        # type: (str) -> None
        self.source_name = source_name
        self.module = module.Module({})   # type: module.Module

        # Adding an empty scope allows naked lets without pain.
        self.var_scopes = deque([deque()])          # type: Scopes[expr.Var]
        self.global_var_scope = deque()             # type: Scope[expr.GlobalVar]
        self.type_param_scopes = deque([deque()])   # type: Scopes[ty.TypeVar]
        self.graph_expr = []                        # type: List[expr.Expr]

        super(ParseTreeToRelayIR, self).__init__()


    def enter_var_scope(self):
        # type: () -> None
        """Enter a new Var scope so it can be popped off later."""

        self.var_scopes.appendleft(deque())

    def exit_var_scope(self):
        # type: () -> Scope[expr.Var]
        """Pop off the current Var scope and return it."""

        return self.var_scopes.popleft()

    def mk_var(self, name, type_):
        # type: (str, ty.Type) -> expr.Var
        """Create a new Var and add it to the Var scope."""

        var = expr.Var(name, type_)
        self.var_scopes[0].appendleft((name, var))
        return var

    def mk_global_var(self, name):
        # type: (str) -> expr.GlobalVar
        """Create a new GlobalVar and add it to the GlobalVar scope."""

        var = expr.GlobalVar(name)
        self.global_var_scope.append((name, var))
        return var

    def enter_type_param_scope(self):
        # type: () -> None
        """Enter a new TypeVar scope so it can be popped off later."""

        self.type_param_scopes.appendleft(deque())

    def exit_type_param_scope(self):
        # type: () -> Scope[ty.TypeVar]
        """Pop off the current TypeVar scope and return it."""

        return self.type_param_scopes.popleft()

    def mk_typ(self, name, kind):
        # (str, ty.Kind) -> ty.TypeVar
        """Create a new TypeVar and add it to the TypeVar scope."""

        typ = ty.TypeVar(name, kind)
        self.type_param_scopes[0].appendleft((name, typ))
        return typ

    def visitTerminal(self, node):
        # type: (TerminalNode) -> Union[expr.Expr, int, float]
        """Visit lexer tokens that aren't ignored or visited by other functions."""

        node_type = node.getSymbol().type
        node_text = node.getText()
        name = node_text[1:]

        # variables
        if node_type == RelayLexer.GLOBAL_VAR:
            return lookup(deque([self.global_var_scope]), node_text[1:])
        elif node_type == RelayLexer.LOCAL_VAR:
            # Remove the leading '%' and lookup the name.
            var = lookup(self.var_scopes, name)
            if var is None:
                raise ParseError("Couldn't resolve `{}`.".format(name))
            return var
        elif node_type == RelayLexer.GRAPH_VAR:
            try:
                return self.graph_expr[int(name)]
            except IndexError:
                raise ParseError("Couldn't resolve `{}`".format(name))

        # data types
        elif node_type == RelayLexer.NAT:
            return int(node_text)
        elif node_type == RelayLexer.FLOAT:
            return float(node_text)
        elif node_type == RelayLexer.BOOL_LIT:
            if node_text == "True":
                return True
            elif node_text == "False":
                return False
            else:
                raise ParseError("Unrecognized BOOL_LIT: `{}`".format(node_text))

        else:
            raise ParseError("todo: {}".format(node_text))

    def visit_list(self, ctx_list):
        # type: (List[ParserRuleContext]) -> List[Any]
        """"Visit a list of contexts."""

        return [self.visit(ctx) for ctx in ctx_list]

    def getType_(self, ctx):
        # type: (Optional[RelayParser.Type_Context]) -> Optional[ty.Type]
        """Return a (possibly None) Relay type."""

        if ctx is None:
            return None

        return self.visit(ctx)

    def visitProg(self, ctx):
        # type: (RelayParser.ProgContext) -> Union[expr.Expr, module.Module]
        if ctx.defn():
            self.visit_list(ctx.defn())
            return self.module

        return self.visit(ctx.expr())

    # Exprs

    def visitOpIdent(self, ctx):
        # type: (RelayParser.OpIdentContext) -> op.Op
        return op.get(ctx.CNAME().getText())

    # pass through
    def visitParens(self, ctx):
        # type: (RelayParser.ParensContext) -> expr.Expr
        return self.visit(ctx.expr())

    # pass through
    def visitBody(self, ctx):
        # type: (RelayParser.BodyContext) -> expr.Expr
        return self.visit(ctx.expr())

    def visitScalarFloat(self, ctx):
        # type: (RelayParser.ScalarFloatContext) -> expr.Constant
        return expr.const(self.visit(ctx.FLOAT()))

    def visitScalarInt(self, ctx):
        # type: (RelayParser.ScalarIntContext) -> expr.Constant
        return expr.const(self.visit(ctx.NAT()))

    def visitScalarBool(self, ctx):
        # type: (RelayParser.ScalarBoolContext) -> expr.Constant
        return expr.const(self.visit(ctx.BOOL_LIT()))

    def visitNeg(self, ctx):
        # type: (RelayParser.NegContext) -> Union[expr.Constant, expr.Call]
        val = self.visit(ctx.expr())
        if isinstance(val, expr.Constant) and val.data.asnumpy().ndim == 0:
            # fold Neg in for scalars
            return expr.const(-val.data.asnumpy().item())

        return op.negative(val)

    def visitTuple(self, ctx):
        # type: (RelayParser.TupleContext) -> expr.Tuple
        tup = self.visit_list(ctx.expr())
        return expr.Tuple(tup)

    # Currently doesn't support mutable sequencing.
    def visitLet(self, ctx):
        # type: (RelayParser.SeqContext) -> expr.Let
        """Desugar various sequence constructs to Relay Let nodes."""
        if ctx.MUT() is not None:
            raise ParseError("Mutation is currently unsupported.")

        if ctx.var() is None or ctx.var().ident() is None:
            # anonymous identity
            ident = "_"
            type_ = None
        else:
            local_var = ctx.var().ident().LOCAL_VAR()
            if local_var is None:
                raise ParseError("Only local ids may be used in `let`s.")
            ident = local_var.getText()[1:]
            type_ = self.getType_(ctx.var().type_())

        var = self.mk_var(ident, type_)

        self.enter_var_scope()
        value = self.visit(ctx.expr(0))
        self.exit_var_scope()

        body = self.visit(ctx.expr(1))

        return expr.Let(var, value, body)

    def visitBinOp(self, ctx):
        # type: (RelayParser.BinOpContext) -> expr.Call
        """Desugar binary operators."""
        arg0, arg1 = self.visit_list(ctx.expr())
        relay_op = BINARY_OPS.get(ctx.op.type)

        if relay_op is None:
            raise ParseError("Unimplemented binary op.")

        return relay_op(arg0, arg1)

    @spanify
    def visitVar(self, ctx):
        # type: (RelayParser.VarContext) -> expr.Var
        """Visit a single variable."""
        ident = ctx.ident().LOCAL_VAR()

        if ident is None:
            raise ParseError("Only local ids may be used in vars.")

        type_ = self.getType_(ctx.type_())

        return self.mk_var(ident.getText()[1:], type_)

    def visitVarList(self, ctx):
        # type: (RelayParser.VarListContext) -> List[expr.Var]
        return self.visit_list(ctx.var())

    # TODO: support a larger class of values than just Relay exprs
    def visitAttr(self, ctx):
        # type: (RelayParser.AttrContext) -> Tuple[str, expr.Expr]
        return (ctx.CNAME().getText(), self.visit(ctx.expr()))

    def visitAttrList(self, ctx):
        # type: (RelayParser.AttrListContext) -> Dict[str, expr.Expr]
        return dict(self.visit_list(ctx.attr()))

    def visitArgList(self,
                     ctx    # type: RelayParser.ArgListContext
                    ):
        # type: (...) -> Tuple[Optional[List[expr.Var]], Optional[Dict[str, expr.Expr]]]
        var_list = self.visit(ctx.varList()) if ctx.varList() else None
        attr_list = self.visit(ctx.attrList()) if ctx.attrList() else None

        return (var_list, attr_list)

    def mk_func(self, ctx):
        # type: (Union[RelayParser.FuncContext, RelayParser.DefnContext]) -> expr.Function
        """Construct a function from either a Func or Defn."""

        # Enter var scope early to put params in scope.
        self.enter_var_scope()
        # Capture type params in params.
        self.enter_type_param_scope()
        var_list, attr_list = self.visit(ctx.argList())
        ret_type = self.getType_(ctx.type_())

        type_params = list(self.exit_type_param_scope())
        if type_params:
            _, type_params = zip(*type_params)

        body = self.visit(ctx.body())
        self.exit_var_scope()

        attrs = tvm.make.node("DictAttrs", **attr_list) if attr_list is not None else None

        return expr.Function(var_list, body, ret_type, type_params, attrs)

    @spanify
    def visitFunc(self, ctx):
        # type: (RelayParser.FuncContext) -> expr.Function
        return self.mk_func(ctx)

    # TODO: how to set spans for definitions?
    # @spanify
    def visitDefn(self, ctx):
        # type: (RelayParser.DefnContext) -> None
        ident = ctx.ident().GLOBAL_VAR()
        if ident is None:
            raise ParseError("Only global ids may be used in `def`s.")
        ident_name = ident.getText()[1:]
        ident = self.mk_global_var(ident_name)

        self.module[ident] = self.mk_func(ctx)

    @spanify
    def visitCall(self, ctx):
        # type: (RelayParser.CallContext) -> expr.Call
        visited_exprs = self.visit_list(ctx.expr())

        func = visited_exprs[0]
        args = visited_exprs[1:]

        return expr.Call(func, args, None, None)

    @spanify
    def visitIfElse(self, ctx):
        # type: (RelayParser.IfElseContext) -> expr.If
        """Construct a Relay If node. Creates a new scope for each branch."""
        cond = self.visit(ctx.expr())

        self.enter_var_scope()
        true_branch = self.visit(ctx.body(0))
        self.exit_var_scope()

        self.enter_var_scope()
        false_branch = self.visit(ctx.body(1))
        self.exit_var_scope()

        return expr.If(cond, true_branch, false_branch)

    @spanify
    def visitGraph(self, ctx):
        # type: (RelayParser.GraphContext) -> expr.Expr
        """Visit a graph variable assignment."""
        if ctx.ident().GRAPH_VAR() is None:
            raise ParseError("Expected a graph var, but got `{}`".format(ctx.ident().getText()))
        graph_nid = int(ctx.ident().GRAPH_VAR().getText()[1:])

        self.enter_var_scope()
        value = self.visit(ctx.expr(0))
        self.exit_var_scope()

        if graph_nid != len(self.graph_expr):
            raise ParseError(
                "Expected new graph variable to be `%{}`,".format(len(self.graph_expr)) + \
                "but got `%{}`".format(graph_nid))
        self.graph_expr.append(value)

        kont = self.visit(ctx.expr(1))
        return kont

    # Types

    # pylint: disable=unused-argument
    def visitIncompleteType(self, ctx):
        # type (RelayParser.IncompleteTypeContext) -> None:
        return None

    def visitIdentType(self, ctx):
        # type: (RelayParser.IdentTypeContext) -> Union[ty.TensorType, str]
        ident_type = ctx.CNAME().getText()

        # look through all type prefixes for a match
        for type_prefix in TYPE_PREFIXES:
            if ident_type.startswith(type_prefix):
                return ty.scalar_type(ident_type)

        raise ParseError("Unknown builtin type: {}".format(ident_type))

    # def visitCallType(self, ctx):
    #     # type: (RelayParser.CallTypeContext) -> Union[expr.Expr, ty.TensorType]
    #     ident_type = ctx.identType().CNAME().getText()

    #     args = self.visit_list(ctx.type_())

    #     if not args:
    #         raise ParseError("Type-level functions must have arguments!")

    #     func_type = TYPE_FUNCS.get(ident_type)(args)

    #     if func_type is None:
    #         raise ParseError("Unknown type-level function: `{}`".format(ident_type))
    #     else:
    #         return func_type

    def visitParensShape(self, ctx):
        # type: (RelayParser.ParensShapeContext) -> int
        return self.visit(ctx.shape())

    def visitShapeSeq(self, ctx):
        # type: (RelayParser.ShapeSeqContext) -> List[int]
        return self.visit_list(ctx.shape())

    def visitTensorType(self, ctx):
        # type: (RelayParser.TensorTypeContext) -> ty.TensorType
        """Create a simple tensor type. No generics."""

        shape = self.visit(ctx.shapeSeq())
        dtype = self.visit(ctx.type_())

        if not isinstance(dtype, ty.TensorType):
            raise ParseError("Expected dtype to be a Relay base type.")

        dtype = dtype.dtype

        return ty.TensorType(shape, dtype)

    def visitTupleType(self, ctx):
        # type: (RelayParser.TupleTypeContext) -> ty.TupleType
        return ty.TupleType(self.visit_list(ctx.type_()))

    def visitFuncType(self, ctx):
        # type: (RelayParser.FuncTypeContext) -> ty.FuncType
        types = self.visit_list(ctx.type_())

        arg_types = types[:-1]
        ret_type = types[-1]

        return ty.FuncType(arg_types, ret_type, [], None)

def make_parser(data):
    # type: (str) -> RelayParser
    """Construct a RelayParser a given data stream."""
    input_stream = InputStream(data)
    lexer = RelayLexer(input_stream)
    token_stream = CommonTokenStream(lexer)
    return RelayParser(token_stream)

__source_name_counter__ = 0

def fromtext(data, source_name=None):
    # type: (str, str) -> Union[expr.Expr, module.Module]
    """Parse a Relay program."""
    global __source_name_counter__

    if source_name is None:
        source_name = "source_file{0}".format(__source_name_counter__)

    if isinstance(source_name, str):
        source_name = SourceName(source_name)

    tree = make_parser(data).prog()
    return ParseTreeToRelayIR(source_name).visit(tree)
