from AlgorithmImports import *
from datetime import timedelta

class Alpha(AlphaModel):

    '''Alpha - Determine equities we should trade.
    
       This is a warping of the Alpha model purpose, but we can get the 
       AroonOscillator to work in this module.
       
       Insights are evaluated every week, but this method is called every day. 
       The indicator is likewise updated every day.
    '''

    def __init__(self, config):
        self.config = config
        self.symbolDataBySymbol = {}

    def Update(self, algorithm, data):
        ''' Update - Scan symbol data to calculate Aroon Score
            Called every day.
        '''
        # List of Insights to emit.
        # Non-null insights are only emitted on selection day.
        insights = []
        # Update Aroon Score for all equities.
        for s, sd in self.symbolDataBySymbol.items():
            if data.Bars.ContainsKey(sd.symbol):
                sd.Update(data.Bars[sd.symbol])
        # Check for selection day.
        if self.config.selectionDay == algorithm.Time.weekday():
            # Sort all the equities by Aroon Score.
            sortedByAroon = sorted(self.symbolDataBySymbol.keys(), key=lambda s: self.symbolDataBySymbol[s].aroonScore, reverse=True)
            # Trim the universe to maxUniverse, if necessary.
            universe = sortedByAroon[:self.config.maxUniverse] if len(sortedByAroon) > self.config.maxUniverse else sortedByAroon
            # Emit an insight for the qualifying entities in the universe.
            for s in universe:
                if self.symbolDataBySymbol[s].aroonScore > 0:
                    # Aroon Score must be greater than zero.
                    insights.append(Insight.Price(s, timedelta(days=1), InsightDirection.Up))
            # Reset the Aroon Score for all equities.
            for s, sd in self.symbolDataBySymbol.items():
                sd.Reset()
        return insights
        
    def OnSecuritiesChanged(self, algorithm, changes):
        for added in changes.AddedSecurities:
            symbol = self.symbolDataBySymbol.get(added.Symbol)
            if symbol is None:
                # algorithm.Debug("Adding ticker {}".format(added.Symbol.Value))
                self.symbolDataBySymbol[added.Symbol] = SymbolData(self.config, algorithm, added.Symbol)
            else:
                # algorithm.Debug("Resetting ticker {}".format(symbol.Value))
                symbol.aroon.Reset()
        for removed in changes.RemovedSecurities:
            # algorithm.Debug("Removing ticker {}".format(removed.Symbol.Value))
            self.symbolDataBySymbol.pop(removed.Symbol, None)
    
    
class SymbolData:
    ''' Indicator for one symbol of our universe.
        The indicator is updated and warmed-up manually.
    '''
    
    def __init__(self, config, algorithm, symbol):
        self.config = config
        self.symbol = symbol
        self.aroonScore = 0
        # Create Aroon Oscillator indicator.
        self.aroon = AroonOscillator(self.config.aroonPeriod, self.config.aroonPeriod)
        # Warm up Aroon Oscillator with history data reaching back from now.
        daysAgo = algorithm.Time - timedelta(days=(self.config.aroonPeriod*2))
        now = algorithm.Time
        history = algorithm.History([self.symbol], daysAgo, now, Resolution.Daily)
        for time, row in history.loc[self.symbol].iterrows():
            tradebar = self.tradebarFromRow(time, row)
            self.aroon.Update(tradebar)

    @property    
    def IsReady(self):
        ''' The symbol data is ready when the AroonOscillator indicator is 
            ready, which means we have completed warm up.
        '''
        return self.aroon.IsReady
        
    def tradebarFromRow(self, time, row):
        ''' Convert a history row to a tradebar
        '''
        tradebar = TradeBar()
        tradebar.EndTime=time 
        tradebar.Close=Decimal(row["close"])
        tradebar.High=Decimal(row["high"])
        tradebar.Low=Decimal(row["low"])
        tradebar.Open=Decimal(row["open"])
        tradebar.Volume=Decimal(row["volume"])
        return tradebar

    def Update(self, bar):
        self.aroon.Update(bar)
        if self.IsReady:
            if self.aroon.Current.Value >= self.config.aroonThreshold:
                self.aroonScore += 1
            elif self.aroon.Current.Value <= -self.config.aroonThreshold:
                self.aroonScore += 1

    def Reset(self):
        self.aroonScore = 0

#
# Portfolio Construction
#

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 *


class PortfolioConstruction(PortfolioConstructionModel):
    
    ''' PortfolioConstruction - Emit log messages with the desired universe of
        equities based on the insights received.
    '''

    def __init__(self, config):
        self.config = config
    
    def CreateTargets(self, algorithm, insights):
        ''' CreateTargets
            Each insight constitutes a tradable equity in our universe.
            Emit a log message for each equity.
        '''
        targets = []
        for insight in insights:
            algorithm.Log("Equity {}".format(insight.Symbol.Value))
        return targets

