#!/usr/bin/python3
# -*- coding: utf-8 -*-
"""A fairly minimal interactive, responsive computer algebra system.

This is intended to be similar to calc.py, but instead of manipulating
scalars and vectors with APL-style broadcasting, we manipulate
algebraic expressions.  The particular thing I was doing that
motivated this was that I wanted to know, given that u⃗ = |x⃗|²x⃗, what
was the (forward) difference u⃗₁ - u⃗₀ when x⃗₁ = x⃗₀ + [1 0]*.  I worked
it out on paper but I thought it would be a great deal more pleasant
to do it in the fashion of calc.py.

It’s not quite usable; it lacks division entirely, but also
abstraction, many important simplifications, all simplifications after
substitution, and factoring.  Also, though, imperative Python with
isinstance is clearly the wrong level at which to write the rewrite
rules I need.  Also I think it would be helpful to have a canonical
form in which x, 1x, x¹, and x + 0 are all the same for purposes of
pattern matching (so you can match x against (const)x or x + (const)
or x**(const)).  But I did get this doomed approach to the point where
I could reproduce my original paper calculation of derivatives!

(However I also tried calculating the first forward differences when
you increment x, and that was totally not a usable experience with
this program.)

If you are going to be applying algebraic manipulations step by step,
though, you really need a way to point at subexpressions inside an
expression.  I think from a keyboard a two-dimensional navigation
style — left and right to move among siblings, up and down to move in
and out — could maybe work.

The key dispatch also leaves a lot to be desired; something IMGUI-like
would eliminate a lot of duplication!

"""
from __future__ import print_function
import os, sys, pickle


def ctrl(ch):
    return chr(0x1f & ord(ch))

class State:
    def __init__(self, filename):
        self.stack = []
        self.vars = {}
        self.deleted = []
        self.done = False
        self.messages = []
        self.filename = filename
        if os.path.exists(filename):
            with open(filename, 'rb') as fo:
                self.stack, self.vars = pickle.load(fo)

    def store(self, var, value):
        self.vars[var] = value

    def fetch(self, var):
        return self.vars[var]

    def has_var(self, var):
        return var in self.vars

    def depth(self):
        return len(self.stack)

    def pop(self):
        return self.stack.pop()

    def push(self, expr):
        self.stack.append(expr)

    def top(self):
        return self.stack[-1]

    def redraw(self):
        print("\033[0H\033[0J", end='')
        if self.vars:
            print('vars:', ' '.join(sorted(var.name for var in self.vars)))
        else:
            print('(no vars defined with ! yet)')

        for expr in self.stack:
            print(expr.display())

        if self.depth() > 0:
            print(self.top().menu(self))

        print(self.generic_menu())
        for message in self.messages:
            print(message)

        self.messages = []
        sys.stdout.flush()

    def generic_menu(self):
        items = ['_ negative']
        if self.depth() > 0:
            items.append('backspace deletes')
        if self.depth() > 1:
            items.append('* multiply  + add  - subtract  ^ exponentiate  s^Wap')
        if self.deleted:
            items.append('^Yundelete')

        items.append('^Quit  or type a number or var')
        return '  '.join(items)

    def dispatch_key(self, key):
        new_expr = None
        if self.depth() > 0:
            new_expr = self.top().dispatch(self, key)

        if new_expr:
            self.stack[-1] = new_expr
            return

        # XXX this part is kind of bogus probably since it sort of
        # implies that you can edit vars and consts on the top of the
        # stack any time

        if 'a' <= key <= 'z':
            self.push(Var(key))
            return

        if key in '0123456789._':
            self.push(Const(key))
            return

        if key == ctrl('Q'):
            self.done = True
            return

        if key in [ctrl('H'), '\177'] and self.depth() > 0:
            self.deleted.append(self.pop())
            return

        if key == ctrl('Y') and self.deleted:
            self.push(self.deleted.pop())

        if key == ctrl('W') and self.depth() >= 2:
            x = self.pop()
            y = self.pop()
            self.push(x)
            self.push(y)
            return

        self.messages.append('(unexpected %r)' % key)

    def save(self):
        with open(self.filename + '.tmp', 'wb') as fo:
            pickle.dump((self.stack, self.vars), fo)
            os.fsync(fo.fileno())

        os.rename(self.filename + '.tmp', self.filename)

