#!/usr/bin/python3
# -*- coding: utf-8 -*-
"""A simple calculator with vectors."""
import os, sys, operator, errno, pickle, math


def ctrl(ch):
    return chr(ord(ch) & 31)

infix_ops = {
    '+': operator.add,
    '-': operator.sub,
    '*': operator.mul,
    '/': operator.truediv,
    '^': operator.pow,
    '%': operator.mod,
    ':': lambda a, b: Vector(range(a[0], b[0])),
    ']': lambda a, b: a.max(b),
    '[': lambda a, b: a.min(b),
    '<': operator.lt,
    '>': operator.gt,
    '@': operator.getitem,
}

unary_ops = {
    ctrl('e'): ('exp', math.exp),
    ctrl('l'): ('ln', math.log),
    ctrl('s'): ('sin', math.sin),
    ctrl('k'): ('cos', math.cos),
    ctrl('a'): ('atan', math.atan),
}

class Prompt:
    def __init__(self):
        self.contents = []

    def is_var(self):
        return any('a' <= c <= 'z' for c in self.contents)

    def is_const(self):
        return any(c in '. ' for c in self.contents)

    def show(self):
        return '> ' + ''.join(self.contents) + '   ' + ctrl('H') * 3

    def flush(self, stack):
        if not self.contents:
            return

        stack.append((Var if self.is_var() else Constant)(''.join(self.contents)))
        self.contents = []

    def dispatch(self, char, stack, vars):
        if '0' <= char <= '9':
            self.contents.append(char)

        elif char == '_':
            self.contents.append('-')

        elif char in '. ':
            if not self.is_var():
                self.contents.append(char)

        elif 'a' <= char <= 'z':
            if not self.is_const():
                self.contents.append(char)

        elif char in {ctrl('H'), '\177'}:
            if self.contents:
                self.contents = self.contents[:-1]
            else:
                if not stack:
                    return

                item = stack.pop()
                stack.extend(item.decompose(vars))

        elif char == '\n':
            self.flush(stack)

        elif char == '\t':
            self.flush(stack)

            if len(stack) < 2:
                return

            a = stack.pop()
            b = stack.pop()
            stack.append(a)
            stack.append(b)

        elif char in infix_ops:
            self.flush(stack)

            if len(stack) < 2:
                return

            b = stack.pop()
            a = stack.pop()
            stack.append(Binop(char, a, b))

        elif char in unary_ops:
            self.flush(stack)
            if not stack:
                return

            stack.append(Unop(unary_ops[char], stack.pop()))

        elif char == '!':
            self.flush(stack)
            if len(stack) < 2:
                return

            name = stack.pop()
            val = stack.pop()
            vars[name.name] = val  # XXX

        elif char == ctrl('r'):
            self.flush(stack)
            stack.append(ReduceFormula())

        elif char == ctrl('t'):
            self.flush(stack)
            stack.append(ScanFormula())

        elif char == ctrl('p'):
            self.flush(stack)
            stack.append(PairsFormula())
        

class Formula:
    def show(self, vars):
        try:
            val = self.eval(vars).show()
        except Exception as e:
            val = repr(e)

        return '%s = %s' % (self.infix(), val)


class Constant(Formula):
    def __init__(self, s):
        self.value = Vector([(float if '.' in s else int)(n) for n in s.split()])

    def eval(self, vars):
        return self.value

    def show(self, vars):
        return self.value.show()

    def infix(self):
        return self.value.show()

    def decompose(self, vars):
        return []


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

    def eval(self, vars):
        return vars[self.name].eval(vars)

    def infix(self):
        return self.name

    def decompose(self, vars):
        return ([vars[self.name]] if self.name in vars else [])

class ReduceFormula(Formula):
    def eval(self, vars):
        return ReduceValue()

    def infix(self):
        return '(reduce)'

    def decompose(self, vars):
        return []

class ScanFormula(Formula):
    def eval(self, vars):
        return ScanValue()

    def infix(self):
        return '(scan)'

    def decompose(self, vars):
        return []

class PairsFormula(Formula):
    def eval(self, vars):
        return PairsValue()

    def infix(self):
        return '(pairs)'

    def decompose(self, vars):
        return []

