# QUANTCONNECT.COM - Democratizing Finance, Empowering Individuals.
# Lean Algorithmic Trading Engine v2.0. Copyright 2014 QuantConnect Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from clr import AddReference
AddReference("System")
AddReference("QuantConnect.Algorithm")
AddReference("QuantConnect.Algorithm.Framework")
AddReference("QuantConnect.Common")

from System import *
from QuantConnect import *
from QuantConnect.Orders import *
from QuantConnect.Algorithm import *
from QuantConnect.Algorithm.Framework import *

from Selection.FundamentalUniverseSelectionModel import FundamentalUniverseSelectionModel
#from Selection.UncorrelatedToSPYUniverseSelectionModel import UncorrelatedToSPYUniverseSelectionModel

from datetime import datetime, timedelta

import pandas as pd
import numpy as np

class UncorrelatedToSPYFrameworkAlgorithm(QCAlgorithmFramework):

    def Initialize(self):
        
        self.UniverseSettings.Resolution = Resolution.Daily
        
        self.SetStartDate(2019,2,2)   # Set Start Date
        self.SetEndDate(2019,3,15)    # Set End Date
        self.SetCash(100000)          # Set Strategy Cash

        self.SetUniverseSelection(UncorrelatedUniverseSelectionModel())
        self.SetAlpha(UncorrelatedToSPYAlphaModel())
        self.SetPortfolioConstruction(EqualWeightingPortfolioConstructionModel())
        self.SetExecution(ImmediateExecutionModel())


    def OnOrderEvent(self, orderEvent):
        if orderEvent.Status == OrderStatus.Filled:
            self.Debug("Purchased Stock: {0}".format(orderEvent.Symbol))
            

class UncorrelatedToSPYAlphaModel(AlphaModel):
    '''Uses ranking of intraday percentage difference between open price and close price to create magnitude and direction prediction for insights'''

    def __init__(self, *args, **kwargs): 
        self.lookback = kwargs['lookback'] if 'lookback' in kwargs else 1
        self.numberOfStocks = kwargs['numberOfStocks'] if 'numberOfStocks' in kwargs else 10
        self.resolution = kwargs['resolution'] if 'resolution' in kwargs else Resolution.Daily
        self.predictionInterval = Time.Multiply(Extensions.ToTimeSpan(self.resolution), self.lookback)
        self.symbolDataBySymbol = {}

    def Update(self, algorithm, data):
        
        insights = []
        ret = []
        symbols = []
        
        activeSec = [x.Key for x in algorithm.ActiveSecurities]
        
        for symbol in activeSec:
            if algorithm.ActiveSecurities[symbol].HasData:
                open = algorithm.Securities[symbol].Open
                close = algorithm.Securities[symbol].Close
                if open != 0:
                    openCloseReturn = close/open - 1
                    ret.append(openCloseReturn)
                    symbols.append(symbol)
                    
                    
        # Intraday price change
        symbolsRet = dict(zip(symbols,ret))
        
        # Rank on price change
        symbolsRanked = dict(sorted(symbolsRet.items(), key=lambda kv: kv[1],reverse=False)[:self.numberOfStocks])
        
        # Emit "up" insight if the price change is positive and "down" insight if the price change is negative
        for key,value in symbolsRanked.items():
            if value > 0:
                insights.append(Insight.Price(key, self.predictionInterval, InsightDirection.Up, value, None))
            else:
                insights.append(Insight.Price(key, self.predictionInterval, InsightDirection.Down, value, None))

        return insights

        
