## from TestLibrary import BasicTemplateLibrary
from numpy import sort
from numpy import mean
from numpy import std
from numpy import array
from numpy import argpartition
from numpy import zeros
from scipy.stats import skew
from collections import deque

## Not using skew as defined in Lemperiere 2016 paper
## because the calculation was questionable

## Description:
## - Calculate momentum, momentum change, skew

class SkewIndicator:
    def __init__(self, name, LookBackLen):
        self.IsReady = False
        self.Name = name
        self.LookBackLen = LookBackLen
        self.Value = 0.0 # dummy for initiating
        self.price_queue = deque(maxlen = LookBackLen + 1)
        self.return_queue = deque(maxlen = LookBackLen)
    def Update(self, input):
        # Fill the queues
        self.price_queue.appendleft(input.Close)
        if len(self.price_queue) > 1:
            self.return_queue.appendleft(self.price_queue[1] / self.price_queue[0] - 1.0)
        # Check if ready, and calculate skew
        if len(self.return_queue) == self.return_queue.maxlen:
            self.IsReady = True
            self.Value = self.calculateSkew()
            
    def calculateSkew(self):
        #  Return vector
        returnVec = list(self.return_queue)
        return skew(returnVec)
        
class LongPositiveSkew(QCAlgorithm):
    LookBackLen = 150
    tradeSize = 1.0
    skewThreshold = 0.1
    selectTopN = 3
    
    def Initialize(self):
        self.SetStartDate(2010, 12, 12)  # Set Start Date
        self.SetCash(13000)  # Set Strategy Cash
        
        # Universe
        self.tickerList = ["XLP","XLU","XLK","XLV","XLY","XLI","XLF","XLB","XLE"]
        self.numTickers = len(self.tickerList)
        for i in range(self.numTickers):
            self.AddEquity(self.tickerList[i], Resolution.Daily)
        self.UniverseSettings.Resolution = Resolution.Daily
        symbols = [ Symbol.Create(self.tickerList[0], SecurityType.Equity, Market.USA), \
                    Symbol.Create(self.tickerList[1], SecurityType.Equity, Market.USA) ]
        self.SetUniverseSelection(ManualUniverseSelectionModel(symbols))
        
        # Indicator
        self.SkewIndicatorList = []
        for i in range(self.numTickers):
            self.SkewIndicatorList.append(SkewIndicator('mySkew_' + self.tickerList[i], self.LookBackLen))
            self.RegisterIndicator(self.tickerList[i], self.SkewIndicatorList[i], Resolution.Daily)
        
        # Chart
        self.chartName = "Skew Chart"
        SkewChart = Chart(self.chartName, ChartType.Stacked)
        for i in range(self.numTickers):
            SkewChart.AddSeries(Series('Skew_' + self.tickerList[i], SeriesType.Line))
        self.AddChart(SkewChart);
        
        self.SetBrokerageModel(BrokerageName.InteractiveBrokersBrokerage, AccountType.Margin)

    def OnData(self, data):
        LatestSkews = zeros(self.numTickers)
        for i in range(self.numTickers):
            self.Plot(self.chartName, 'Skew_' + self.tickerList[i], self.SkewIndicatorList[i].Value)
            LatestSkews[i] = self.SkewIndicatorList[i].Value
            
        topSkewIndices = argpartition(LatestSkews, self.numTickers - self.selectTopN)[(self.numTickers - self.selectTopN):]
        
        for i in range(self.selectTopN):
            thisIndex = topSkewIndices[i]
            if self.SkewIndicatorList[thisIndex].Value < self.skewThreshold:
                self.SetHoldings(self.tickerList[thisIndex], 0.0)
            elif self.SkewIndicatorList[thisIndex].Value > self.skewThreshold:
                self.SetHoldings(self.tickerList[thisIndex], self.tradeSize / self.selectTopN)