# region imports
from AlgorithmImports import *
# endregion

class OpeningRangeBreakoutUniverseAlgorithm(QCAlgorithm):

    def initialize(self):
        self.set_start_date(2016, 1, 1)
        self.set_end_date(2017, 1, 1)
        self.set_cash(10_000_000)
        
        # Set the universe parameters.
        self._indicator_period = 14 # days
        self._mean_volume_threshold = 1_000_000 # Shares
        self._universe_size = 20

        # Set the trading parameters.
        self._opening_range_minutes = 5

        # Add SPY so there is at least 1 asset at minute resolution to step the algorithm along.
        self._spy = self.add_equity('SPY').symbol 
        # Create the objects and Scheduled Events we need to select the universe and place trades.
        self._selection_data_by_symbol = {}
        self.universe_settings.schedule.on(self.date_rules.week_start(self._spy))
        self._universe = self.add_universe(self._get_fundamentals)
        self.schedule.on(self.date_rules.every_day(self._spy), self.time_rules.after_market_open(self._spy, self._opening_range_minutes), self._scan_for_entries)
        self.schedule.on(self.date_rules.every_day(self._spy), self.time_rules.before_market_close(self._spy, 1), self.liquidate)

    def _get_fundamentals(self, fundamentals):
        # Select 1000 largest Equities.
        fundamentals = sorted(fundamentals, key=lambda f: f.market_cap)[-1_000:]
        
        # Create and warm-up SelectionData objects for each Equity.
        self._selection_data_by_symbol = {f.symbol: SelectionData(self._indicator_period) for f in fundamentals}
        for bars in self.history[TradeBar](list(self._selection_data_by_symbol.keys()), self._indicator_period, Resolution.DAILY):
            for bar in bars.values():
                self._selection_data_by_symbol[bar.symbol].update(bar)

        # Select assets based on price, mean trading volume, and ATR ratio.
        return sorted(
            [
                symbol for symbol, selection_data in self._selection_data_by_symbol.items() 
                if (selection_data.sma.current.value > 5 and 
                    selection_data.is_ready and 
                    selection_data.mean_volume.current.value > self._mean_volume_threshold)
            ],
            key=lambda symbol: self._selection_data_by_symbol[symbol].atr.current.value / self._selection_data_by_symbol[symbol].sma.current.value
        )[-self._universe_size:]

    def on_securities_changed(self, changes):
        for security in changes.added_securities:
            security.max = self.max(security.symbol, self._opening_range_minutes)
            security.min = self.min(security.symbol, self._opening_range_minutes)
            security.roc = self.roc(security.symbol, self._opening_range_minutes)

    def _scan_for_entries(self):
        for symbol in self._universe.selected:
            security = self.securities[symbol]
            if security.roc.current.value > 0 and security.price > security.max.current.value:
                self.set_holdings(security.symbol, 1/self._universe_size)
            elif security.roc.current.value < 0 and security.price < security.max.current.value:
                self.set_holdings(security.symbol, -1/self._universe_size)


class SelectionData:

    def __init__(self, period):
        self.mean_volume = SimpleMovingAverage(period)
        self.atr = AverageTrueRange(period)
        self.sma = SimpleMovingAverage(period)

    def update(self, bar):
        self.mean_volume.update(bar.end_time, bar.volume)
        self.atr.update(bar)
        self.sma.update(bar.end_time, bar.volume)

    @property
    def is_ready(self):
        return self.atr.is_ready and self.mean_volume.is_ready and self.sma.is_ready