class UncorrelatedUniverseSelectionModel(FundamentalUniverseSelectionModel):
    '''This universe selection model picks stocks that currently have their correlation to a benchmark deviated from the mean.'''

    def __init__(self,
                 benchmark = Symbol.Create("SPY", SecurityType.Equity, Market.USA),
                 numberOfSymbolsCoarse = 100,
                 numberOfSymbols = 10,
                 windowLength = 5,
                 historyLength = 25):
        '''Initializes a new default instance of the OnTheMoveUniverseSelectionModel
        Args:
            benchmark: Symbol of the benchmark
            numberOfSymbolsCoarse: Number of coarse symbols
            numberOfSymbols: Number of symbols selected by the universe model
            windowLength: Rolling window length period for correlation calculation
            historyLength: History length period'''
        super().__init__(False)

        self.benchmark = benchmark 
        self.numberOfSymbolsCoarse = numberOfSymbolsCoarse
        self.numberOfSymbols = numberOfSymbols
        self.windowLength = windowLength
        self.historyLength = historyLength
        
        # Symbols in universe
        self.symbols = []
        self.coarseSymbols = []
        self.cor = None
        
        self.removedSymbols = []
        self.symbolDataBySymbol = {}
        
    
        self.lastSymbols = []

    def SelectCoarse(self, algorithm, coarse):

        #if not self.coarseSymbols:
        # The stocks must have fundamental data
        # The stock must have positive previous-day close price
        # The stock must have positive volume on the previous trading day
        filtered = [x for x in coarse if x.HasFundamentalData and x.Volume > 0 and x.Price > 0]
        sortedByDollarVolume = sorted(filtered, key = lambda x: x.DollarVolume, reverse=True)[:self.numberOfSymbolsCoarse]

        self.coarseSymbols = [x.Symbol for x in sortedByDollarVolume]
        
        self.corRanked(algorithm, self.coarseSymbols)

        return self.coarseSymbols

    def corRanked(self, algorithm, symbols):

        # Not enough symbols to filter
        if len(symbols) <= self.numberOfSymbols:
            return symbols
        
        for symbol in symbols:
            if not symbol in self.symbolDataBySymbol.keys():
                symbolData = self.SymbolData(symbol=symbol,algorithm=algorithm, benchmark = self.benchmark,windowLength = self.windowLength, historyLength = self.historyLength)
                symbolData.WarmUp()
                self.symbolDataBySymbol[symbol] = symbolData
            elif symbol in self.lastSymbols:
                self.symbolDataBySymbol[symbol].update()
            else:
                self.symbolDataBySymbol.pop(symbol, None)            
        
        self.lastSymbols = self.symbols
        
        zScore = {}
        for symbol in self.symbolDataBySymbol.keys():
            obj = self.symbolDataBySymbol[symbol]

            if abs(obj.getMu()) > 0.5:
                zScore.update({symbol : obj.getScore()})
        
        return sorted(zScore, key=lambda symbol: zScore[symbol],reverse=True)[:self.numberOfSymbols]
            
    class SymbolData:
        '''Contains data specific to a symbol required by this model'''
        def __init__(self,algorithm, symbol, benchmark,windowLength, historyLength):
            self.algorithm = algorithm
            self.symbol = symbol
            self.benchmark = benchmark
            self.windowLength = windowLength
            self.historyLength = historyLength
            
            self.zScore = None
            self.returns = None
            self.cor = None
            self.corMu = None

        def calcScore(self):
            
            # Calculate the mean of correlation
            self.corMu = self.cor.mean()[0]

            # Calculate the standard deviation of correlation
            corStd = self.cor.std()[0]
    
            # Current correlation
            corCur = self.cor.tail(1).unstack()[0]
    
            # Calculate absolute value of Z-Score for stocks in the Coarse Universe. 
            self.zScore = (abs(corCur - self.corMu) / corStd)
            
        def WarmUp(self):
            returns=self.getReturns()
            # Calculate stdev(correlation) using rolling window for all history
            corMat=returns.rolling(self.windowLength,min_periods = self.windowLength).corr().dropna()
            
            # Correlation of all securities against SPY
            self.cor = corMat[str(self.benchmark)].unstack()
            
            self.calcScore()
        
        def update(self):
            self.getReturns()
            corRow=returns.tail(self.windowLength).corr()[str(self.benchmark)]
            self.cor = self.cor.append(corRow).tail(self.historyLength)
            self.calcScore()

            
        def getReturns(self):
            hist = self.algorithm.History([self.symbol,self.benchmark], self.historyLength, Resolution.Daily)
        
            # Calculate returns
            self.returns=hist.close.unstack(level=0).pct_change()
    
            return self.returns
        
        ## Accessors    
        def getScore(self):
            return self.zScore
            
        def getMu(self):
            return self.corMu