#!/usr/bin/env python
"""Usage: `python assessment2.py {subcommand} {expressions}`

If expressions are not provided, they are taken from each line of stdin.
Documentation of subcommands is available from `python assessment2.py -h`.

Examples for marking criteria:
  1. Eval: `python assessment2.py eval "(((2*(3+2))+5)/2)"`
  2. Tree: `python assessment2.py render-tree "(((2*(3+2))+5)/2)"`
  3. Preorder:  `python assessment2.py format-prefix "(((2*(3+2))+5)/2)"`
     In-order:  `python assessment2.py format-infix "(((2*(3+2))+5)/2)"`
     Postorder: `python assessment2.py format-postfix "(((2*(3+2))+5)/2)"`
  4. Errors: `python assessment2.py eval "(4*3*2)" "(4*(2))" "(4*(3+2)*(2+1))" "(2*4)*(3+2)" "((2+3)*(4*5)" "((2+3)*(4*5)" "(2+5)*(4/(2+2)))" "(2+5)*(4/(2+2)))" "(((2+3)*(4*5))+(1(2+3)))"`
  5. Tests: `python assessment2.py test`

Tests can be run with `pytest assessment2.py` or `python assessment2.py test`.

All of the formatting is implemented with post-order traversals. This is
necessary to output parenthesis in the pre- and in-order cases. It could also be
done with a buffer and a combined traversal (with callbacks before, between, and
after subtree visits), but that requires mutation.

Note: this program requires Python 3.12 due to the generic syntax used to define
`BTree.traverse`. Pytest is also required. These should be installed already
on Codio.
"""
# stdlib
from dataclasses import dataclass
from typing import Literal, Tuple, Callable, cast, override
from operator import add, sub, mul, truediv as div
from argparse import ArgumentParser
import sys
import re
from textwrap import dedent

# pytest, from https://pypi.org/project/pytest/
# authors can be found at https://github.com/pytest-dev/pytest/graphs/contributors,
import pytest

OPERATORS: dict[Literal["+", "-", "*", "/"], Callable[[float, float], float]] = {
    "+": add,
    "-": sub,
    "*": mul,
    "/": div,
}


class BTree:
    value: float

    def traverse[T](
        self,
        value: Callable[["Value", int], T],
        operator: Callable[["Operator", T, T, int], T],
        depth: int = 0,
    ) -> T:
        """Traverse the binary tree with a post-order traversal

        `value` takes the value node and its depth from the root node
        `operator` takes the operator node, the results of traversing the left
            and right subtrees, and the depth of the node

        The tree could, for example, be evaluated with `traverse`:
            ```
                (...).traverse(
                    lambda node, _: node.value,
                    lambda node, left, right, _: OPERATORS[node.operator](left, right)
                )
            ```
        """
        raise NotImplementedError

    def as_preorder_str(self) -> str:
        """Format the tree in prefix notation"""
        return self.traverse(
            lambda value, _: str(value.value),
            lambda op, left, right, _: f"({op.operator} {left} {right})",
        )

    def as_inorder_str(self) -> str:
        """Format the tree as a typical parenthesized expression"""
        return self.traverse(
            lambda value, _: str(value.value),
            lambda op, left, right, _: f"({left} {op.operator} {right})",
        )

    __str__ = as_inorder_str

    def as_inorder_lines(self) -> str:
        """Format the tree, visually as a tree"""
        return self.traverse(
            lambda value, depth: " " * depth + str(value.value),
            lambda op, left, right, depth: f"{left}\n{' ' * depth}{op.operator}\n{right}",
        )

    def as_postorder_str(self) -> str:
        """Format the tree in postfix notation (RPN)"""
        return self.traverse(
            lambda value, _: str(value.value),
            lambda op, left, right, _: f"{left} {right} {op.operator}",
        )