class Var:
    def __init__(self, name, editing=True):
        self.name = name
        self.editing = editing

    def __eq__(self, other):
        return isinstance(other, Var) and self.name == other.name

    def __hash__(self):
        return hash(self.name) + 3

    def display(self):
        return self.expr(0, 0)

    def expr(self, lprec, rprec):
        return self.name

    def menu(self, state):
        items = ['(entering variable)' if self.editing else '(variable)']
        if state.has_var(self):
            items.append('@fetch')
        if state.depth() > 2:
            items.append('^Substitute')
        if state.depth() > 1 and self in state.stack[-2]:
            items.append('^Differentiate')
        if state.depth() > 1:
            items.append('!store')
        return '  '.join(items)

    def freeze(self):
        return Var(self.name, editing=False)

    def dispatch(self, state, key):
        if self.editing and ('a' <= key <= 'z' or '0' <= key <= '9'):
            return Var(self.name + key)

        if self.editing and key in (' ', ctrl('J')):
            return self.freeze()

        if self.editing and key in (ctrl('H'), '\177'):
            name = self.name[:-1]
            if not name:
                return None     # let the top-level backspace handler delete it

            return Var(name)

        if state.depth() > 2 and key == ctrl('S'):
            me = state.pop()
            replacement = state.pop()
            victim = state.top()
            # kind of bogus but we don’t want the original to go away
            # because there’s not really any undo
            state.push(replacement)
            return victim.subst(me, replacement)

        if state.depth() > 1 and key == ctrl('D') and self in state.stack[-2]:
            # Similarly, here we also don’t want the original to go away
            return state.stack[-2].derivative(self)

        if state.depth() >= 2 and key == '!':
            me = state.pop()
            state.store(self, state.top())
            return state.top()

        if key == '@' and state.has_var(self):
            return state.fetch(self)

        return generic_expr_dispatch(state, key)

    def subst(self, var, replacement):
        return replacement if self == var else self

    def __contains__(self, other):
        return self == other

    def derivative(self, other):
        return Const('1') if self == other else Const('0')


class Const:
    def __init__(self, text, editing=True):
        self.text = '-' if text == '_' else text  # XXX '-' by itself is not okay
        self.editing = editing

    def __eq__(self, other):
        return isinstance(other, Const) and self.eval({}) == other.eval({})

    def __hash__(self):
        return hash(self.eval({})) + 4

    def expr(self, lprec, rprec):
        return self.text

    def display(self):
        return self.expr(0, 0)

    def menu(self, state):
        return '(entering constant)' if self.editing else '(constant)'

    def freeze(self):
        return Const(self.text, editing=False)

    def dispatch(self, state, key):
        if self.editing:
            if '0' <= key <= '9':
                return Const(self.text + key)

            if key == '_':
                if self.text.startswith('-'):
                    return Const(self.text[1:])
                else:
                    return Const('-' + self.text)

            if key == '.' and '.' not in self.text:
                    return Const(self.text + '.')

            if key in (' ', ctrl('J')):
                return self.freeze()

            if key in (ctrl('h'), '\177'):
                text = self.text[:-1]
                if not text:
                    return None

                return Const(text)

        return generic_expr_dispatch(state, key)

    def eval(self, context):
        return int(self.text) if '.' not in self.text else float(self.text)

    def subst(self, var, replacement):
        return self

    def __contains__(self, other):
        return self == other

    def derivative(self, other):
        return Const('0')


