# This class is a subclass FoxySheepListener that rewrites the
# parse tree generated by ANTLR4. It "flattens" flat operators that ANTLR
# parses as left associative.

from antlr4.ParserRuleContext import ParserRuleContext
from antlr4.tree.Tree import TerminalNodeImpl, ParseTree
from antlr4.Token import CommonToken
from FoxySheep.generated.InputFormListener import InputFormListener
from FoxySheep.generated.InputFormParser import InputFormParser


def addChild(parent:ParserRuleContext, child:ParseTree, i:int = None):
    """Does what RuleContext.addChild is supposed to do. The ANTLR4 Python3 target
    doesn't follow the same API as the Java target. Some classes in the hierarchy
    don't have addChild, while others do."""
    try:
        if i is None:
            parent.children.append(child)
        else:
            parent.children.insert(i, child)
    except:
        parent.children = [child]


def adopt(parent:ParserRuleContext, child:ParseTree):
    """Convenience function to create a bidirectional parent-child relationship."""
    addChild(parent, child)
    child.parentCtx = parent


def makeNumber(parent:ParserRuleContext, n:int):
    """This node represents a virtual ParseTree node that does not come from
    the parser but rather is constructed via a rewriting rule. For example,
    the expression "a-b" is parsed as "Plus[a, Times[-1, b]]" in exitPlusOp(),
    so a node for "-1" needs to be created even though "-1" does not appear as
    a token in the token stream.

    Note that makeNumber sets the NumberContext's parent but does not add
    anything to parent's children."""

    # The hierarchy is:
    #   CommonToken->TerminalNodeImpl->NumberLiteralContext
    #       ->NumberBaseTenContext->NumberContext->parent

    digits_token = CommonToken(type=InputFormParser.DIGITS)
    digits_token._text = str(n)

    number = InputFormParser.NumberContext(None, InputFormParser.ExprContext(None, parent=parent))

    number_literal = InputFormParser.NumberBaseTenContext(None, InputFormParser.NumberLiteralContext(None, parent=number))

    number_literal.addTokenNode(digits_token)
    addChild(number, number_literal)

    return number


def flatten(ctx:ParserRuleContext):
    """Takes a ParserRuleContext of a binary operator and "flattens"
    the operator if one of its operands is the same binary operator context.

    This function only flattens if the operator is the same and also
    keeps the operators intact."""

    # If the child isn't the same construct, nothing to do.
    if ctx.getChild(0).__class__ != ctx.__class__:
        return

    lhs = ctx.getChild(0)
    op = ctx.getChild(1)
    rhs = ctx.getChild(2)
    lhsop = lhs.getChild(1)

    # If the operator of the nested Context isn't the same, nothing to do.
    # The operator is always in position 1 for infix operators. We do this
    # check because some Contexts that use the same context for multiple
    # operators.
    if op.getSymbol().type != lhsop.getSymbol().type:
        return

    # Clear all children.
    ctx.children.clear()
    # Add all children of lhs. (Also adds the operator of the lhs.)
    ctx.children.extend(lhs.children)
    # Finally, add the rhs back in.
    ctx.children.append(rhs)