@dataclass(frozen=True, slots=True)
class Value(BTree):
    value: float

    @override
    def traverse[T](
        self,
        value: Callable[["Value", int], T],
        operator: Callable[["Operator", T, T, int], T],
        depth: int = 0,
    ) -> T:
        """Traverse the binary tree with a post-order traversal

        `value` takes the value node and its depth from the root node
        `operator` takes the operator node, the results of traversing the left
            and right subtrees, and the depth of the node

        The tree could, for example, be evaluated with `traverse`:
            ```
                (...).traverse(
                    lambda node, _: node.value,
                    lambda node, left, right, _: OPERATORS[node.operator](left, right)
                )
            ```
        """

        return value(self, depth)


@dataclass(frozen=True, slots=True)
class Operator(BTree):
    operator: Literal["+", "-", "*", "/"]
    left: BTree
    right: BTree

    @property
    def value(self) -> float:  # type: ignore - the field should be read-only
        return OPERATORS[self.operator](self.left.value, self.right.value)

    @override
    def traverse[T](
        self,
        value: Callable[["Value", int], T],
        operator: Callable[["Operator", T, T, int], T],
        depth: int = 0,
    ) -> T:
        """Traverse the binary tree with a post-order traversal

        `value` takes the value node and its depth from the root node
        `operator` takes the operator node, the results of traversing the left
            and right subtrees, and the depth of the node

        The tree could, for example, be evaluated with `traverse`:
            ```
                (...).traverse(
                    lambda node, _: node.value,
                    lambda node, left, right, _: OPERATORS[node.operator](left, right)
                )
            ```
        """

        return operator(
            self,
            self.left.traverse(value, operator, depth + 1),
            self.right.traverse(value, operator, depth + 1),
            depth,
        )


def parse(expr: str) -> BTree:
    """Parse a parenthesized expression as in the spec.

    Requires a single parenthesized string, potentially with whitespace
    Raises ValueError on parse errors
    """

    def number(expr: str) -> Tuple[BTree, str] | None:
        if not expr[0].isdigit():
            return None
        rest = expr.lstrip("0123456789")
        return Value(int(expr[: len(expr) - len(rest)])), rest.lstrip()

    def parenthesized(expr: str) -> Tuple[BTree, str] | None:
        if not expr.startswith("("):
            return None
        expr = expr[1:].lstrip()

        left, expr = operand(expr)

        operator, expr = expr[0], expr[1:]
        if operator not in OPERATORS:
            raise ValueError(f"Unknown operator {operator}")
        operator = cast(Literal["+", "-", "*", "/"], operator)

        right, expr = operand(expr.lstrip())

        if not expr.startswith(")"):
            if expr and expr[0] in OPERATORS:
                raise ValueError("Too many operands in expression")
            raise ValueError("Expected closing parenthesis")
        return Operator(operator, left, right), expr[1:].lstrip()

    def operand(expr: str) -> Tuple[BTree, str]:
        v = number(expr) or parenthesized(expr)
        if not v:
            raise ValueError(
                "Expected a number or the beginning of a parenthesized expression"
            )
        return v

    result = parenthesized(expr.lstrip())
    if not result or result[1]:
        # If rest is non-empty, then something was after the parenthesis starting the expression
        # i.e. the full expression was not parenthesized.
        raise ValueError("Expected parenthesized expression at the top level")

    return result[0]


@pytest.mark.parametrize(
    ("expr", "ast"),
    [
        ("(1 + 1)", Operator("+", Value(1), Value(1))),
        (
            "(2*(3+ 2))",
            Operator("*", Value(2), Operator("+", Value(3), Value(2))),
        ),
        (" (1 +1) ", Operator("+", Value(1), Value(1))),
    ],
)
def test_parse(expr: str, ast: BTree):
    """Test parsing of expressions"""
    assert parse(expr) == ast


@pytest.mark.parametrize(
    ("expr", "value"),
    [
        ("(1 + 2)", 3),
        ("(12 + (((4 / 2) * 3) + (3 * 2)))", 24),
        ("(((2 * (3 + 2)) + 5) / 2)", 7.5),
    ],
)
def test_eval(expr: str, value: float):
    """Test evaluation"""
    assert parse(expr).value == value


