from pyparsing import *
from base64 import b64decode
import pprint
import math
import statistics
import sys

#
# Code for parsing s expressions (not mine)
#
# Copyright 2007-2011, by Paul McGuire
#

def verifyLen(s,l,t):
    t = t[0]
    if t.len is not None:
        t1len = len(t[1])
        if t1len != t.len:
            raise ParseFatalException(s,l,\
                    "invalid data of length %d, expected %s" % (t1len, t.len))
    return t[1]

# define punctuation literals
LPAR, RPAR, LBRK, RBRK, LBRC, RBRC, VBAR = map(Suppress, "()[]{}|")

decimal = Regex(r'0|[1-9]\d*').setParseAction(lambda t: int(t[0]))
hexadecimal = ("#" + OneOrMore(Word(hexnums)) + "#")\
                .setParseAction(lambda t: int("".join(t[1:-1]),16))
bytes = Word(printables)
raw = Group(decimal("len") + Suppress(":") + bytes).setParseAction(verifyLen)
token = Word(alphanums + "-./_:*+=")
base64_ = Group(Optional(decimal|hexadecimal,default=None)("len") + VBAR
    + OneOrMore(Word( alphanums +"+/=" )).setParseAction(lambda t: b64decode("".join(t)))
    + VBAR).setParseAction(verifyLen)

qString = Group(Optional(decimal,default=None)("len") +
                        dblQuotedString.setParseAction(removeQuotes)).setParseAction(verifyLen)
simpleString = base64_ | raw | decimal | token | hexadecimal | qString

# extended definitions
decimal = Regex(r'-?0|[1-9]\d*').setParseAction(lambda t: int(t[0]))
real = Regex(r"[+-]?\d+\.\d*([eE][+-]?\d+)?").setParseAction(lambda tokens: float(tokens[0]))
token = Word(alphanums + "-./_:*+=!<>")

simpleString = (real | base64_ | raw | decimal | token | hexadecimal | qString)

display = LBRK + simpleString + RBRK
string_ = Optional(display) + simpleString

sexp = Forward()
sexpList = Group(LPAR + ZeroOrMore(sexp) + RPAR)
sexp << ( string_ | sexpList )

test = "(add 20 5)"
input = "(mul (add 1 2) (log 8))"
input2 = "(max (data 0) (data 1))"

data = [1.0, 2.0]
n = 2


## My code ##
def doexpr(expr):
    if isinstance(expr, list):
        head = expr[0]
        tail = expr[1:]
        if (type(head) == str):
            if (head == "add"):
                return doexpr(tail[0]) + doexpr(tail[1])
            elif (head == "sub"):
                return doexpr(tail[0]) - doexpr(tail[1])
            elif (head == "mul"):
                return doexpr(tail[0]) * doexpr(tail[1])
            elif (head == "div"):
                if (doexpr(tail[1]) == 0):
                    return 0
                else:
                    return doexpr(tail[0]) / doexpr(tail[1])
            elif (head == "pow"):
                if (tail[0] < 0 and tail[1] % 1 != 0)
                    return 0
                else:
                    return doexpr(tail[0]) ** doexpr(tail[1])
            elif (head == "sqrt"):
                if doexpr(tail[0]) < 0:
                    return 0
                else:
                    return math.sqrt(doexpr(tail[0]))
            elif (head == "log"):
                if (doexpr(tail[0]) <= 0):
                    return 0
                else:
                    return math.log(doexpr(tail[0]), 2)
            elif (head == "exp"):
                return math.exp(doexpr(tail[0]))
            elif (head == "max"):
                return max(doexpr(tail[0]), doexpr(tail[1]))
            elif (head == "ifleq"):
                if (doexpr(tail[0]) <= doexpr(tail[1])):
                    return doexpr(tail[2])
                else:
                    return doexpr(tail[3])
            elif (head == "data"):
                 return data[int(abs(math.floor(doexpr(tail[0]))) % n)]
            elif (head == "diff"):
                return data[int(abs(math.floor(doexpr(tail[0]))) % n)] - data[int(abs(math.floor(doexpr(tail[1]))) % n)]
            elif (head == "avg"):
                idx0 = int(abs(math.floor(doexpr(tail[0]))) % n)
                idx1 = int(abs(math.floor(doexpr(tail[1]))) % n)
                if (idx0 != idx1):
                    sset = data[min(idx0,idx1):max(idx0,idx1)]
                    return statistics.mean(sset)
                else:
                    return 0
            else:
                sys.exit("Bad input at: %s" % expr)
        elif (type(head) == list):
            return doexpr(head)
        elif (type(head) == int or type(head) == float):
            return head
        else:
            sys.exit("malformed expression")
    elif (type(expr) == int or type(expr) == float):
        return expr
    else:
        sys.exit("malformed expression")

if (len(sys.argv) < 9):
    sys.exit("bad argument format")
for i in range(9):
    if (sys.argv[i] == "-question"):
        question = int(sys.argv[i+1])
        if (question < 1 or question > 5):
            print ("was expecting value between 1 and 5, but got: %d" % question)
            sys.exit("bad question number")
    if (sys.argv[i] == "-n"):
        n = int(sys.argv[i+1])
        if (n < 1):
            print ("was expecting value greater than 0, but got: %d" % n)
            sys.exit("bad dimension n")
    if (sys.argv[i] == "-x"):
        data = sys.argv[i+1].split(" ")
        for j in range(len(data)):
            data[j] = float(data[j])
        if (len(data) != n):
            print ("was expecting input data of length: %d, but was length: %d" % (n, len(data)))
            sys.exit("data dimension mismatch")
    if (sys.argv[i] == "-expr"):
        inputexpr = sys.argv[i+1]

if (question is None
        or n is None
        or data is None
        or inputexpr is None):
    sys.exit("Not all arguments specified")

sexpr = sexp.parseString(inputexpr, parseAll=True).asList()
print(doexpr(sexpr))