from __future__ import annotations
import ast
import copy
import random

def _is_pure(node: ast.AST) -> bool:
    if isinstance(node, ast.Name):
        return True
    if isinstance(node, ast.Constant):
        return isinstance(node.value, (int, bool))
    if isinstance(node, ast.Attribute):
        return _is_pure(node.value)
    if isinstance(node, ast.Subscript):
        return _is_pure(node.value) and _is_pure(node.slice)
    if isinstance(node, ast.UnaryOp) and isinstance(node.op, (ast.USub, ast.UAdd, ast.Invert)):
        return _is_pure(node.operand)
    return False

def _dup(node: ast.AST) -> ast.AST:
    return copy.deepcopy(node)

class MBATransformer(ast.NodeTransformer):

    def __init__(self, rng: random.Random, prob: float=0.7):
        self.rng = rng
        self.prob = prob

    def visit_BinOp(self, node: ast.BinOp) -> ast.AST:
        self.generic_visit(node)
        if not (_is_pure(node.left) and _is_pure(node.right)):
            return node
        if self.rng.random() > self.prob:
            return node
        a, b = (node.left, node.right)
        op = node.op
        if isinstance(op, ast.BitXor):
            return self._xor(a, b)
        if isinstance(op, ast.BitOr):
            return self._or(a, b)
        if isinstance(op, ast.BitAnd):
            return self._and(a, b)
        if isinstance(op, ast.LShift) and isinstance(b, ast.Constant) and isinstance(b.value, int):
            return self._lshift(a, b.value)
        if isinstance(op, ast.RShift) and isinstance(b, ast.Constant) and isinstance(b.value, int):
            return self._rshift(a, b.value)
        return node

    def visit_UnaryOp(self, node: ast.UnaryOp) -> ast.AST:
        self.generic_visit(node)
        if not _is_pure(node.operand):
            return node
        if self.rng.random() > self.prob:
            return node
        if isinstance(node.op, ast.Invert):
            return self._invert(node.operand)
        return node

    def _xor(self, a: ast.AST, b: ast.AST) -> ast.AST:
        choice = self.rng.choice(('subor', 'ornotand', 'addsub'))
        if choice == 'subor':
            return ast.BinOp(left=ast.BinOp(left=_dup(a), op=ast.BitOr(), right=_dup(b)), op=ast.Sub(), right=ast.BinOp(left=_dup(a), op=ast.BitAnd(), right=_dup(b)))
        if choice == 'ornotand':
            return ast.BinOp(left=ast.BinOp(left=_dup(a), op=ast.BitOr(), right=_dup(b)), op=ast.BitAnd(), right=ast.UnaryOp(op=ast.Invert(), operand=ast.BinOp(left=_dup(a), op=ast.BitAnd(), right=_dup(b))))
        return ast.BinOp(left=ast.BinOp(left=_dup(a), op=ast.Add(), right=_dup(b)), op=ast.Sub(), right=ast.BinOp(left=ast.Constant(value=2), op=ast.Mult(), right=ast.BinOp(left=_dup(a), op=ast.BitAnd(), right=_dup(b))))

    def _or(self, a: ast.AST, b: ast.AST) -> ast.AST:
        if self.rng.random() < 0.5:
            return ast.BinOp(left=ast.BinOp(left=_dup(a), op=ast.Add(), right=_dup(b)), op=ast.Sub(), right=ast.BinOp(left=_dup(a), op=ast.BitAnd(), right=_dup(b)))
        return ast.BinOp(left=ast.BinOp(left=_dup(a), op=ast.BitXor(), right=_dup(b)), op=ast.BitXor(), right=ast.BinOp(left=_dup(a), op=ast.BitAnd(), right=_dup(b)))

    def _and(self, a: ast.AST, b: ast.AST) -> ast.AST:
        if self.rng.random() < 0.5:
            return ast.BinOp(left=ast.BinOp(left=_dup(a), op=ast.Add(), right=_dup(b)), op=ast.Sub(), right=ast.BinOp(left=_dup(a), op=ast.BitOr(), right=_dup(b)))
        return ast.UnaryOp(op=ast.Invert(), operand=ast.BinOp(left=ast.UnaryOp(op=ast.Invert(), operand=_dup(a)), op=ast.BitOr(), right=ast.UnaryOp(op=ast.Invert(), operand=_dup(b))))

    def _invert(self, a: ast.AST) -> ast.AST:
        return ast.BinOp(left=ast.UnaryOp(op=ast.USub(), operand=_dup(a)), op=ast.Sub(), right=ast.Constant(value=1))

    def _lshift(self, a: ast.AST, n: int) -> ast.AST:
        if n < 0 or n > 60:
            return ast.BinOp(left=_dup(a), op=ast.LShift(), right=ast.Constant(value=n))
        return ast.BinOp(left=_dup(a), op=ast.Mult(), right=ast.Constant(value=1 << n))

    def _rshift(self, a: ast.AST, n: int) -> ast.AST:
        if n < 0 or n > 60:
            return ast.BinOp(left=_dup(a), op=ast.RShift(), right=ast.Constant(value=n))
        return ast.BinOp(left=_dup(a), op=ast.FloorDiv(), right=ast.Constant(value=1 << n))

    def visit_match_case(self, node: ast.match_case) -> ast.AST:
        if node.guard is not None:
            node.guard = self.visit(node.guard)
        node.body = [self.visit(stmt) for stmt in node.body]
        return node