#region imports
from AlgorithmImports import *
#endregion
from QuantConnect import Resolution, Extensions
from QuantConnect.Algorithm.Framework.Alphas import *
from QuantConnect.Algorithm.Framework.Portfolio import *
from itertools import groupby
from datetime import datetime, timedelta
from pytz import utc
UTCMIN = datetime.min.replace(tzinfo=utc)
#endregion
class InsightWeigtedPortfolio(PortfolioConstructionModel):


    def __init__(self, rebalancingFunc = Expiry.EndOfMonth):
        
        
        self.insightCollection = InsightCollection()
        self.removedSymbols = []
        
        """
        self.nextRebalance = None
        self.rebalancingFunc = rebalancingFunc
        """

    def CreateTargets(self, algorithm, insights):


        targets = []
            
        if len(insights) == 0:
            return targets
        
        # apply rebalancing logic
        """
        if self.nextRebalance is not None and algorithm.Time < self.nextRebalance:
            return targets
        self.nextRebalance = self.rebalancingFunc(algorithm.Time)
        """

        # here we get the new insights and add them to our insight collection
        for insight in insights:
            self.insightCollection.Add(insight)
            
        # create flatten target for each security that was removed from the universe
        if len(self.removedSymbols) > 0:
            #get the invested tickers
            invested = [x.Symbol.Value for x in algorithm.Portfolio.Values if x.Invested]
            #check if the tickers is in invested, otherwise, do nothing
            universeDeselectionTargets = [ PortfolioTarget(symbol, 0) for symbol in self.removedSymbols if symbol.Value in invested]
            targets.extend(universeDeselectionTargets)
            algorithm.Log('(Portfolio module) liquidating: ' + str([x.Value for x in self.removedSymbols]) + ' if they are active, due to not being in the universe')
            self.removedSymbols = []

        expiredInsights = self.insightCollection.RemoveExpiredInsights(algorithm.UtcTime)

        expiredTargetsLog = []
        expiredTargets = []
        for symbol, f in groupby(expiredInsights, lambda x: x.Symbol):
            if not self.insightCollection.HasActiveInsights(symbol, algorithm.UtcTime):
                expiredTargets.append(PortfolioTarget(symbol, 0))
                expiredTargetsLog.append(symbol)
                continue
        
        algorithm.Log(f'(Portfolio module) sold {expiredTargetsLog} due to insight being expired')
        targets.extend(expiredTargets)

        # get insight that have not expired of each symbol that is still in the universe
        activeInsights = self.insightCollection.GetActiveInsights(algorithm.UtcTime)
        

        # get the last generated active insight for each insight, and not symbol
        lastActiveInsights = []
        for symbol, g in groupby(activeInsights):
            lastActiveInsights.append(sorted(g, key = lambda x: x.GeneratedTimeUtc)[-1])

        calculatedTargets = {}
        for insight in lastActiveInsights:
            if insight.Symbol not in calculatedTargets.keys():
                calculatedTargets[insight.Symbol] = insight.Direction
            else:
                calculatedTargets[insight.Symbol] += insight.Direction

        
        # determine target percent for the given insights
        weightFactor = 1.0
        weightSums = sum(abs(direction) for symbol, direction in calculatedTargets.items())

        boughtTargetsLog = []

        if weightSums > 1:
            weightFactor = 1 / weightSums
            
        for symbol, weight in calculatedTargets.items():
            allocationPercent = weight * weightFactor
            target = PortfolioTarget.Percent(algorithm, symbol, allocationPercent)
            boughtTargetsLog.append(symbol)
            targets.append(target)

        algorithm.Log(f'(Portfolio module) Bought {boughtTargetsLog} stocks, that expires at {Expiry.EndOfMonth}')

        return targets
        
    def OnSecuritiesChanged(self, algorithm, changes):
        
        
        newRemovedSymbols = [x.Symbol for x in changes.RemovedSecurities if x.Symbol not in self.removedSymbols]
        
        # get removed symbol and invalidate them in the insight collection
        self.removedSymbols.extend(newRemovedSymbols)
        self.insightCollection.Clear(self.removedSymbols)
            
        removedList = [x.Value for x in self.removedSymbols]
        algorithm.Log('(Portfolio module) securities removed from Universe: ' + str(removedList))

