# fx-92 Scientifique Collège+ language interpreter: AST interpreter

from fx92.ast import N, Node
import math

class Context:
    """
    An execution context, tied to a window, in which a script can be run.
    Provides both variables, state management, and interpretation.
    """

    def __init__(self, window):
        """
        Create a new execution context with a clean variable state and a
        window.
        """

        self.vars = { name: 0 for name in "MABCDEF" }
        self.w = window

        self.x = float(0)
        self.y = float(0)
        self.theta = float(0)
        self.pen = False

        # TODO: Also specify the style!

        self.w.clear(self.w.WHITE)
        self.w.update()

    def settheta(self, newtheta):
        self.theta = newtheta % 360

    def run(self, n):
        """
        Execute a full program tree. For expression nodes, returns a value; for
        statements returns None. Also draws as a side-effect.
        """

        # Core nodes

        if n.type == N.PROGRAM:
            for stmt in n.args:
                self.run(stmt)

        # Basic statements

        elif n.type == N.FORWARD:
            dist, = self.runs(n)

            if self.pen:
                self.w.linef(self.x, self.y, self.theta, dist, self.w.BLACK)
                self.w.update()

            self.x += dist * math.cos(self.theta * math.pi / 180)
            self.y += dist * math.sin(self.theta * math.pi / 180)

        elif n.type == N.ROTATE:
            angle, = self.runs(n)
            self.settheta(self.theta + angle)

        elif n.type == N.ORIENT:
            angle, = self.runs(n)
            self.settheta(angle)

        elif n.type == N.GOTO:
            x, y = self.runs(n)

            if self.pen:
                self.w.line(self.x, self.y, x, y, self.w.BLACK)
                self.w.update()

            self.x = x
            self.y = y

        elif n.type == N.PENDOWN:
            self.pen = True
        elif n.type == N.PENUP:
            self.pen = False

        elif n.type == N.ASSIGN:
            val = self.run(n.args[0])
            var = n.args[1]

            if var == "x":
                self.x = val
            elif var == "y":
                self.y = val
            elif var == "theta":
                self.theta = val
            else:
                self.vars[var] = val

        elif n.type == N.INPUT:
            # TODO: Input stmt should call into the lexer+parser (urg)
            raise Exception("Input statement not supported yet x_x")

        elif n.type == N.PRINT:
            val, = self.runs(n)
            print(val)

        elif n.type == N.STYLE:
            # TODO: Support style
            raise Exception("Style statement not supported yet x_x")

        elif n.type == N.WAIT:
            print("---- pause")
            print("x={} y={} theta={}".format(self.x, self.y, self.theta))
            print("A={} B={}".format(self.vars['A'], self.vars['B']))
            print("C={} D={}".format(self.vars['C'], self.vars['D']))
            print("E={} F={}".format(self.vars['E'], self.vars['F']))
            print("M={}".format(self.vars['M']))
            self.w.pause()

        # Flow control

        elif n.type == N.REPEAT:
            count = self.run(n.args[0])
            count = int(count)

            for i in range(count):
                self.run(n.args[1])

        elif n.type == N.WHILE:
            while self.run(n.args[0]):
                self.run(n.args[1])

        elif n.type == N.IF:
            if self.run(n.args[0]):
                self.run(n.args[1])
            elif n.args[2] is not None:
                self.run(n.args[2])

        # Expression

        elif n.type == N.ADD:
            return sum(self.runs(n))

        elif n.type == N.SUB:
            x, y = self.runs(n)
            return x - y

        elif n.type == N.MUL:
            prod = 1
            for x in self.runs(n):
                prod *= x
            return prod

        elif n.type == N.DIV:
            x, y = self.runs(n)
            return x / y

        elif n.type == N.MINUS:
            x, = self.runs(n)
            return -x

        elif n.type == N.EXP:
            x, y = self.runs(n)
            return x ** y

        elif n.type == N.VAR:
            var = n.args[0]
            if var == "x":
                return self.x
            if var == "y":
                return self.y
            if var == "theta":
                return self.theta
            return self.vars[var]

        elif n.type == N.CONST:
            return n.value

        elif n.type == N.FUN:
            arg = self.run(n.args[1])

            f = {
                "Abs":     abs,
                "Rnd":     None,
                "sinh":    math.sinh,
                "cosh":    math.cosh,
                "tanh":    math.tanh,
                "asinh":   math.asinh,
                "acosh":   math.acosh,
                "atanh":   math.atanh,
                "sqrt":    math.sqrt,
                "log":     math.log,
                "cbrt":    lambda x: x ** (1/3),
                "sin":     math.sin,
                "cos":     math.cos,
                "tan":     math.tan,
                "asin":    math.asin,
                "acos":    math.acos,
                "atan":    math.atan,
                "log10":   math.log10,
                "Ent":     lambda x: float(int(x)),
                "EntEx":   lambda x: float(int(x)) - (x < 0),
                "RanInt":  None,
                "GCD":     None,
                "LCM":     None,
                "Arond":   None,
            }[n.args[0]]

            if f is None:
                raise Exception(
                    "Function {} not there yet x_x".format(n.args[0]))
            return f(arg)

        elif n.type == N.REL:
            left  = self.run(n.args[0])
            right = self.run(n.args[2])

            return {
                ">":  left > right,
                "<":  left < right,
                ">=": left >= right,
                "<=": left <= right,
                "=":  left == right,
                "!=": left != right
            }[n.args[1]]

        else:
            t = N(n.type).name
            raise Exception("Unhandled node of type {}".format(t))

    def runs(self, n):
        """Evaluate all the arguments of a node."""
        return [self.run(e) for e in n.args]