from TrendRevAlphaModel import EmaCrossAlphaModel
from Selection.QC500UniverseSelectionModel import QC500UniverseSelectionModel


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




class CalculatingTanAnguilline(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2020, 10, 1)  # Set Start Date
        self.SetEndDate(2020,10,19)

        self.SetCash(100000)  # Set Strategy Cash
        # self.AddEquity("SPY", Resolution.Minute)
        
        self.window_size = 15
        self.SetWarmUp(self.window_size)
        
        self.AddAlpha(EmaCrossAlphaModel(fastPeriod = 5,
                 slowPeriod = 10,
                 resolution = Resolution.Daily, 
                 window_size = self.window_size))

        #self.SetUniverseSelection(QC500UniverseSelectionModel())
        
        symbols = [ Symbol.Create("TSLA", SecurityType.Equity, Market.USA) , Symbol.Create("HES", SecurityType.Equity, Market.USA) ]
        self.SetUniverseSelection( ManualUniverseSelectionModel(symbols) )

        
        self.SetExecution(ImmediateExecutionModel())


        self.SetPortfolioConstruction(EqualWeightingPortfolioConstructionModel())
        
        

        self.SetRiskManagement(MaximumDrawdownPercentPerSecurity(0.01))



    def OnData(self, data):
        '''OnData event is the primary entry point for your algorithm. Each new data point will be pumped in here.
            Arguments:
                data: Slice object keyed by symbol containing the stock data
        '''

        # if not self.Portfolio.Invested:
        #    self.SetHoldings("SPY", 1)

import pandas as pd
import numpy as np

class EmaCrossAlphaModel(AlphaModel):
    '''Alpha model that uses an EMA cross to create insights'''

    def __init__(self,
                 fastPeriod = 5,
                 slowPeriod = 10,
                 resolution = Resolution.Daily, 
                 window_size = 15):
        '''Initializes a new instance of the EmaCrossAlphaModel class
        Args:
            fastPeriod: The fast EMA period
            slowPeriod: The slow EMA period'''
        self.window_size = window_size
        self.fastPeriod = fastPeriod
        self.slowPeriod = slowPeriod
        self.resolution = resolution
        self.predictionInterval = Time.Multiply(Extensions.ToTimeSpan(resolution), fastPeriod)
        
        self.rolling_window = pd.DataFrame()

        
        self.symbolDataBySymbol = {}
        

        resolutionString = Extensions.GetEnumString(resolution, Resolution)
        self.Name = '{}({},{},{})'.format(self.__class__.__name__, fastPeriod, slowPeriod, resolutionString)


    def Update(self, algorithm, data):
        '''Updates this alpha model with the latest data from the algorithm.
        This is called each time the algorithm receives data for subscribed securities
        Args:
            algorithm: The algorithm instance
            data: The new data available
        Returns:
            The new insights generated'''
        insights = []
        for symbol, symbolData in self.symbolDataBySymbol.items():
            
            close = data[symbol].Close
            high = data[symbol].High
            low = data[symbol].Low
            open = data[symbol].Open
            volume = data[symbol].Volume
        
            row = pd.DataFrame({"close": [close], "low": [low], "high": [high], "open": [open], "volume" : [volume]}, index=[data.Time])
            symbolData.rolling_window = symbolData.rolling_window.append(row).iloc[-self.window_size:]

            algorithm.Log(f"\nRolling Window:\n{symbolData.rolling_window.to_string()}\n")
            if symbolData.Fast.IsReady and symbolData.Slow.IsReady:
                #algorithm.Debug(str(algorithm.Time))

                if symbolData.FastIsOverSlow:
                    if symbolData.Slow > symbolData.Fast:
                        insights.append(Insight.Price(symbolData.Symbol, self.predictionInterval, InsightDirection.Down))
                        algorithm.Debug(str(algorithm.Time) + " " +str(symbol) + " Buy " + str(symbolData.rolling_window))

                elif symbolData.SlowIsOverFast:
                    if symbolData.Fast > symbolData.Slow:
                        insights.append(Insight.Price(symbolData.Symbol, self.predictionInterval, InsightDirection.Up))
                        algorithm.Debug(str(algorithm.Time) + " " +str(symbol) + " Sell")

            symbolData.FastIsOverSlow = symbolData.Fast > symbolData.Slow

        return insights

    def OnSecuritiesChanged(self, algorithm, changes):
        '''Event fired each time the we add/remove securities from the data feed
        Args:
            algorithm: The algorithm instance that experienced the change in securities
            changes: The security additions and removals from the algorithm'''
        for added in changes.AddedSecurities:
            symbolData = self.symbolDataBySymbol.get(added.Symbol)
            if symbolData is None:
                # create fast/slow EMAs
                symbolData = SymbolData(added)
                symbolData.Fast = algorithm.EMA(added.Symbol, self.fastPeriod, self.resolution)
                symbolData.Slow = algorithm.EMA(added.Symbol, self.slowPeriod, self.resolution)
                symbolData.rolling_window = algorithm.History(added.Symbol, self.window_size)
                algorithm.Log(f"\nRolling Window:\n{symbolData.rolling_window.to_string()}\n")
                self.symbolDataBySymbol[added.Symbol] = symbolData
            else:
                # a security that was already initialized was re-added, reset the indicators
                symbolData.Fast.Reset()
                symbolData.Slow.Reset()


class SymbolData:
    '''Contains data specific to a symbol required by this model'''
    def __init__(self, security):
        self.Security = security
        self.Symbol = security.Symbol
        self.Fast = None
        self.Slow = None
        self.rolling_window = pd.DataFrame()

        # True if the fast is above the slow, otherwise false.
        # This is used to prevent emitting the same signal repeatedly
        self.FastIsOverSlow = False

    @property
    def SlowIsOverFast(self):
        return not self.FastIsOverSlow