##-------------------Class to hold variable weights---------------------------------------##        
class Weights(Object):
    MarginMultiplier = float(1.0)
    Short = float(0.805)

##-------------------Imports-------------------------------------------------------------------##

import numpy as np
import pandas as pd
import re
from datetime import *

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

from System import *
from QuantConnect import *
from QuantConnect import Resolution, Extensions
from QuantConnect.Algorithm.Framework.Execution import *
from QuantConnect.Data.UniverseSelection import *
from QuantConnect.Data.Custom.CBOE import *
from Selection.FundamentalUniverseSelectionModel import FundamentalUniverseSelectionModel
from QuantConnect.Algorithm.Framework.Selection import *

##-------------------Alpha Files and Objects--------------------------------------------------##

from DynamicWeight import Weights

#from LongAlpha import Long
#from ShortAlpha import Short


##-------------------Portfolio Construction Files---------------------------------------------##

from SourceModelPortfolioConstruction import SourceModelPortfolioConstructionModel


##-------------------Risk Management Files----------------------------------------------------##

#from RiskManagement import ManageRisk


##-------------------Global variables---------------------------------------------------------##

resolution = Resolution.Minute


class Project(QCAlgorithm):

##-------------------Initialize variables, lists, asset subscriptions, etc--------------------##

    def Initialize(self):
        # Set Start Date so that backtest has 5+ years of data
        self.SetStartDate(2014, 1, 1)

        # Set Cash
        self.SetCash(1000000)
        
        #Variables
        self.Zero = int(0)

        # Selected Securitites
        self.Assets = [
                self.AddEquity("SPY", Resolution = resolution),
                self.AddEquity("SPXL", Resolution = resolution),
                self.AddEquity("SPXS", Resolution = resolution)]

        ManualSymbols = []
        SymbolList = []
        
        for x in self.Assets:
            ManualSymbols.append(Symbol.Create(x.Symbol, SecurityType.Equity, Market.USA))
            SymbolList.append(x.Symbol)

##-------------------Construct Alpha Model----------------------------------------------------##

        #Universe Selection
        self.AddUniverseSelection(ManualUniverseSelectionModel(ManualSymbols))
        
        #Alpha Models
    #    self.AddAlpha(LongAlpha())
    #    self.AddAlpha(ShortAlpha())
        
        #Portfolio Construction
        self.SetPortfolioConstruction(SourceModelPortfolioConstructionModel(SymbolList))
        
        #Risk Management
    #    self.SetRiskManagement(ManageRisk())
        
        #Brokerage Model
        self.SetBrokerageModel(AlphaStreamsBrokerageModel())

##-----------------Imports----------------------------------------------------------------------------------------##
import numpy as np

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

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

from DynamicWeight import Weights

##-----------------Global variables-------------------------------------------------------------------------------##

Zero = int(0)
UTCMIN = datetime.min.replace(tzinfo=utc)

class SourceModelPortfolioConstructionModel(PortfolioConstructionModel):

##-----------------Initialize variables, lists, etc---------------------------------------------------------------##

    def __init__(self, SymbolList):
        '''Receives a list of the universe symbols.  Designed for manual universe selection'''
        
        # Static Variables
        self.CummulativeInsightCollection = InsightCollection()
        self.FlatInsightCollection = InsightCollection()
        self.RemovedSymbols = []
        self.SymbolList = SymbolList
        self.ErrorSymbols = {}
        self.Percents = {}
        
        # Long Only Variables
        self.LongInsightCollection = InsightCollection()
        self.LongNextExpiryTime = UTCMIN
        self.LongRebalancingTime = UTCMIN
        self.LongRebalancingFunc = lambda dt: dt + Extensions.ToTimeSpan(Resolution.Minute)
        self.LongSourceModel = None
        
        # Short Only Variables
        self.ShortInsightCollection = InsightCollection()
        self.ShortNextExpiryTime = UTCMIN
        self.ShortRebalancingTime = UTCMIN
        self.ShortRebalancingFunc = lambda dt: dt + Extensions.ToTimeSpan(Resolution.Hour)
        self.ShortSourceModel = None