from datetime import datetime, timedelta
from QuantConnect.Data.UniverseSelection import * 
from Selection.FundamentalUniverseSelectionModel import FundamentalUniverseSelectionModel


class UniverseSelection(FundamentalUniverseSelectionModel):
    
    '''Select a new universe of equities'''
    
    def __init__(self, config, filterFineData=True, universeSettings=None):
        super().__init__(filterFineData, universeSettings)
        self.config = config
        
    def SelectCoarse(self, algorithm, coarse):
        '''Initial universe selection by looking for tickers with fundamental 
           data, and then testing for desirable qualities. The resulting 
           selection is sorted for price to enhance leverage.
        '''
        # Only run on the selected day at midnight.
        # Otherwise, the universe has not changed.
        if self.config.selectionDay != algorithm.Time.weekday():
            return Universe.Unchanged
            
        algorithm.Debug("SelectCoarse: start " + str(algorithm.Time.strftime("%Y-%m-%d %H:%M:%S")))
        # Must have fundamental data.
        # Must be an equity.
        # Must be in a US market.
        # Equity price must be above the minimum.
        # Equity market volume must be above the minimum.
        # Sort for price, because price movements for a low priced equity
        # offer greater leverage over more expensive equities.
        sortedByPrice = sorted([x for x in coarse if x.HasFundamentalData
            and x.Symbol.SecurityType == SecurityType.Equity
            and x.Symbol.ID.Market == Market.USA
            and x.AdjustedPrice >= self.config.minPrice
            and x.DollarVolume >= self.config.minVolume],
                key=lambda x: x.Price)
        universe = [x.Symbol for x in sortedByPrice]
        algorithm.Debug("SelectCoarse: universe length: " + str(len(universe)))
        #if len(universe) == 0:
        #    return Universe.Unchanged
        
        return universe

    def SelectFine(self, algorithm, fine):
        '''Refine universe selection by ensuring equities are at least a 
           configured time past their IPO.
        '''
        algorithm.Debug("SelectFine: start at {}".format(algorithm.Time.strftime("%Y-%m-%d %H:%M:%S")))
        universe = []
        for f in fine:
            # Equity IPO must have happened some time ago.
            if int((algorithm.Time - f.SecurityReference.IPODate).days) < self.config.minDaysSinceIPO:
                continue
            universe.append(f.Symbol)
        algorithm.Debug("SelectFine: universe length {}: ".format(len(universe)))
        if len(universe) == 0:
            return Universe.Unchanged
    
        # There is a new universe
        return universe

from datetime import datetime, timedelta
from QuantConnect.Data.UniverseSelection import * 
from Selection.FundamentalUniverseSelectionModel import FundamentalUniverseSelectionModel
from config import Config
from universe import UniverseSelection
from alpha import Alpha
from portfolio import PortfolioConstruction


class AroonSieve(QCAlgorithm):
    
    ''' AroonSieve
    
        Find a universe of equities which are definitely moving up or moving down
        according to the Aroon indicator.
        
        We speculate that these equities will do well in an Opening Range Breakout
        algorithm.
        
        The universe is written to a log file, suitable for further processing 
        with an aim to support upload or local backtesting.
    '''

    def Initialize(self):
        # Configuration parameters controlling this algorithm
        self.config = Config()

        # Backtesting
        self.SetStartDate(2021, 6, 1)  # Set backtest Start Date
        self.SetEndDate(2021, 8, 30)   # Set backtest End Date
        self.SetCash(10000)            # Set backtest Strategy Cash

        # Brokerage
        self.SetBrokerageModel(BrokerageName.InteractiveBrokersBrokerage, AccountType.Margin)

        # Universe Construction
        # Discover prelimary universe of tradable equities.
        self.UniverseSettings.Resolution = Resolution.Daily
        self.AddUniverseSelection(UniverseSelection(self.config))
        
        self.SetWarmup(self.config.aroonPeriod)
        
        # Alpha Model
        # Refine tradable universe by Aroon Score.
        self.SetAlpha(Alpha(self.config))
        
        # Portfolio Construction
        # Log tradable universe.
        self.SetPortfolioConstruction(PortfolioConstruction(self.config))

# Since QuantConnect doesn't support config files, let's do this kludge.

class Config():
    
    def __init__(self):
        #
        # Coarse Selection Parameters
        #
        # Saturday night at midnight
        self.selectionDay = 4
        # Minimum price
        self.minPrice = 1.00
        # Minimum dollar volume
        # 10 million dollars a day
        self.minVolume = 1e7
        #
        # Fine Selection Parameters
        #
        # Minimum days since IPO
        self.minDaysSinceIPO = 60
        # Maximum size of the universe
        self.maxUniverse = 20
        # Aroon Oscillator period
        self.aroonPeriod = 25
        # Threshold for up and down trends
        self.aroonThreshold = 90