from AlgorithmImports import *


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)

        # Parameters
        self.max_positions = 20
        self.risk = 0.01            # equity risk per position
        self.entry_gap = 0.1        # fraction of ATR from close price for entry order stop level
        self._universe_size = self.get_parameter("universeSize", 1000)
        self._atr_threshold = 0.5
        self._indicator_period = 14  # days
        self._opening_range_minutes = self.get_parameter("openingRangeMinutes", 1)
        self._leverage = 4
        self._symbol_data_by_symbol = {}

        # Add SPY as a benchmark and stepping asset
        self._spy = self.add_equity("SPY").symbol

        # Add a universe of the most liquid US equities
        self.universe_settings.leverage = self._leverage
        # self.universe_settings.asynchronous = True  # won't work correct since self.time will be wrong
        self.last_month = None
        self._universe = self.add_universe(self.filter_universe)

        # Schedule daily liquidation before market close
        self.schedule.on(
            self.date_rules.every_day(self._spy),
            self.time_rules.before_market_close(self._spy, 1),
            self.liquidate
        )

        # Warm-up indicators
        self.set_warm_up(timedelta(days=2 * self._indicator_period))

    def filter_universe(self, fundamentals):
        # only update universe on first day of month
        if self.time.month == self.last_month:
            return Universe.UNCHANGED

        self.last_month = self.time.month
        return [
            f.symbol for f in sorted(
                [f for f in fundamentals if f.price > 5 and f.symbol != self._spy],
                key=lambda f: f.dollar_volume,
                reverse=True
            )[:self._universe_size]
        ]

    def on_securities_changed(self, changes):
        # Add indicators for each asset that enters the universe
        for security in changes.added_securities:
            self._symbol_data_by_symbol[security.symbol] = SymbolData(self, security, self._opening_range_minutes,
                                                                      self._indicator_period)

    def on_data(self, slice):
        if self.is_warming_up or not (self.time.hour == 9 and self.time.minute == 30 + self._opening_range_minutes):
            return

        # Select stocks in play
        filtered = sorted(
            [self._symbol_data_by_symbol[s] for s in self.active_securities.keys
             if self.active_securities[s].price > 0 and s in self._universe.selected  # not sure about last condition, if it is needed?
             and self._symbol_data_by_symbol[s].relative_volume > 1
             and self._symbol_data_by_symbol[s].ATR.current.value > self._atr_threshold],
            key=lambda x: x.relative_volume,
            reverse=True
        )[:self.max_positions]

        # Look for trade entries
        for symbolData in filtered:
            symbolData.scan()

    def on_order_event(self, orderEvent):
        if orderEvent.status != OrderStatus.FILLED:
            return

        if orderEvent.symbol in self._symbol_data_by_symbol:
            self._symbol_data_by_symbol[orderEvent.symbol].on_order_event(orderEvent)


class SymbolData:
    def __init__(self, algorithm: QCAlgorithm, security, openingRangeMinutes, indicatorPeriod):
        self.algorithm = algorithm
        self.security = security
        self.opening_bar = None
        self.relative_volume = 0
        self.ATR = algorithm.ATR(security.symbol, indicatorPeriod, resolution=Resolution.DAILY)
        self.volumeSMA = SimpleMovingAverage(indicatorPeriod)
        self.stop_loss_price = None
        self.entry_ticket = None
        self.stop_loss_ticket = None

        self.consolidator = algorithm.consolidate(
            security.symbol, TimeSpan.from_minutes(openingRangeMinutes), self.consolidation_handler
        )

    def consolidation_handler(self, bar):
        if self.opening_bar and self.opening_bar.time.date() == bar.time.date():
            return

        self.relative_volume = bar.volume / self.volumeSMA.current.value if self.volumeSMA.is_ready and self.volumeSMA.current.value > 0 else 0
        self.volumeSMA.update(bar.end_time, bar.volume)
        self.opening_bar = bar

    def scan(self):
        if not self.opening_bar:
            return

        if self.opening_bar.close > self.opening_bar.open:
            self.place_trade(self.opening_bar.high, self.opening_bar.high - self.algorithm.entry_gap * self.ATR.current.value)
        elif self.opening_bar.close < self.opening_bar.open:
            self.place_trade(self.opening_bar.low, self.opening_bar.low + self.algorithm.entry_gap * self.ATR.current.value)

    def place_trade(self, entryPrice, stopPrice):
        risk_per_position = (self.algorithm.portfolio.total_portfolio_value * self.algorithm.risk) / self.algorithm.max_positions
        quantity = int(risk_per_position / (entryPrice - stopPrice))

        # Limit quantity to what is allowed by portfolio allocation
        quantity_limit = self.algorithm.calculate_order_quantity(self.security.symbol, 1 / self.algorithm.max_positions)
        if quantity > 0:
            sign = 1
        elif quantity < 0:
            sign = -1
        else:
            sign = 0
        quantity = int(min(abs(quantity), quantity_limit) * sign)

        if quantity != 0:
            self.stop_loss_price = stopPrice
            self.entry_ticket = self.algorithm.stop_market_order(self.security.symbol, quantity, entryPrice, "Entry")

    def on_order_event(self, orderEvent):
        if self.entry_ticket and orderEvent.order_id == self.entry_ticket.order_id:
            self.stop_loss_ticket = self.algorithm.stop_market_order(
                self.security.symbol, -self.entry_ticket.quantity, self.stop_loss_price, "Stop Loss"
            )