# Derek M Tishler - 2017
# https://tishlercapital.com/
# Basic TensorFlow Softmax Classification Example
# Based on https://www.tensorflow.org/get_started/mnist/pros

import random
import numpy as np
import pandas as pd

import tensorflow as tf

seed  = 1
random.seed(seed)
np.random.seed(seed)

class Model():

    def __init__(self):

        # len of hitory
        self.eval_lookback  = 252*4 + 1# input batch size will be eval_lookback+n_features-1 
        
        # We will feed in the past n open-to-open price changes
        self.n_features = 15
        
        # How much historical data do we need?
        self.warmup_count   = self.eval_lookback + self.n_features

        # define our tensorflow model/network
        self.network_setup()

    def network_setup(self):
        # Tensorflow Turorial does a great job(with illustrations) so comments left out here mostly: https://www.tensorflow.org/get_started/mnist/pros
        self.sess               = tf.InteractiveSession()

        self.x                  = tf.placeholder(tf.float32, shape=[None, self.n_features])
        self.y_                 = tf.placeholder(tf.float32, shape=[None, 2])

        self.W                  = tf.Variable(tf.zeros([self.n_features,2]))
        self.b                  = tf.Variable(tf.zeros([2]))

        self.y                  = tf.matmul(self.x,self.W) + self.b

        self.y_pred = tf.nn.softmax(self.y)

        self.cross_entropy      = tf.reduce_mean(
            tf.nn.softmax_cross_entropy_with_logits(labels=self.y_, logits=self.y))

        self.train_step         = tf.train.AdamOptimizer(1e-3).minimize(self.cross_entropy)

        # Some handy metric ops
        self.correct_prediction = tf.equal(tf.argmax(self.y,1), tf.argmax(self.y_,1))
        self.accuracy           = tf.reduce_mean(tf.cast(self.correct_prediction, tf.float32))
        #self.logloss            = tf.contrib.losses.log_loss(self.y_pred, self.y_)

        # This is done later vs Tensorflow Tutorial because of AdamOptimizer usage, which needs its own vars to be init'ed
        self.sess.run(tf.global_variables_initializer())

    def preproessing(self):
        
        # Create our input feature dataset and corresponding labels
        all_data = np.append(self.hist_data.open.values.flatten().astype(np.float32), self.current_price)
        features   = []
        labels     = []
        for i in range(self.n_features+1, len(all_data)-1):
            # input is change in priice
            features.append( np.diff(all_data[i-self.n_features-1:i].copy()) )
            # label is change in price from last day in input to the next day
            dp = 100.*(all_data[i+1]-all_data[i])/all_data[i]
            if dp > 0.0:
                dp = 1
            else:
                dp = 0
            labels.append(dp)
        self.features = np.array(features)
        self.labels   = np.array(labels)

        # convert to one hot for tensorflow
        oh = np.zeros((len(labels),2))
        oh[np.arange(len(labels)),labels] = 1.0
        self.labels = oh

    def train(self):
        # Perform  training step(s) and check train accuracy. This is really lame, use a test/train split and measure OOS data for good info about test/validation accuracy.
        for _ in range(1):
            self.train_step.run(session=self.sess, feed_dict={self.x: self.features, self.y_: self.labels})
        self.current_accuracy = self.accuracy.eval(session=self.sess, feed_dict={self.x: self.features, self.y_: self.labels})
        print("\nTrain LogLoss: %0.5f"%self.current_accuracy)
        

    def predict(self):
        # Perform inference
        pred_feat  =  np.append(self.hist_data.open.values.flatten().astype(np.float32), self.current_price)[-self.n_features-1:]
        pred_feat  = np.diff(pred_feat)
        pred_proba = self.y_pred.eval(session=self.sess, feed_dict={self.x: [pred_feat]})
        
        print("Forecast: Long p: %0.3f\tCashh p: %0.3f"%(pred_proba[0][0], pred_proba[0][1]))
        self.current_forecast = pred_proba[0]
        
        # Cash or Long
        return np.argmax(pred_proba[0])
        
        # Short or Long
        """if pred_proba[0][0] > 0.5:
            return -1.
        else:
            return 1."""
        
    


class BasicTemplateAlgorithm(QCAlgorithm):

    def Initialize(self):

        self.model = Model()

        self.SetStartDate(2005,1,1)  #Set Start Date
        self.SetEndDate(2017,11,1)    #Set End Date
        self.SetCash(100000)         #Set Strategy Cash
        
        # Find more symbols here: http://quantconnect.com/data
        self.symbol       = "SPY"
        self.model.symbol = self.symbol
        self.granularity  = Resolution.Minute

        self.SetBrokerageModel(BrokerageName.InteractiveBrokersBrokerage, AccountType.Margin)
        self.AddEquity(self.symbol, self.granularity)
        self.SetWarmUp(self.model.warmup_count)

        sPlot = Chart('Strategy Equity')
        sPlot.AddSeries(Series('Training_Model_Accuracy', SeriesType.Line, 2))
        sPlot.AddSeries(Series('Cash_Probability', SeriesType.Scatter, 2))
        sPlot.AddSeries(Series('Long_Probability', SeriesType.Scatter, 2))
        self.AddChart(sPlot)


        self.model.hist_data = self.History([self.symbol,], self.model.warmup_count, Resolution.Daily).astype(np.float32)
        self.do_once         = True

        # prevent order spam
        self.target          = 0.0
        
        self.Schedule.On(self.DateRules.EveryDay(self.symbol), \
            self.TimeRules.AfterMarketOpen(self.symbol), \
            Action(self.Rebalance))

    def OnData(self, data):
        if self.IsWarmingUp:
            return

    def Rebalance(self):
        
        self.model.current_price = float(self.Securities[self.symbol].Price)
        
        # Accrew history over time, is this faster than a big history call each time? (more for use in OnData when spamming call)
        if not self.do_once:
            new_hist      = self.History([self.symbol,], 1, Resolution.Daily).astype(np.float32)
            self.model.hist_data = self.model.hist_data.append(new_hist).iloc[1:] #append and pop stack   
        else:
            self.do_once  = False
        
        self.model.preproessing()
        self.model.train()
        signal = self.model.predict()
        self.Checkpoint()

        if signal != self.target:

            self.target = signal
            self.SetHoldings(self.symbol, self.target, liquidateExistingHoldings = True)


    def Checkpoint(self):
        self.Plot("Strategy Equity",'Cash_Probability', 100.*self.model.current_forecast[0])
        self.Plot("Strategy Equity",'Long_Probability', 100.*self.model.current_forecast[1])
        self.Plot("Strategy Equity",'Training_Model_Accuracy', 100.*self.model.current_accuracy)