def rewriteSpan(ctx:ParserRuleContext):
    opIndex = [i for (i, c) in enumerate(ctx.children) if c.getText() == ';;']
    spanExpressions = []

    opIndexSize = len(opIndex)
    childrenSize = len(ctx.children)
    # We emulate a for-loop with a while-loop. We can't use python's for
    # because we need to change the index variables in the loop body.
    i = 0
    nextOp = 0
    while True:
        if nextOp + 1 < opIndexSize \
            and opIndex[nextOp + 1] + 1 < childrenSize \
                and isinstance(ctx.children[opIndex[nextOp + 1] + 1], InputFormParser.ExprContext):
            # There is a next ";;"
            # and there is a node after the next ";;"
            # and there is a second ";;" followed by an expr.
            i = opIndex[nextOp + 1] + 1
            spanExpressions.append(i)
            # We want nextOp to end at the last ";;" of the current expression.
            nextOp += 1
        else:
            # There is no second ";;" belonging to this expression.
            if opIndex[nextOp] + 1 < childrenSize \
                    and isinstance(ctx.children[opIndex[nextOp + 1] + 1], InputFormParser.ExprContext):
                # There is a node after ";;"
                # and this span expression ends in an expr.
                i = opIndex[nextOp] + 1
                spanExpressions.append(i)
            else:
                # This span expression ends in the current ";;".
                i = opIndex[nextOp]
                spanExpressions.append(i)
        # Check for end of for-loop
        if i >= childrenSize or nextOp >= opIndexSize:
            break
        # Set up the next loop.
        i += 1
        nextOp +=1

    # At this point spanExpressions holds the index of the last child of each
    # span expression. It might be that after all of this there is nothing to do.
    if len(spanExpressions) == 1:
        return
    # Otherwise there is more than one span expression, and we need to rewrite
    # the tree replacing the Span?Context this method was invoked on with a
    # TimesContext.
    timesctx = InputFormParser.TimesContext(None, ctx)
    timesctx.children = [] # Clear children

    # Add each span expression as a child to timesctx.
    j = 0
    for i, spanExp in enumerate(spanExpressions):
        # i is the index of the current span expression in spanExpressions,
        # and j is the index to the beginning of the new span expression's
        # children in ctx.children. We make new SpanAContext objects for each
        # span expression.
        span = InputFormParser.SpanAContext(None, ctx)
        adopt(timesctx, span)
        span.children = ctx.children[j:spanExpressions[i]]
        # update j to be the beginning of the next expression.
        j = spanExpressions[i] + 1

    # Finally, detach the span this method was invoked on from its parent and
    # replace with the TimesContext.
    parentsChildren = ctx.parentCtx.children
    ctxIndex = parentsChildren.index(ctx)
    parentsChildren[ctxIndex] = timesctx
    ctx.parentCtx = timesctx