##-----------------Creates target weights---------------------------------------------------------------------------##

    def DetermineTargetPercent(self, algorithm, insights):

        for insight in insights:
            self.CummulativeInsightCollection.Add(insight)
        
        ActiveInsights = self.CummulativeInsightCollection.GetActiveInsights(algorithm.UtcTime)
        
        LastActiveInsights = []
        for symbol, g in groupby(ActiveInsights, lambda x: x.Symbol):
            LastActiveInsights.append(sorted(g, key = lambda x: x.GeneratedTimeUtc)[-1])
        
        # Give equal weighting to each security
        count = sum(x.Direction != InsightDirection.Flat for x in LastActiveInsights)
        percent = Zero if count == Zero else (Weights.MarginMultiplier / count)
        
        for insight in LastActiveInsights:
            
            if insight.Direction == InsightDirection.Down:
                self.Percents[insight] = insight.Direction * np.minimum(percent, Weights.Short)
                
            else:
                 self.Percents[insight] = insight.Direction * percent


##-----------------Creates and returns Portfolio Targets------------------------------------------------------------##

    def CreateTargets(self, algorithm, insights):
            
        targets = []
            
        if self.LongSourceModel == None or self.ShortSourceModel == None:
            self.DetermineSourceModels(algorithm, insights)  

        self.DetermineTargetPercent(algorithm, insights)
        
        targets.extend(self.CreateShortPositions(algorithm, insights))
        targets.extend(self.CreateLongPositions(algorithm, insights))
        targets.extend(self.CreateFlatPositions(algorithm,  insights))
        
        # Create flatten target for each security that was removed from the universe
        if self.RemovedSymbols is not None:
            universeDeselectionTargets = [ PortfolioTarget(symbol, Zero) for symbol in self.RemovedSymbols ]
            targets.extend(universeDeselectionTargets)
            self.RemovedSymbols = None
        
        return targets  


##-----------------Captures the source models' hash codes-----------------------------------------------------------##

    def DetermineSourceModels(self, algorithm, insights):
        
        for insight in insights:
            
            if insight.Symbol == self.SymbolList[0] and insight.Direction == InsightDirection.Up and self.LongSourceModel == None and insight.SourceModel != self.ShortSourceModel:
                self.LongSourceModel = insight.SourceModel
                algorithm.Debug("Long Source Model: {0}".format(self.LongSourceModel))
                
            elif insight.Symbol == self.SymbolList[0] and insight.Direction == InsightDirection.Down and insight.Period == timedelta(days=3) and self.ShortSourceModel == None:
                self.ShortSourceModel = insight.SourceModel
                algorithm.Debug("Short Source Model: {0}".format(self.ShortSourceModel))
                
                
                
##-----------------Generates Long  targets separately------------------------------------------------------------##

    def CreateLongPositions(self, algorithm, insights):
        
        LongTargets = []
        
        if (algorithm.UtcTime <= self.LongNextExpiryTime and algorithm.UtcTime <= self.LongRebalancingTime and len(insights) == Zero and self.RemovedSymbols is None):
            return LongTargets
            
        for insight in insights:
            if insight.SourceModel == self.LongSourceModel:
                self.LongInsightCollection.Add(insight) 
        
        # Get insight that haven't expired of each symbol that is still in the universe
        ActiveInsights = self.LongInsightCollection.GetActiveInsights(algorithm.UtcTime)

        # Get the last generated active insight
        LastActiveInsights = []
        for symbol, g in groupby(ActiveInsights, lambda x: x.Symbol):
            LastActiveInsights.append(sorted(g, key = lambda x: x.GeneratedTimeUtc)[-1])
        
        self.ErrorSymbols = {}
        for insight in LastActiveInsights:
            symbol = insight.Symbol
            target = PortfolioTarget.Percent(algorithm, symbol, self.Percents[insight])
            if not target is None:
                LongTargets.append(target)
            else:
                self.ErrorSymbols[symbol] = symbol

        # Get expired insights and create flatten targets for each symbol
        ExpiredInsights = self.LongInsightCollection.RemoveExpiredInsights(algorithm.UtcTime)
        
        for insight in ExpiredInsights:
            self.Percents.pop(insight)

        ExpiredTargets = []
        for symbol, f in groupby(ExpiredInsights, lambda x: x.Symbol):
            if not self.LongInsightCollection.HasActiveInsights(symbol, algorithm.UtcTime) and not symbol in self.ErrorSymbols:
                ExpiredTargets.append(PortfolioTarget(symbol, Zero))
                continue

        LongTargets.extend(ExpiredTargets)

        self.LongNextExpiryTime = self.LongInsightCollection.GetNextExpiryTime()
        
        if self.LongNextExpiryTime is None:
            self.LongNextExpiryTime = UTCMIN

        self.LongRebalancingTime = self.LongRebalancingFunc(algorithm.UtcTime)  
            
        return LongTargets
        
        
        
