import numpy as np

### <summary>
### Basic template algorithm simply initializes the date range and cash. This is a skeleton
### framework you can use for designing an algorithm.
### </summary>
class BasicTemplateAlgorithm(QCAlgorithm):
    '''Basic template algorithm simply initializes the date range and cash'''
    _tkr1 = 'SPY'
    _tkr2 = 'QQQ'
    _spread = -1


    def calc_sma(self,tkr1_hist, tkr2_hist, mult1, mult2, num_periods):
        if len(tkr1_hist)<num_periods:
            return(None)
        if len(tkr2_hist)<num_periods:
            return(None)
        series = tkr1_hist*mult1+tkr2_hist*mult2
        return(series.mean())

    def Initialize(self):
        '''Initialise the data and resolution required, as well as the cash and start-end dates for your algorithm. All algorithms must initialized.'''

        self.SetStartDate(2013,10,07)  #Set Start Date
        self.SetEndDate(2013,10,11)    #Set End Date
        self.SetCash(100000)           #Set Strategy Cash
        # Find more symbols here: http://quantconnect.com/data
        # Add securities and get the data
        self.AddEquity(self._tkr1, Resolution.Minute)
        self.AddEquity(self._tkr2, Resolution.Minute)
        self.Debug("numpy test >>> print numpy.pi: " + str(np.pi))

    def OnData(self, data):
        '''OnData event is the primary entry point for your algorithm. Each new data point will be pumped in here.

        Arguments:
            data: Slice object keyed by symbol containing the stock data
        '''
        num_periods = 15
        mult1 = 1
        mult2 = -1
        history_leg1 = self.History([self._tkr1], num_periods, Resolution.Minute)
        history_leg2 = self.History([self._tkr2], num_periods, Resolution.Minute)
        ma = self.calc_sma(history_leg1, history_leg2, mult1, mult2, num_periods)
        if ma is None:
            return
        leg1 = data[self._tkr1].Close
        leg2 = data[self._tkr2].Close
        self._spread = leg1*mult1 + leg2*mult2
        self.Plot("Spread", self._spread)
        if not self.Portfolio.Invested:
            self.SetHoldings(self._tkr1,1)
            self.SetHoldings(self._tkr2,-1)
        else:
            # need to close position
            pass
    
    def OnEndOfDay(self):
        self.Debug("End of day event.")
        self.Debug(str(self._spread))