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

from System import *
from QuantConnect import *
from QuantConnect.Algorithm import *
from QuantConnect.Indicators import *
from QuantConnect.Algorithm.Framework import *
from QuantConnect.Algorithm.Framework.Risk import *
from QuantConnect.Algorithm.Framework.Alphas import *
from QuantConnect.Algorithm.Framework.Execution import *
from QuantConnect.Algorithm.Framework.Portfolio import *
from QuantConnect.Algorithm.Framework.Selection import *
from QuantConnect.Algorithm.Framework.Selection import *
from QuantConnect.Indicators import RollingWindow, SimpleMovingAverage

import numpy as np
from datetime import timedelta

from Execution.ImmediateExecutionModel import ImmediateExecutionModel
from Portfolio.EqualWeightingPortfolioConstructionModel import EqualWeightingPortfolioConstructionModel
from Risk.MaximumDrawdownPercentPerSecurity import MaximumDrawdownPercentPerSecurity
#from HealthcareUniverse import HealthcareUniverse



class SMACrossAlphaModel(AlphaModel):
    
    def __init__(self):
        # Initialize Indicators
        self.smaslow = []
        self.smafast = []
        # Create Rolling Windows
        self.checkWindow = RollingWindow[float](2)
        self.smaslowWindow = None
        self.smafastWindow = None
        self.month = None
        
    def OnSecuritiesChanged(self, algorithm, changes):
        # Modify data if securities changed
        for security in changes.AddedSecurities:
            symbol = security.Symbol
            self.smaslow.append({'symbol':symbol, 'indicator':algorithm.SMA(symbol, 21, Resolution.Daily)})
            self.smafast.append({'symbol':symbol, 'indicator':algorithm.SMA(symbol, 9, Resolution.Daily)})
        self.numSecurities = len(self.smaslow)
        # Works appropriately for manual universe, health universe fails with half empty dictionary {'symbol':symbol, }
        
    def Update(self, algorithm, data):
        insights = []
        if self.month == algorithm.Time.month: return []
        self.month = algorithm.Time.month
        if not self.smaslow[0]['indicator'].IsReady: return insights
        self.UpdateWindows(data)
        if not self.checkWindow.IsReady: return insights
        insights += self.CrossIndicator(data)
        algorithm.Debug(len(insights))
        return insights
        
    def UpdateWindows(self, data):
        if self.smaslowWindow is None: self.InitializeWindows()
        self.checkWindow.Add(1)
        self.smaslowWindow[1, :] = self.smaslowWindow[0, :]
        self.smafastWindow[1, :] = self.smafastWindow[0, :]
        for i in range(self.numSecurities):
            self.smaslowWindow[0, i] = self.smaslow[i]['indicator'].Current.Value
            self.smafastWindow[0, i] = self.smafast[i]['indicator'].Current.Value
        
    def InitializeWindows(self):
        self.smaslowWindow = np.zeros(shape=(2, self.numSecurities))
        self.smafastWindow = np.zeros(shape=(2, self.numSecurities))
        
    def CrossIndicator(self, data):
        insights = []
        for i in range(self.numSecurities):
            if (self.smafastWindow[0, i] < self.smaslowWindow[0, i]) or (self.smaslowWindow[0, i] < self.smaslowWindow[1, i]):
                insights.append(Insight.Price(self.smaslow[i]['symbol'], Expiry.EndOfMonth, InsightDirection.Down))
            else:
                insights.append(Insight.Price(self.smaslow[i]['symbol'], Expiry.EndOfMonth, InsightDirection.Up))
        return insights
        
        
class SMACrossFramework(QCAlgorithm):
    
    def Initialize(self):
        self.SetStartDate(2020, 1, 1)  # Set Start Date
        self.SetCash(1000000)  # Set Strategy Cash
        self.SetBrokerageModel(AlphaStreamsBrokerageModel())
        #self.spy = self.AddEquity('SPY')
        self.SetBenchmark(self.AddEquity('SPY', Resolution.Daily).Symbol)
        
        
        tickers = ['SPY', 'BND']
        symbols = [ Symbol.Create(ticker, SecurityType.Equity, Market.USA) for ticker in tickers ]
        self.SetUniverseSelection(ManualUniverseSelectionModel(symbols))
        self.UniverseSettings.Resolution = Resolution.Daily
        #self.SetUniverseSelection(HealthcareUniverse())
        self.SetAlpha(SMACrossAlphaModel())
        self.SetExecution(ImmediateExecutionModel())
        self.SetPortfolioConstruction(EqualWeightingPortfolioConstructionModel(lambda time: Expiry.EndOfMonth(time), portfolioBias = PortfolioBias.LongShort))
        self.SetRiskManagement(MaximumDrawdownPercentPerSecurity(0.1))