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

from 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.pen = False

        self.angle = float(0)
        self.sin = float(0)
        self.cos = float(1)

        # TODO: Also specify the style!

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

    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)
            dx = self.cos * dist
            dy = self.sin * dist

            if self.pen:
                coord = self.x, self.y, self.x + dx, self.y + dy
                coord = [int(round(x)) for x in coord]
                x1, y1, x2, y2 = coord

                x1 += self.w.width // 2
                x2 += self.w.width // 2
                y1 = self.w.height // 2 - y1
                y2 = self.w.height // 2 - y2

                self.w.line(x1, y1, x2, y2, self.w.BLACK)
                self.w.update()

            self.x += dx
            self.y += dy

        elif n.type == N.ROTATE:
            angle, = self.runs(n)
            self.angle += angle
            self.sin = math.sin(math.pi * self.angle / 180)
            self.cos = math.cos(math.pi * self.angle / 180)

        elif n.type == N.ORIENT:
            angle, = self.runs(n)
            self.angle = angle
            self.sin = math.sin(math.pi * self.angle / 180)
            self.cos = math.cos(math.pi * self.angle / 180)

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

            if self.pen:
                coord = self.x, self.y, x, y
                coord = [int(round(x)) for x in coord]
                x1, y1, x2, y2 = coord

                x1 += self.w.width // 2
                x2 += self.w.width // 2
                y1 = self.w.height // 2 - y1
                y2 = self.w.height // 2 - y2

                self.w.line(x1, y1, x2, y2, 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]

            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(f"x={self.x} y={self.y} theta={self.angle}")
            print(f"A={self.vars['A']} B={self.vars['B']}")
            print(f"C={self.vars['C']} D={self.vars['D']}")
            print(f"E={self.vars['E']} F={self.vars['F']}")
            print(f"M={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
            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]