from AlgorithmImports import *

class MarketOrderModel(ExecutionModel):

    def __init__(self):
        self.targetsCollection = PortfolioTargetCollection()

    def Execute(self, algorithm, targets):

        # for performance we check count value, OrderByMarginImpact and ClearFulfilled are expensive to call
        self.targetsCollection.AddRange(targets)
        if self.targetsCollection.Count > 0:
            for target in self.targetsCollection.OrderByMarginImpact(algorithm):
                security = algorithm.Securities[target.Symbol]
                # calculate remaining quantity to be ordered
                quantity = OrderSizing.GetUnorderedQuantity(algorithm, target, security)
                if quantity != 0:
                    aboveMinimumPortfolio = BuyingPowerModelExtensions.AboveMinimumOrderMarginPortfolioPercentage(security.BuyingPowerModel, security, quantity, algorithm.Portfolio, algorithm.Settings.MinimumOrderMarginPortfolioPercentage)
                    if aboveMinimumPortfolio:
                        algorithm.MarketOrder(security, quantity)

            self.targetsCollection.ClearFulfilled(algorithm)

from AlgorithmImports import *
import pandas as pd
import numpy as np
import statsmodels.api as sm
from enum import Enum
from collections import deque


class PairsTradingAlphaModel(AlphaModel):
    def __init__(self, lookback, resolution, prediction,  minimumCointegration, std, stoplossStd):
        self.resolution = resolution
        self.lookback = lookback
        self.prediction = prediction
        self.minimumCointegration = minimumCointegration
        self.upperStd = std
        self.lowerStd = -abs(std)
        self.upperStoploss = stoplossStd
        self.lowerStoploss = -abs(stoplossStd)
        self.mean = 0


        self.pairs = {}
        self.Securities = []


    def Update(self, algorithm, data):
        #implement the update features here. Update the RollingWindow
        insights = []

        for symbol in self.Securities:
            if not algorithm.IsMarketOpen(symbol.Symbol):
                return []
        
        #update the rolling window with same slices, or ols wont fit
        for keys, symbolData in self.pairs.items():
            if keys[0] in data.Bars and keys[1] in data.Bars:
                symbolData.symbol1Bars.Add(data[keys[0]])
                symbolData.symbol2Bars.Add(data[keys[1]])

            #Rebalacing logic here?
            if symbolData.symbol1Bars.IsReady and symbolData.symbol2Bars.IsReady:
            
                state = symbolData.state

                S1 = algorithm.PandasConverter.GetDataFrame[IBaseDataBar](symbolData.symbol1Bars).close.unstack(level=0)
                S2 = algorithm.PandasConverter.GetDataFrame[IBaseDataBar](symbolData.symbol2Bars).close.unstack(level=0)

                S1 = S1.dropna(axis=1)
                S2 = S2.dropna(axis=1)
                
                S1 = sm.add_constant(S1)
                results = sm.OLS(S2, S1).fit()
                S1 = S1[keys[0]]
                S2 = S2[keys[1]]
                b = results.params[keys[0]]
                #If S2 moves higher, the spread becomes higher. Therefore, short S2, long S1 if spread moves up, mean reversion
                spread = S2 - b * S1
                zscore = self.ZScore(spread)

                insight, state = self.TradeLogic(keys[0], keys[1], zscore[-1], state)

                #if we have changed state, append insight
                if symbolData.state != state:
                    insights.extend(insight)
                    symbolData.state = state
                else:
                    continue

        return insights

    
    def TradeLogic(self, stock1, stock2, zscore, state):
        
        insights = []

        if state == State.FlatRatio:
            if zscore > self.upperStd:
                longS1 = Insight.Price(stock1, self.prediction, InsightDirection.Up, weight=1)
                shortS2 = Insight.Price(stock2, self.prediction, InsightDirection.Down, weight=-1)
                return Insight.Group(longS1, shortS2), State.LongRatio

            elif zscore < self.lowerStd:
                shortS1 = Insight.Price(stock1, self.prediction, InsightDirection.Down, weight=-1)
                longS2 = Insight.Price(stock2, self.prediction, InsightDirection.Up, weight=1)
                return Insight.Group(shortS1, longS2), State.ShortRatio

            else:
                return [], State.FlatRatio

        if state == State.ShortRatio:
            if zscore > self.mean:
                #liquidate
                flatS1 = Insight.Price(stock1, self.prediction, InsightDirection.Flat, weight=0)
                flatS2 = Insight.Price(stock1, self.prediction, InsightDirection.Flat, weight=0)
                return Insight.Group(flatS1, flatS2), State.FlatRatio

            else:
                return [], State.ShortRatio

            """
            elif zscore > self.lowerStoploss:
                #stop loss
                #when stop loss is trickered, we dont send the state to flat, as we will wait for the spread to cross below the mean, before doing that
                flatS1 = Insight.Price(stock1, self.prediction, InsightDirection.Flat, weight=0)
                flatS2 = Insight.Price(stock1, self.prediction, InsightDirection.Flat, weight=0)
                return Insight.Group(flatS1, flatS2), AlphaSymbolData.State.ShortRatio
            """

        if state == State.LongRatio:
            if zscore < self.mean:
                #liquidate
                flatS1 = Insight.Price(stock1, self.prediction, InsightDirection.Flat, weight=0)
                flatS2 = Insight.Price(stock1, self.prediction, InsightDirection.Flat, weight=0)
                return Insight.Group(flatS1, flatS2), State.ShortRatio

            else:
                return [], State.LongRatio

            """
            elif zscore < self.lowerStoploss:
                #stop loss
                #when stop loss is trickered, we dont send the state to flat, as we will wait for the spread to cross below the mean, before doing that
                flatS1 = Insight.Price(stock1, self.prediction, InsightDirection.Flat, weight=0)
                flatS2 = Insight.Price(stock1, self.prediction, InsightDirection.Flat, weight=0)
                return Insight.Group(flatS1, flatS2), AlphaSymbolData.State.LongRatio
            """
            
        
    def ZScore(self, series):
        return (series - series.mean()) / np.std(series)


    def OnSecuritiesChanged(self, algorithm, changes):

        for security in changes.AddedSecurities:
            self.Securities.append(security)

        for security in changes.RemovedSecurities:
            if security in self.Securities:
                self.Securities.remove(security)
        
        symbols = [x.Symbol for x in self.Securities]

        history = algorithm.History(symbols, self.lookback, self.resolution).close.unstack(level=0)

        #method to calculate how cointegrated the stocks are
        n = history.shape[1]
        keys = history.columns 

        #smart looping technique. Looks at every stocks and tests it cointegration. It is kind of slow though, could be improved. 
        for i in range(n):
            for ii in range(i+1, n): 
                stock1 = history[keys[i]]
                stock2 = history[keys[ii]]

                asset1 = keys[i]
                asset2 = keys[ii]

                pair_symbol = (asset1, asset2)
                invert = (asset2, asset1)

                if pair_symbol in self.pairs or invert in self.pairs:
                    continue
                
                if stock1.hasnans or stock2.hasnans:
                    algorithm.Debug(f'WARNING! {asset1} and {asset2} has Nans. Did not perform coint')
                    continue

                #The cointegration part, that calculates cointegration between 2 stocks
                result = sm.tsa.stattools.coint(stock1, stock2) 
                pvalue = result[1] 
                if pvalue < self.minimumCointegration:
                    self.pairs[pair_symbol] = AlphaSymbolData(algorithm, asset1, asset2, 500, self.resolution)

        for security in changes.RemovedSecurities:
            keys = [k for k in self.pairs.keys() if security.Symbol in k]

            for key in keys:
                self.pairs.pop(key)