class Binop:
    def __init__(self, x, y):
        self.kids = x, y

    def __eq__(self, other):
        return isinstance(other, self.__class__) and other.kids == self.kids

    def __hash__(self):
        return hash(self.kids) + 5

    def display(self):
        return self.expr(0, 0)

    def menu(self, state):
        return 'd^Estructure'

    def freeze(self):
        return self

    def dispatch(self, state, key):
        return self.generic_binop_dispatch(state, key)

    def generic_binop_dispatch(self, state, key):
        if key == ctrl('E'):
            state.pop()
            state.push(self.kids[0])
            state.push(self.kids[1]) # XXX bogus
            return self.kids[1]

        return generic_expr_dispatch(state, key)

    def expr(self, lprec, rprec):
        return self.generic_binop_expr(lprec, rprec)

    def generic_binop_expr(self, lprec, rprec):
        x, y = self.kids
        xe, ye = x.expr(lprec, self.prec), y.expr(self.prec, rprec)
        e = '%s %s %s' % (xe, self.op, ye)
        return e if self.prec >= lprec and self.prec > rprec else '(%s)' % e

    def subst(self, var, replacement):
        return self.__class__(self.kids[0].subst(var, replacement),
                              self.kids[1].subst(var, replacement))

    def __contains__(self, other):
        return any(other in kid for kid in self.kids)


class Product(Binop):
    op = '·'
    prec = 2

    def menu(self, state):
        s = ['d^Estructure']
        if isinstance(self.kids[0], Sum):
            s.append('^Left-distribute')
        if isinstance(self.kids[1], Sum):
            s.append('^Right-distribute')

        return '  '.join(s)

    def dispatch(self, state, key):
        if key == ctrl('L') and isinstance(self.kids[0], Sum):
            return distribute_left(self.kids[0], self.kids[1])

        if key == ctrl('R') and isinstance(self.kids[1], Sum):
            return distribute_right(self.kids[0], self.kids[1])

        return self.generic_binop_dispatch(state, key)

    def derivative(self, var):
        return add(multiply(self.kids[0].derivative(var), self.kids[1]),
                   multiply(self.kids[0], self.kids[1].derivative(var)))


def multiply(x, y):
    x, y = x.freeze(), y.freeze()
    if isinstance(x, Const) and isinstance(y, Const):
        return Const(str(x.eval({}) * y.eval({})), editing=False)

    if isinstance(x, Product):
        return Product(x.kids[0], multiply(x.kids[1], y))

    if isinstance(y, Product) and isinstance(x, Const) and isinstance(y.kids[0], Const):
        return multiply(multiply(x, y.kids[0]), y.kids[1])

    if x == y:
        return exponentiate(x, Const('2'))

    if isinstance(x, Power) and x.kids[0] == y:
        return exponentiate(y, add(x.kids[1], Const('1')))

    if y == Const('1'):
        return x

    if x == Const('1'):
        return y

    if x == Const('0'):
        return x

    if y == Const('0'):
        return y

    return Product(x, y)

def distribute_left(x, y):
    if isinstance(x, Sum):
        return add(multiply(x.kids[0], y), distribute_left(x.kids[1], y))
    else:
        return multiply(x, y)

def distribute_right(x, y):
    if isinstance(y, Sum):
        return add(multiply(x, y.kids[0]), distribute_right(x, y.kids[1]))
    else:
        return multiply(x, y)


class Power(Binop):
    op = '**'
    prec = 3

    def menu(self, state):
        items = ['d^Estructure']
        if isinstance(self.kids[1], Const) and isinstance(self.kids[1].eval({}), int) and isinstance(self.kids[0], Sum):
            items.append('^Left-distribute  ^Right-distribute')
        return '  '.join(items)

    def dispatch(self, state, key):
        if isinstance(self.kids[1], Const) and isinstance(self.kids[1].eval({}), int) and isinstance(self.kids[0], Sum):
            if key == ctrl('L'):
                newexp = add(self.kids[1], Const('-1'))
                return add(distribute_left(self.kids[0].kids[0], exponentiate(self.kids[0], newexp)),
                           distribute_left(self.kids[0].kids[1], exponentiate(self.kids[0], newexp)))

            if key == ctrl('R'):
                newexp = add(self.kids[1], Const('-1'))
                return add(distribute_right(exponentiate(self.kids[0], newexp), self.kids[0].kids[0]),
                           distribute_left(exponentiate(self.kids[0], newexp), self.kids[0].kids[1]))

        return self.generic_binop_dispatch(state, key)
        

    def expr(self, lprec, rprec):
        if isinstance(self.kids[1], Const):
            return self.kids[0].expr(lprec, self.prec) + superscript(self.kids[1].expr(0, 0))

        return self.generic_binop_expr(lprec, rprec)

    def derivative(self, var):
        if var not in self:
            return Const('0')
        
        if var == self.kids[0] and isinstance(self.kids[1], Const):
            return multiply(self.kids[1], exponentiate(self.kids[0], add(self.kids[1], Const('-1'))))

        raise ValueError('Derivatives not implemented for general powers', self)