class PostParser(InputFormListener):

    def exitComparison(self, ctx:InputFormParser.ComparisonContext):
        """Inequality[]

        This function flattens and keeps the operators intact. It differs from
        flatten() in that we flatten if the class is the same but don't check
        if the operator is the same."""

        # If the child isn't the same construct, nothing to do.
        if ctx.getChild(0).__class__ != ctx.__class__:
            return

        lhs = ctx.getChild(0)
        op = ctx.getChild(1)
        rhs = ctx.getChild(2)
        lhsop = lhs.getChild(1)

        # This is where we differ from flatten(). We don't do the following
        # check.
        # if op.getSymbol().getType() != lhsop.getSymbol().getType():
        #     return

        # Clear all children.
        ctx.children.clear()
        # Add all children of lhs. (Also adds the operator of the lhs.)
        ctx.children.extend(lhs.children)
        # Keep the operator.
        ctx.children.append(op)
        # Finally, add the rhs back in.
        ctx.children.append(rhs)

    def exitCompoundExpression(self, ctx:InputFormParser.CompoundExpressionContext):
        """Composition[expr1,expr2]	e@*e@*e.

        ANTLR4 parses this rule as right associative for some reason, so
        we cannot use flatten(). The code is actually much simpler than
        flatten because of the right associativity."""

        childCount = ctx.getChildCount()

        # If there is no RHS, nothing to do.
        if childCount < 3:
            return

        rhs = ctx.getChild(childCount-1)

        # If the RHS child isn't the same construct, nothing to do.
        if rhs.__class__ != ctx.__class__:
            return

        # Remove RHS child.
        ctx.removeLastChild()
        # Add all children of rhs. (Also adds the operator of the rhs.)
        ctx.children.extend(rhs.children)

    def exitComposition(self, ctx:InputFormParser.CompositionContext):
        flatten(ctx)

    def exitRightComposition(self, ctx:InputFormParser.RightCompositionContext):
        flatten(ctx)

    def exitStringJoin(self, ctx:InputFormParser.StringJoinContext):
        flatten(ctx)

    def exitSmallCircle(self, ctx:InputFormParser.SmallCircleContext):
        flatten(ctx)

    def exitSpan(self, ctx:InputFormParser.ExprContext):
        """Span[expr1,expr2,expr3] e;;e;;e
        Parsing Span nodes is a complete mess."""

        # The only difference between the code for exitSpanA and exitSpanB is
        # the index of the rhs.
        rhsIndex = 1
        if isinstance(ctx, InputFormParser.SpanBContext):
            rhsIndex = 0


        # If there is no RHS expr, nothing to do.
        if len(ctx.expr()) == rhsIndex:
            return

        # Get the RHS expr.
        rhs = ctx.expr(rhsIndex)

        # If the RHS child isn't a SpanA, nothing to do.
        if not isinstance(rhs, InputFormParser.SpanAContext):
            return

        # Remove the last expr
        ctx.removeLastChild()
        # Replace it with its children
        ctx.children.extend(rhs.children)

        # If this is the topmost Span context, rewrite the tree.
        if not (isinstance(ctx.parentCtx, InputFormParser.SpanAContext)
                or isinstance(ctx.parentCtx, InputFormParser.SpanBContext)):
            rewriteSpan(ctx)

    def exitSpanA(self, ctx:InputFormParser.SpanAContext):
        """Span[expr1,expr2,expr3] e;;e;;e
        Parsing Span nodes is a complete mess."""
        self.exitSpan(ctx)

    def exitSpanB(self, ctx:InputFormParser.SpanBContext):
        """Span[expr1,expr2,expr3] e;;e;;e
        Parsing Span nodes is a complete mess."""
        self.exitSpan(ctx)

    def exitCircleDot(self, ctx:InputFormParser.CircleDotContext):
        flatten(ctx)

    def exitNonCommutativeMultiply(self, ctx:InputFormParser.NonCommutativeMultiplyContext):
        flatten(ctx)

    def exitCross(self, ctx:InputFormParser.CrossContext):
        flatten(ctx)

    def exitDot(self, ctx:InputFormParser.DotContext):
        flatten(ctx)

    def exitDiamond(self, ctx:InputFormParser.DiamondContext):
        flatten(ctx)

    def exitWedge(self, ctx:InputFormParser.WedgeContext):
        flatten(ctx)

    def exitVee(self, ctx:InputFormParser.VeeContext):
        flatten(ctx)

    def exitCircleTimes(self, ctx:InputFormParser.CircleTimesContext):
        flatten(ctx)

    def exitCenterDot(self, ctx:InputFormParser.CenterDotContext):
        flatten(ctx)

    def exitTimes(self, ctx:InputFormParser.TimesContext):
        """Times[expr1,expr2]

        We need to flatten over both Times and implicit Times. So
        flatten() isn't going to cut it because sometimes there is
        no operator."""

        # If the child isn't the same construct, nothing to do.
        if not isinstance(ctx.getChild(0), InputFormParser.TimesContext):
            return

        lhs = ctx.expr(0)
        rhs = ctx.expr(1)
        # Replace children with all children of lhs. (Also adds the operator of the lhs.)
        ctx.children = lhs.children
        # Add the rhs back in.
        ctx.children.append(rhs)

    def exitStar(self, ctx:InputFormParser.StarContext):
            flatten(ctx)

    def exitVerticalTilde(self, ctx:InputFormParser.VerticalTildeContext):
        flatten(ctx)

    def exitCoproduct(self, ctx:InputFormParser.CoproductContext):
        flatten(ctx)

    def exitCap(self, ctx:InputFormParser.CapContext):
        flatten(ctx)

    def exitCup(self, ctx:InputFormParser.CupContext):
        flatten(ctx)

    def exitCirclePlus(self, ctx:InputFormParser.CirclePlusContext):
        flatten(ctx)

    def exitPlusOp(self, ctx:InputFormParser.PlusOpContext):
        """PlusOp[expr1,expr2]

        We have to treat PlusOp special, because we have to keep the
        operators intact, and only  plus and minus (not PlusMinus or
        MinusPlus) are flat. The situation is complicated by the fact
        that Mathematica parses "a-b" as "Plus[a, Times[-1, b]]". We
        Rewrite the parse tree, inserting the Times context and
        changing BINARYMINUS to BINARYPLUS."""

        # If the op isn't Plus or Minus, nothing to do.
        if ctx.BINARYMINUS() is None and ctx.BINARYPLUS() is None:
            return

        # Since ANTLR4 parses this operator as left associative, we only
        # need to check the left hand side expr.
        rhs = ctx.getChild(2)

        # If the operator of the PlusOp is BINARYMINUS, we rewrite the tree as
        # "Plus[lhs, Times[-1, rhs]]". Note that if rhs is TIMES, we have to
        # keep that TIMES flat.
        if ctx.BINARYMINUS() is not None:
            # Construct Times, or obtain it from the rhs.
            times = None
            if isinstance(rhs, InputFormParser.TimesContext):
                times = rhs
            else:
                # If rhs is already a times, keep it flat.
                times = InputFormParser.TimesContext(None, InputFormParser.ExprContext(None))
                ctx.children.remove(rhs)
                adopt(ctx, times)
                adopt(times, rhs)
            # Add "-1" as the first child of Times.
            addChild(times, makeNumber(times, -1), 0)

            # Finally, we have to change operator to BINARYPLUS.
            plustoken = CommonToken(type=InputFormParser.BINARYPLUS)
            plustoken.text = '+'
            plus = TerminalNodeImpl(plustoken)
            # Replace minus token with plus.
            ctx.children[1] = plus
            plus.parentCtx = ctx

        # Flatten
        flatten(ctx)

    def exitIntersection(self, ctx:InputFormParser.IntersectionContext):
        flatten(ctx)

    def exitUnion(self, ctx:InputFormParser.UnionContext):
        flatten(ctx)

    def exitVerticalBar(self, ctx:InputFormParser.VerticalBarContext):
        flatten(ctx)

    def exitSame(self, ctx:InputFormParser.SameContext):
        flatten(ctx)

    def exitSetContainment(self, ctx:InputFormParser.SetContainmentContext):
        flatten(ctx)

    def exitAndOr(self, ctx:InputFormParser.ExprContext, opText:str):
        """Unifies the code for flattening And/Nand and Or/Nor."""

        # If the child isn't the same construct, nothing to do.
        if ctx.getChild(0).__class__ != ctx.__class__:
            return

        lhs = ctx.getChild(0)
        op = ctx.getChild(1)
        rhs = ctx.getChild(2)
        lhsop = lhs.getChild(1).getText()
        # Here's the part that's different from flatten().
        # If childOp is an Nand or parentOp is a Nand, then we need child==parent.
        if (lhsop == opText or op.getText() == opText) and op.getText() != lhsop:
            return

        # Clear all children.
        ctx.children.clear()
        # Add all children of lhs. (Also adds the operator of the lhs.)
        ctx.children.extend(lhs.children)
        # Finally, add the rhs back in.
        ctx.children.append(rhs)

    def exitAnd(self, ctx:InputFormParser.AndContext):
        self.exitAndOr(ctx, '\u22bc')

    def exitOr(self, ctx:InputFormParser.OrContext):
        self.exitAndOr(ctx, '\u22bd')

    def exitXor(self, ctx:InputFormParser.XorContext):
        flatten(ctx)

    def exitEquivalent(self, ctx:InputFormParser.EquivalentContext):
        flatten(ctx)

    def exitAlternatives(self, ctx:InputFormParser.AlternativesContext):
        flatten(ctx)

    def exitStringExpression(self, ctx:InputFormParser.StringExpressionContext):
        flatten(ctx)

    def exitColon(self, ctx:InputFormParser.ColonContext):
        flatten(ctx)

    def exitVerticalSeparator(self, ctx:InputFormParser.VerticalSeparatorContext):
        flatten(ctx)