class AlphaSymbolData:
    def __init__(self, algorithm, symbol1, symbol2, lookback, resolution):

        self.state = State.FlatRatio
        self.lookback = lookback

        self.symbol1 = symbol1
        self.symbol2 = symbol2

        self.symbol1Bars = RollingWindow[IBaseDataBar](lookback)
        self.symbol2Bars = RollingWindow[IBaseDataBar](lookback)



class State(Enum):
    ShortRatio = -1
    FlatRatio = 0
    LongRatio = 1


        

#region imports
from AlgorithmImports import *
#endregion

class NoRiskManagment(RiskManagementModel):
    
    def ManageRisk(self, algorithm, targets):
        return []

from AlgorithmImports import *
from EqualPCM import InsightWeigtedPortfolio
from PairsTradingAlpha import PairsTradingAlphaModel
from ExecutionModel import MarketOrderModel
from RiskModel import NoRiskManagment
from datetime import timedelta

### <summary>
### Framework algorithm that uses the PearsonCorrelationPairsTradingAlphaModel.
### This model extendes BasePairsTradingAlphaModel and uses Pearson correlation
### to rank the pairs trading candidates and use the best candidate to trade.
### </summary>
class PairsTradingV2(QCAlgorithm):
    '''Framework algorithm that uses the PearsonCorrelationPairsTradingAlphaModel.
    This model extendes BasePairsTradingAlphaModel and uses Pearson correlation
    to rank the pairs trading candidates and use the best candidate to trade.'''

    def Initialize(self):
        self.Debug('Algorithm started. Wait for warmup')
        self.SetStartDate(2018, 12, 1)
        self.SetEndDate(2019, 2, 1)
        self.SetWarmup(100)

        self.num_coarse = 20

        self.UniverseSettings.Resolution = Resolution.Hour

        self.AddUniverse(self.CoarseUniverse)
        self.SetAlpha(PairsTradingAlphaModel(lookback = 200,
                                            resolution = Resolution.Hour,
                                            prediction = timedelta(days=20),
                                            minimumCointegration = .05,
                                            std=2,
                                            stoplossStd=2.5
                                            ))
        self.SetPortfolioConstruction(InsightWeigtedPortfolio())
        self.SetExecution(MarketOrderModel())
        self.SetRiskManagement(NoRiskManagment())

        self.rebalancingFunc = Expiry.EndOfMonth
        self.nextRebalance = None


    def CoarseUniverse(self, coarse):

        if self.nextRebalance is not None and self.Time < self.nextRebalance:
            return Universe.Unchanged
        self.nextRebalance == self.rebalancingFunc(self.Time)
        #Exclude stocks like BRKA that cost 500.000 dollars
        selected = sorted([x for x in coarse if x.HasFundamentalData and x.Price > 15], 
                        key = lambda x: x.DollarVolume, reverse = True)
        
        self.Log(f'(Universe module)Sent {len([x.Symbol for x in selected[:self.num_coarse]])} symbols to the alpha module')

        return [x.Symbol for x in selected[:self.num_coarse]]

    def OnEndOfDay(self):
        self.Plot("Positions", "Num", len([x.Symbol for x in self.Portfolio.Values if self.Portfolio[x.Symbol].Invested]))
        self.Plot(f"Margin", "Used", self.Portfolio.TotalMarginUsed)
        self.Plot(f"Margin", "Remaning", self.Portfolio.MarginRemaining)
        self.Plot(f"Cash", "Remaining", self.Portfolio.Cash)