class Binop(Formula):
    def __init__(self, op, a, b):
        self.op, self.a, self.b = op, a, b

    def eval(self, vars):
        a = self.a.eval(vars)
        b = self.b.eval(vars)
        return infix_ops[self.op](a, b)

    def decompose(self, vars):
        return [self.a, self.b]

    def infix(self):
        return '(%s %s %s)' % (self.a.infix(), self.op, self.b.infix())

class Unop(Formula):
    def __init__(self, op, arg):
        self.name, self.f = op
        self.arg = arg

    def eval(self, vars):
        arg = self.arg.eval(vars)
        return arg.apply(self.f)

    def decompose(self, vars):
        return [self.arg]

    def infix(self):
        return '%s(%s)' % (self.name, self.arg.infix())


class ScanValue:
    def show(self):
        return 'scan'

    def magic_broadcast(self, op, vec):
        result = [vec[0]]
        for item in vec[1:]:
            result.append(op(result[-1], item))

        return Vector(result)

class ReduceValue:
    def show(self):
        return 'reduce'

    def magic_broadcast(self, op, vec):
        result = vec[0]
        for item in vec[1:]:
            result = op(result, item)

        return Vector([result])

class PairsValue:
    def show(self):
        return 'pairs'

    def magic_broadcast(self, op, vec):
        return Vector([vec[0]] + [op(vec[i], vec[i+1])
                                  for i in range(len(vec)-1)])

class Vector:
    def __init__(self, contents):
        self.contents = list(contents)

    def show(self):
        return ' '.join(str(n) for n in self.contents)

    def __add__(self, other):
        return self.broadcast(operator.add, other)

    def __sub__(self, other):
        return self.broadcast(operator.sub, other)
    
    def __mul__(self, other):
        return self.broadcast(operator.mul, other)

    def __truediv__(self, other):
        return self.broadcast(operator.truediv, other)

    def __pow__(self, other):
        return self.broadcast(operator.pow, other)

    def __mod__(self, other):
        return self.broadcast(operator.mod, other)

    def min(self, other):
        return self.broadcast(min, other)
    
    def max(self, other):
        return self.broadcast(max, other)

    def __gt__(self, other):
        return self.broadcast(operator.gt, other)

    def __lt__(self, other):
        return self.broadcast(operator.lt, other)

    def broadcast(self, op, other):
        if hasattr(other, 'magic_broadcast'):
            return other.magic_broadcast(op, self)

        if len(other) == 1:
            return Vector([op(item, other[0]) for item in self])
        elif len(self) == 1:
            return Vector([op(self[0], item) for item in other])
        elif len(self) == len(other):
            return Vector([op(a, b) for a, b in zip(self, other)])
        else:
            raise IndexingError(self, other)

    def apply(self, unop):
        return Vector([unop(x) for x in self])

    def __iter__(self):
        return iter(self.contents)

    def __len__(self):
        return len(self.contents)

    def __getitem__(self, item):
        if isinstance(item, Vector):
            return Vector([self[i] for i in item])

        return self.contents[item]

    def __repr__(self):
        return 'Vector(%s)' % self.contents


class IndexingError(Exception): pass


def main():
    stack = []
    vars = {}
    prompt = Prompt()

    try:
        with open('.calcpickle', 'rb') as f:
            stack, vars = pickle.load(f)
    except IOError as e:
        if e.errno != errno.ENOENT:
            raise

    while True:
        print("\033[2J\033[H" + # ANSI clear and home escape sequences
              '^Exp ^Ln ^Sin ^Kos ^Atan ^Reduce ^Tscan ^Pairs [min ]max @index :range')
        print('vars:', ' '.join(sorted(vars)))
        for item in stack[-23:]:
            print(item.show(vars))

        print(prompt.show(), end='')
        sys.stdout.flush()
        prompt.dispatch(sys.stdin.read(1), stack, vars)

        with open('.calcpickle.tmp', 'wb') as f:
            pickle.dump((stack, vars), f)
        os.rename('.calcpickle.tmp', '.calcpickle')

if __name__ == '__main__':
    settings = os.popen('stty -g').read().strip()
    os.system('stty cbreak -ixon')
    try:
        main()
    finally:
        os.spawnv(os.P_WAIT, '/bin/stty', ['stty', settings])