from universe_selection import ScalpingUniverseSelectionModel
from alpha import ScalpingAlphaModel
from portfolio import FixedWeightingPortfolioConstructionModel

class HyperActiveGreenSnake(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2021, 2, 1)  # Set Start Date
        self.SetEndDate(2021, 2, 2)  # Set End Date
        self.SetCash(100000)  # Set Strategy Cash
        
        # Data resolution
        self.UniverseSettings.Resolution = Resolution.Minute
        
        #Universe Selection
        self.securities = []
        self.CustomUniverseSelectionModel = ScalpingUniverseSelectionModel(self)
        self.AddUniverse(self.CustomUniverseSelectionModel.SelectCoarse)
        
        #Alpha Model
        self.SetAlpha(ScalpingAlphaModel())

        #Porfolio Construction
        self.SetPortfolioConstruction(FixedWeightingPortfolioConstructionModel())
        
        #Risk Management
        self.SetRiskManagement(NullRiskManagementModel())
        
        #Execution
        self.SetExecution( ImmediateExecutionModel() ) 

    def OnData(self, data):
        pass

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

    def __init__(self,
                 fastPeriod = 5,
                 slowPeriod = 8,
                 resolution = Resolution.Minute):
        '''Initializes a new instance of the EmaCrossAlphaModel class
        Args:
            fastPeriod: The fast EMA period
            slowPeriod: The slow EMA period'''
        self.fastPeriod = fastPeriod
        self.slowPeriod = slowPeriod
        self.resolution = resolution
        self.predictionInterval = Time.Multiply(Extensions.ToTimeSpan(resolution), fastPeriod)
        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():
            if symbolData.Fast.IsReady and symbolData.Slow.IsReady:

                if symbolData.FastIsOverSlow:
                    if symbolData.Slow > symbolData.Fast:
                        insights.append(Insight.Price(symbolData.Symbol, self.predictionInterval, InsightDirection.Down))

                elif symbolData.SlowIsOverFast:
                    if symbolData.Fast > symbolData.Slow:
                        insights.append(Insight.Price(symbolData.Symbol, self.predictionInterval, InsightDirection.Up))

            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.SMA(added.Symbol, self.fastPeriod, self.resolution)
                symbolData.Slow = algorithm.SMA(added.Symbol, self.slowPeriod, self.resolution)
                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

        # 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

class FixedWeightingPortfolioConstructionModel(PortfolioConstructionModel):
    '''Provides an implementation of IPortfolioConstructionModel that gives equal weighting to all securities.
    The target percent holdings of each security is 1/N where N is the number of securities. 
    For insights of direction InsightDirection.Up, long targets are returned and
    for insights of direction InsightDirection.Down, short targets are returned.'''
    def __init__(self):
        self.insightCollection = InsightCollection()
        self.removedSymbols = []

    def CreateTargets(self, algorithm, insights):
        '''Create portfolio targets from the specified insights
        Args:
            algorithm: The algorithm instance
            insights: The insights to create portoflio targets from
        Returns:
            An enumerable of portoflio targets to be sent to the execution model'''
        self.insightCollection.AddRange(insights)

        targets = []

        if self.removedSymbols is not None:
            # zero out securities removes from the universe
            for symbol in self.removedSymbols:
                targets.append(PortfolioTarget(symbol, 0))
                self.removedSymbols = None

        if len(insights) == 0:
            return targets

        # Get symbols that have emit insights and still in the universe
        symbols = list(set([x.Symbol for x in self.insightCollection if x.CloseTimeUtc > algorithm.UtcTime]))

        # give equal weighting to each security
        percent = 0.1
        for symbol in symbols:
            activeInsights = [ x for x in self.insightCollection if x.Symbol == symbol ]
            direction = activeInsights[-1].Direction
            targets.append(PortfolioTarget.Percent(algorithm, symbol, direction * percent))

        return targets

    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'''

        # save securities removed so we can zero out our holdings
        self.removedSymbols = [x.Symbol for x in changes.RemovedSecurities]

        # remove the insights of the removed symbol from the collection
        for removedSymbol in self.removedSymbols:
            if self.insightCollection.ContainsKey(removedSymbol):
                for insight in self.insightCollection[removedSymbol]:
                    self.insightCollection.Remove(insight)

class ScalpingUniverseSelectionModel():
    
    def __init__(self, algorithm):
        
        self.algorithm = algorithm
        
    def SelectCoarse(self, coarse):
        
        #Initial filtering returns around 900 stocks/day
        universe = [c for c in coarse if 1 < c.Price < 15 and c.Volume > 1000000] 
        
        std_data = {}
        
        for equity in universe:
            symbol = equity.Symbol
            history = self.algorithm.History(symbol, 60, Resolution.Minute)
            price = equity.Price
            std = StandardDeviation(60)
            if not history.empty:
                for index, row in history.loc[str(symbol)].iterrows():
                    std.Update(index, row["close"])
            if (std.Current.Value / price) > 0.02:
                std_data[symbol] = std.Current.Value / price

        sortedByStd = sorted(std_data, key=lambda x: std_data[x], reverse=True)
        
        #self.algorithm.Debug(f'Number of Securities: {len(universe)}')
        
        return sortedByStd[:10]