@pytest.mark.parametrize(
    ("expr", "tree"),
    [
        (
            "(1 + 2)",
            """
                 1
                +
                 2
            """,
        ),
        (
            "(1 + (((4 / 2) * 3) + (3 * 2)))",
            """
                 1
                +
                    4
                   /
                    2
                  *
                   3
                 +
                   3
                  *
                   2
            """,
        ),
        (
            "(((2*(3+2))+5)/2)",
            """
                   2
                  *
                    3
                   +
                    2
                 +
                  5
                /
                 2
            """
        )
    ],
)
def test_inorder_lines(expr: str, tree: str):
    """Test tree rendering"""
    assert parse(expr).as_inorder_lines() == dedent(tree).strip("\n")


@pytest.mark.parametrize(
    ("expr", "prefix"),
    [
        ("(1 + 2)", "(+ 1 2)"),
        ("(1 + (((4 / 2) * 3) + (3 * 2)))", "(+ 1 (+ (* (/ 4 2) 3) (* 3 2)))"),
    ],
)
def test_prefix(expr: str, prefix: str):
    """Test prefix formatting"""
    assert parse(expr).as_preorder_str() == prefix


@pytest.mark.parametrize(
    ("expr"),
    [
        "(1 + 2)",
        "(1 + (((4 / 2) * 3) + (3 * 2)))",
        "(((3 + 2) * 2) + (((7 / 3) * 5) - 3))",
    ],
)
def test_infix_roundtrip(expr: str):
    """Test infix formatting roundtrips"""
    assert str(parse(expr)) == expr


@pytest.mark.parametrize(
    ("expr", "postfix"),
    [
        ("(1 + 2)", "1 2 +"),
        ("(1 + (((4 / 2) * 3) + (3 * 2)))", "1 4 2 / 3 * 3 2 * + +"),
    ],
)
def test_postfix(expr: str, postfix: str):
    """Test postfix formatting"""
    assert parse(expr).as_postorder_str() == postfix


@pytest.mark.parametrize(
    ("expr", "err"),
    [
        ("(4 * 3 * 2)", "Too many operands in expression"),
        ("(4 * (2))", "Unknown operator )"),
        ("(4 * (3 + 2) * (2 + 1))", "Too many operands in expression"),
        ("(2 *4)*(3+2)", "Expected parenthesized expression at the top level"),
        ("(2+5)*(4/(2+2))", "Expected parenthesized expression at the top level"),
        ("(((2+3)*(4*5))+(1(2+3)))", "Unknown operator ("),
    ],
)
def test_error(expr: str, err: str):
    """Test error messages"""
    with pytest.raises(ValueError, match=re.escape(err)):
        parse(expr)


if __name__ == "__main__":
    parser = ArgumentParser(usage=__doc__)
    subcommands = parser.add_subparsers(title="Subcommands", required=True)

    def with_expressions(parser):
        parser.add_argument(
            "expressions",
            nargs="*",
            help="The expressions to operate on. If none are provided, operate on lines of stdin.",
        )
        return parser

    with_expressions(
        subcommands.add_parser("eval", help="Evaluate the expression")
    ).set_defaults(func=lambda expr: expr.value)
    with_expressions(
        subcommands.add_parser("render-tree", help="Render the expression as a tree")
    ).set_defaults(func=BTree.as_inorder_lines)
    with_expressions(
        subcommands.add_parser(
            "format-prefix", help="Format the expression in prefix notation"
        )
    ).set_defaults(func=BTree.as_preorder_str)
    with_expressions(
        subcommands.add_parser(
            "format-infix",
            help="Format the expression in typical parenthesized infix notation",
        )
    ).set_defaults(func=BTree.as_inorder_str)
    with_expressions(
        subcommands.add_parser(
            "format-postfix",
            help="Format the expression in postfix notation, i.e. as RPN",
        )
    ).set_defaults(func=BTree.as_postorder_str)
    subcommands.add_parser("test", help="Run tests").set_defaults(
        func=lambda: pytest.main([__file__])
    )
    args = parser.parse_args()
    if "expressions" in args:
        for expression in args.expressions or sys.stdin:
            if len(args.expressions) > 1:
                print(expression + ":")
            try:
                parsed = parse(expression)
            except ValueError as e:
                print(f"Error: {e}")
                continue
            print(args.func(parsed))
    else:
        args.func()