##-----------------Generates Short  targets separately-----------------------------------------------------------##

    def CreateShortPositions(self, algorithm, insights):
        
        ShortTargets = []
        
        if (algorithm.UtcTime <= self.ShortNextExpiryTime and algorithm.UtcTime <= self.ShortRebalancingTime and len(insights) == Zero and self.RemovedSymbols is None):
            return ShortTargets
        
        for insight in insights:
            if insight.SourceModel == self.ShortSourceModel:
                self.ShortInsightCollection.Add(insight) 
        
        # Get insight that haven't expired of each symbol that is still in the universe
        ActiveInsights = self.ShortInsightCollection.GetActiveInsights(algorithm.UtcTime)

        # Get the last generated active insight
        LastActiveInsights = []
        for symbol, g in groupby(ActiveInsights, lambda x: x.Symbol):
            LastActiveInsights.append(sorted(g, key = lambda x: x.GeneratedTimeUtc)[-1])

        self.ErrorSymbols = {}
        for insight in LastActiveInsights:
            symbol = insight.Symbol
            target = PortfolioTarget.Percent(algorithm, symbol, self.Percents[insight])
            if not target is None:
                ShortTargets.append(target)
            else:
                self.ErrorSymbols[symbol] = symbol

        # Get expired insights and create flatten targets for each symbol
        ExpiredInsights = self.ShortInsightCollection.RemoveExpiredInsights(algorithm.UtcTime)

        for insight in ExpiredInsights:
            self.Percents.pop(insight)
            
        ExpiredTargets = []
        for symbol, f in groupby(ExpiredInsights, lambda x: x.Symbol):
            if not self.ShortInsightCollection.HasActiveInsights(symbol, algorithm.UtcTime) and not symbol in self.ErrorSymbols:
                ExpiredTargets.append(PortfolioTarget(symbol, Zero))
                continue

        ShortTargets.extend(ExpiredTargets)

        self.ShortNextExpiryTime = self.ShortInsightCollection.GetNextExpiryTime()
        
        if self.ShortNextExpiryTime is None:
            self.ShortNextExpiryTime = UTCMIN

        self.ShortRebalancingTime = self.ShortRebalancingFunc(algorithm.UtcTime)
        
        return ShortTargets
        
        

##-----------------Generates Flat targets separately----------------------------------------------------------------##    

    def CreateFlatPositions(self, algorithm, insights):
        
        FlatTargets = []
        
        self.ErrorSymbols = {}
        for insight in insights:
            if insight.Direction == InsightDirection.Flat:
                symbol = insight.Symbol
                target = PortfolioTarget.Percent(algorithm, symbol, Zero)
                if not target is None:
                    FlatTargets.append(target)
                else:
                    self.ErrorSymbols[symbol] = symbol    
                
        return FlatTargets   



##-----------------Handles changes to the universe------------------------------------------------------------------##

    def OnSecuritiesChanged(self, algorithm, changes):
    
        # Get removed symbol and invalidate them in the insight collections
        self.RemovedSymbols = [x.Symbol for x in changes.RemovedSecurities]
        
        self.LongInsightCollection.Clear(self.RemovedSymbols)
        self.ShortInsightCollection.Clear(self.RemovedSymbols)
        self.FlatInsightCollection.Clear(self.RemovedSymbols)