from deap import creator, base, tools, algorithms, gp
import operator
import re
from random import randint
from zipline.api import order_target, record, symbol
from zipline import run_algorithm
import matplotlib.pyplot as plt
import pandas as pd

def protectedDiv(left, right):
    try:
        return left / right
    except ZeroDivisionError:
        return 1


pset = gp.PrimitiveSetTyped("MAIN", [], bool)

# arithmetic operators
pset.addPrimitive(operator.add, [float, float], float)
pset.addPrimitive(operator.sub, [float, float], float)
pset.addPrimitive(protectedDiv, [float, float], float)
pset.addPrimitive(operator.mul, [float, float], float)


def and_(a, b):
    return (a and b)


def or_(a, b):
    return (a or b)


def not_(a):
    return not a


# boolean operators
pset.addPrimitive(and_, [bool, bool], bool)
pset.addPrimitive(or_, [bool, bool], bool)
pset.addPrimitive(not_, [bool], bool)

# relational operators
pset.addPrimitive(operator.lt, [float, float], bool)
pset.addPrimitive(operator.gt, [float, float], bool)


def if_then_else(input, output1, output2):
    if input:
        return output1
    else:
        return output2


pset.addPrimitive(if_then_else, [bool, bool, bool], bool)


def average(string, days):
    if days > 300:
        days = 300
    average = data.history(context.asset, string, bar_count=int(
        days), frequency="1d").mean()

    return average


def maximum(type, days):
    return 100


def minimum(type, days):
    return 100


def norm(left, right):
    return abs(left - right)


def idem(string):
    return string


pset.addPrimitive(idem, [str], str)

pset.addPrimitive(norm, [float, float], float)
pset.addPrimitive(average, [str, float], float)
pset.addPrimitive(maximum, [str, float], float)
pset.addPrimitive(average, [str, float], float)
pset.addTerminal("price", str)
pset.addTerminal("volume", str)
pset.addTerminal(False, bool)
pset.addTerminal(True, bool)
pset.addEphemeralConstant('a', lambda: float(randint(0, 300)), float)
expr = gp.genFull(pset, min_=1, max_=4)

tree = gp.PrimitiveTree(expr)
treeString = str(tree)
treeString = re.sub(r"idem\([^']*", '', treeString)
treeString = re.sub(r"(?<='price')[^,]*", '', treeString)
treeString = re.sub(r"(?<='vol')[^,]*", '', treeString)
tree2 = gp.PrimitiveTree.from_string(treeString, pset)
print(tree2)

function = gp.compile(tree2, pset)


def initialize(context):
    context.i = 0
    context.asset = symbol('AAPL')


def handle_data(context, data):
    # Skip first 300 days to get full windows
    context.i += 1
    if context.i < 300:
        return

    print(average('price', 100))

    # Trading lo
    # gic
    if function:
        # order_target orders as many shares as needed to
        # achieve the desired number of shares.
        order_target(context.asset, 100)


def analyze(context, perf):
    fig = plt.figure()
    ax1 = fig.add_subplot(211)
    perf.portfolio_value.plot(ax=ax1)
    ax1.set_ylabel('portfolio value in $')

    plt.legend(loc=0)
    plt.show()


start = pd.Timestamp('2014-01-01', tz='utc')
end = pd.Timestamp('2018-01-01', tz='utc')

run_algorithm(start=start, end=end, initialize=initialize,
                  capital_base=100000, handle_data=handle_data,
                  bundle='quantopian-quandl', analyze=analyze)