superdigits = {
    '.': '·',
    '0': '⁰',
    '1': '¹',
    '2': '²',
    '3': '³',
    '4': '⁴',
    '5': '⁵',
    '6': '⁶',
    '7': '⁷',
    '8': '⁸',
    '9': '⁹',
    '-': '⁻',
}

def superscript(text):
    return ''.join(superdigits[c] for c in text)

def exponentiate(x, y):
    x, y = x.freeze(), y.freeze()
    if isinstance(x, Const) and isinstance(y, Const):
        return Const(str(x.eval({}) ** y.eval({})), editing=False)

    if isinstance(x, Power):
        return exponentiate(x.kids[0], multiply(x.kids[1], y))

    if y == Const('1'):
        return x

    if y == Const('0'):
        return Const('1')

    return Power(x, y)


class Sum(Binop):
    op = '+'
    prec = 1

    def derivative(self, var):
        return add(self.kids[0].derivative(var),
                   self.kids[1].derivative(var))

def add(x, y):
    x, y = x.freeze(), y.freeze()
    if isinstance(x, Const) and isinstance(y, Const):
        return Const(str(x.eval({}) + y.eval({})), editing=False)

    if isinstance(x, Sum):
        return Sum(x.kids[0], add(x.kids[1], y))

    if isinstance(y, Sum) and isinstance(x, Const) and isinstance(y.kids[0], Const):
        return add(add(x, y.kids[0]), y.kids[1])

    if isinstance(y, Sum) and x == y.kids[0]:
        return add(multiply(Const('2'), x), y.kids[1])

    if isinstance(y, Sum) and isinstance(y.kids[0], Product) and isinstance(y.kids[0].kids[0], Const) and x == y.kids[0].kids[1]:
        return add(multiply(add(Const('1'), y.kids[0].kids[0]), x), y.kids[1])

    # XXX 3x + 2x + 1 not handled

    if isinstance(x, Product) and x.kids[1] == y and isinstance(x.kids[0], Const):
        return multiply(add(x.kids[0], Const('1')), y)

    if isinstance(y, Product) and y.kids[1] == x and isinstance(y.kids[0], Const):
        return multiply(add(y.kids[0], Const('1')), x)

    if isinstance(x, Product) and isinstance(y, Product) and x.kids[1] == y.kids[1] and isinstance(x.kids[0], Const) and isinstance(y.kids[0], Const):
        return multiply(add(x.kids[0], y.kids[0]), x.kids[1])

    if x == y:
        return multiply(Const('2'), x)

    if x == Const('0'):
        return y

    if y == Const('0'):
        return x

    return Sum(x, y)


def generic_expr_dispatch(state, key):
    if state.depth() >= 2:
        if key == '*':
            y = state.pop()
            return multiply(state.top(), y)
        if key == '^':
            y = state.pop()
            return exponentiate(state.top(), y)
        if key == '+':
            y = state.pop()
            return add(state.top(), y)
        if key == '-':
            y = state.pop()
            return add(state.top(), multiply(Const('-1'), y))

def repl():
    state = State('.alg')
    while not state.done:
        state.redraw()
        k = sys.stdin.read(1)
        state.dispatch_key(k)
        state.save()

def main():
    state = os.popen('stty -g').read().strip()
    os.system('stty cbreak -ixon')
    try:
        repl()
    finally:
        os.system('stty ' + state)

if __name__ == '__main__':
    main()